From 6910e5947e3f00b7ef9c9ad2c21be66e5245fca3 Mon Sep 17 00:00:00 2001
From: Adam Fowler <adamfowler71@gmail.com>
Date: Tue, 13 Aug 2019 15:40:37 +0100
Subject: [PATCH] Added Middleware editing of response bodies

Needed to add this to fix S3.GetBucketLocation does not return a location https://github.com/swift-aws/aws-sdk-swift/issues/157
---
 Sources/AWSSDKSwiftCore/AWSClient.swift          | 11 ++++++++---
 Sources/AWSSDKSwiftCore/Encoder/XMLDecoder.swift |  2 +-
 Sources/AWSSDKSwiftCore/Message/AWSRequest.swift | 16 +++++++++++++---
 3 files changed, 22 insertions(+), 7 deletions(-)

diff --git a/Sources/AWSSDKSwiftCore/AWSClient.swift b/Sources/AWSSDKSwiftCore/AWSClient.swift
index 32ea57f0a..164fcf68e 100644
--- a/Sources/AWSSDKSwiftCore/AWSClient.swift
+++ b/Sources/AWSSDKSwiftCore/AWSClient.swift
@@ -44,7 +44,7 @@ public struct AWSClient {
 
     let partitionEndpoint: String?
 
-    public let middlewares: [AWSRequestMiddleware]
+    public let middlewares: [AWSServiceMiddleware]
 
     public var possibleErrorTypes: [AWSErrorType.Type]
 
@@ -68,7 +68,7 @@ public struct AWSClient {
 
     public static let eventGroup: EventLoopGroup = MultiThreadedEventLoopGroup(numberOfThreads: 1)
 
-    public init(accessKeyId: String? = nil, secretAccessKey: String? = nil, sessionToken: String? = nil, region givenRegion: Region?, amzTarget: String? = nil, service: String, serviceProtocol: ServiceProtocol, apiVersion: String, endpoint: String? = nil, serviceEndpoints: [String: String] = [:], partitionEndpoint: String? = nil, middlewares: [AWSRequestMiddleware] = [], possibleErrorTypes: [AWSErrorType.Type]? = nil) {
+    public init(accessKeyId: String? = nil, secretAccessKey: String? = nil, sessionToken: String? = nil, region givenRegion: Region?, amzTarget: String? = nil, service: String, serviceProtocol: ServiceProtocol, apiVersion: String, endpoint: String? = nil, serviceEndpoints: [String: String] = [:], partitionEndpoint: String? = nil, middlewares: [AWSServiceMiddleware] = [], possibleErrorTypes: [AWSErrorType.Type]? = nil) {
         let credential: CredentialProvider
         if let accessKey = accessKeyId, let secretKey = secretAccessKey {
             credential = Credential(accessKeyId: accessKey, secretAccessKey: secretKey, sessionToken: sessionToken)
@@ -492,12 +492,17 @@ extension AWSClient {
 
         try validateCode(response: response, members: Output._members)
 
-        let responseBody = try validateBody(
+        var responseBody = try validateBody(
             for: response,
             payloadPath: Output.payloadPath,
             members: Output._members
         )
 
+        // do we need to fix up the response before processing it
+        for middleware in middlewares {
+            responseBody = try middleware.chain(responseBody: responseBody)
+        }
+        
         let decoder = DictionaryDecoder()
 
         var responseHeaders: [String: String] = [:]
diff --git a/Sources/AWSSDKSwiftCore/Encoder/XMLDecoder.swift b/Sources/AWSSDKSwiftCore/Encoder/XMLDecoder.swift
index 214837386..3684bf0c3 100644
--- a/Sources/AWSSDKSwiftCore/Encoder/XMLDecoder.swift
+++ b/Sources/AWSSDKSwiftCore/Encoder/XMLDecoder.swift
@@ -693,7 +693,7 @@ fileprivate class _XMLDecoder : Decoder {
     }
 
     func unbox(_ element : XML.Element, as type: String.Type) throws -> String {
-        guard let unboxValue = element.stringValue else { throw DecodingError._typeMismatch(at: codingPath, expectation: Bool.self, reality: element.stringValue ?? "nil") }
+        guard let unboxValue = element.stringValue else { throw DecodingError._typeMismatch(at: codingPath, expectation: String.self, reality: element.stringValue ?? "nil") }
         return unboxValue
     }
 
diff --git a/Sources/AWSSDKSwiftCore/Message/AWSRequest.swift b/Sources/AWSSDKSwiftCore/Message/AWSRequest.swift
index cf4aa70bd..7177fa637 100644
--- a/Sources/AWSSDKSwiftCore/Message/AWSRequest.swift
+++ b/Sources/AWSSDKSwiftCore/Message/AWSRequest.swift
@@ -11,8 +11,18 @@ import NIO
 import NIOTLS
 import NIOHTTP1
 
-public protocol AWSRequestMiddleware {
+public protocol AWSServiceMiddleware {
     func chain(request: AWSRequest) throws -> AWSRequest
+    func chain(responseBody: Body) throws -> Body
+}
+
+public extension AWSServiceMiddleware {
+    func chain(request: AWSRequest) throws -> AWSRequest {
+        return request
+    }
+    func chain(responseBody: Body) throws -> Body {
+        return responseBody
+    }
 }
 
 extension URL {
@@ -51,9 +61,9 @@ public struct AWSRequest {
     public let httpMethod: String
     public var httpHeaders: [String: Any?] = [:]
     public var body: Body
-    public let middlewares: [AWSRequestMiddleware]
+    public let middlewares: [AWSServiceMiddleware]
 
-    public init(region: Region = .useast1, url: URL, serviceProtocol: ServiceProtocol, service: String, amzTarget: String? = nil, operation: String, httpMethod: String, httpHeaders: [String: Any?] = [:], body: Body = .empty, middlewares: [AWSRequestMiddleware] = []) {
+    public init(region: Region = .useast1, url: URL, serviceProtocol: ServiceProtocol, service: String, amzTarget: String? = nil, operation: String, httpMethod: String, httpHeaders: [String: Any?] = [:], body: Body = .empty, middlewares: [AWSServiceMiddleware] = []) {
         self.region = region
         self.url = url
         self.serviceProtocol = serviceProtocol