Add AWSEditHeadersMiddleware (#527)

* Add AWSEditHeadersMiddleware

Middleware for adding, replacing, removing headers.
Also moved some code about.

* Fix Sendable error in swift 5.4

* comment update

Sources/SotoCore/Message/AWSMiddleware.swift

@@ -12,9 +12,6 @@
 //
 //===----------------------------------------------------------------------===//
 
-import Logging
-import NIOHTTP1
-
 /// Context object sent to `AWSServiceMiddleware` `chain` functions
 public struct AWSMiddlewareContext {
     public let options: AWSServiceConfig.Options
@@ -39,85 +36,3 @@ public extension AWSServiceMiddleware {
         return response
     }
 }
-
-/// Middleware that outputs the contents of requests being sent to AWS and the contents of the responses received.
-public struct AWSLoggingMiddleware: AWSServiceMiddleware {
-    #if compiler(>=5.6)
-    public typealias LoggingFunction = @Sendable (String) -> Void
-    #else
-    public typealias LoggingFunction = (String) -> Void
-    #endif
-    /// initialize AWSLoggingMiddleware
-    /// - parameters:
-    ///     - log: Function to call with logging output
-    public init(log: @escaping LoggingFunction = { print($0) }) {
-        self.log = { log($0()) }
-    }
-
-    /// initialize AWSLoggingMiddleware to use Logger
-    /// - Parameters:
-    ///   - logger: Logger to use
-    ///   - logLevel: Log level to output at
-    public init(logger: Logger, logLevel: Logger.Level = .info) {
-        self.log = { logger.log(level: logLevel, "\($0())") }
-    }
-
-    func getBodyOutput(_ body: Body) -> String {
-        var output = ""
-        switch body {
-        case .xml(let element):
-            output += "\n  "
-            output += element.description
-        case .json(let buffer):
-            output += "\n  "
-            output += buffer.getString(at: buffer.readerIndex, length: buffer.readableBytes) ?? "Failed to convert JSON response to UTF8"
-        case .raw(let payload):
-            output += "raw (\(payload.size?.description ?? "unknown") bytes)"
-        case .text(let string):
-            output += "\n  \(string)"
-        case .empty:
-            output += "empty"
-        }
-        return output
-    }
-
-    func getHeadersOutput(_ headers: HTTPHeaders) -> String {
-        if headers.count == 0 {
-            return "[]"
-        }
-        var output = "["
-        for header in headers {
-            output += "\n    \(header.name) : \(header.value)"
-        }
-        return output + "\n  ]"
-    }
-
-    /// output request
-    public func chain(request: AWSRequest, context: AWSMiddlewareContext) throws -> AWSRequest {
-        self.log(
-            "Request:\n" +
-                "  \(request.operation)\n" +
-                "  \(request.httpMethod) \(request.url)\n" +
-                "  Headers: \(self.getHeadersOutput(request.httpHeaders))\n" +
-                "  Body: \(self.getBodyOutput(request.body))"
-        )
-        return request
-    }
-
-    /// output response
-    public func chain(response: AWSResponse, context: AWSMiddlewareContext) throws -> AWSResponse {
-        self.log(
-            "Response:\n" +
-                "  Status : \(response.status.code)\n" +
-                "  Headers: \(self.getHeadersOutput(HTTPHeaders(response.headers.map { ($0, "\($1)") })))\n" +
-                "  Body: \(self.getBodyOutput(response.body))"
-        )
-        return response
-    }
-
-    #if compiler(>=5.6)
-    let log: @Sendable (@autoclosure () -> String) -> Void
-    #else
-    let log: (@autoclosure () -> String) -> Void
-    #endif
-}

Sources/SotoCore/Message/AWSRequest.swift

@@ -384,7 +384,7 @@ extension AWSRequest {
 
     /// Add headers standard to all requests "content-type" and "user-agent"
     private mutating func addStandardHeaders() {
-        httpHeaders.replaceOrAdd(name: "user-agent", value: "Soto/6.0")
+        httpHeaders.add(name: "user-agent", value: "Soto/6.0")
         guard httpHeaders["content-type"].first == nil else {
             return
         }

Sources/SotoCore/Message/Middleware/EditHeadersMiddleware.swift

@@ -0,0 +1,54 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the Soto for AWS open source project
+//
+// Copyright (c) 2022 the Soto project authors
+// Licensed under Apache License v2.0
+//
+// See LICENSE.txt for license information
+// See CONTRIBUTORS.txt for the list of Soto project authors
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+//===----------------------------------------------------------------------===//
+
+import NIOHTTP1
+
+/// Middleware for editing header values sent to AWS service.
+public struct AWSEditHeadersMiddleware: AWSServiceMiddleware {
+    public enum HeaderEdit {
+        case add(name: String, value: String)
+        case replace(name: String, value: String)
+        case remove(name: String)
+    }
+
+    let edits: [HeaderEdit]
+
+    init(_ edits: [HeaderEdit]) {
+        self.edits = edits
+    }
+
+    init(_ edits: HeaderEdit...) {
+        self.init(edits)
+    }
+
+    public func chain(request: AWSRequest, context: AWSMiddlewareContext) throws -> AWSRequest {
+        var request = request
+        for edit in self.edits {
+            switch edit {
+            case .add(let name, let value):
+                request.httpHeaders.add(name: name, value: value)
+            case .replace(let name, let value):
+                request.httpHeaders.replaceOrAdd(name: name, value: value)
+            case .remove(let name):
+                request.httpHeaders.remove(name: name)
+            }
+        }
+        return request
+    }
+}
+
+#if compiler(>=5.6)
+extension AWSEditHeadersMiddleware: Sendable {}
+extension AWSEditHeadersMiddleware.HeaderEdit: Sendable {}
+#endif

Sources/SotoCore/Message/Middleware/LoggingMiddleware.swift

@@ -0,0 +1,98 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the Soto for AWS open source project
+//
+// Copyright (c) 2017-2022 the Soto project authors
+// Licensed under Apache License v2.0
+//
+// See LICENSE.txt for license information
+// See CONTRIBUTORS.txt for the list of Soto project authors
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+//===----------------------------------------------------------------------===//
+
+import Logging
+import NIOHTTP1
+
+/// Middleware that outputs the contents of requests being sent to AWS and the contents of the responses received.
+public struct AWSLoggingMiddleware: AWSServiceMiddleware {
+    #if compiler(>=5.6)
+    public typealias LoggingFunction = @Sendable (String) -> Void
+    #else
+    public typealias LoggingFunction = (String) -> Void
+    #endif
+    /// initialize AWSLoggingMiddleware
+    /// - parameters:
+    ///     - log: Function to call with logging output
+    public init(log: @escaping LoggingFunction = { print($0) }) {
+        self.log = { log($0()) }
+    }
+
+    /// initialize AWSLoggingMiddleware to use Logger
+    /// - Parameters:
+    ///   - logger: Logger to use
+    ///   - logLevel: Log level to output at
+    public init(logger: Logger, logLevel: Logger.Level = .info) {
+        self.log = { logger.log(level: logLevel, "\($0())") }
+    }
+
+    func getBodyOutput(_ body: Body) -> String {
+        var output = ""
+        switch body {
+        case .xml(let element):
+            output += "\n  "
+            output += element.description
+        case .json(let buffer):
+            output += "\n  "
+            output += buffer.getString(at: buffer.readerIndex, length: buffer.readableBytes) ?? "Failed to convert JSON response to UTF8"
+        case .raw(let payload):
+            output += "raw (\(payload.size?.description ?? "unknown") bytes)"
+        case .text(let string):
+            output += "\n  \(string)"
+        case .empty:
+            output += "empty"
+        }
+        return output
+    }
+
+    func getHeadersOutput(_ headers: HTTPHeaders) -> String {
+        if headers.count == 0 {
+            return "[]"
+        }
+        var output = "["
+        for header in headers {
+            output += "\n    \(header.name) : \(header.value)"
+        }
+        return output + "\n  ]"
+    }
+
+    /// output request
+    public func chain(request: AWSRequest, context: AWSMiddlewareContext) throws -> AWSRequest {
+        self.log(
+            "Request:\n" +
+                "  \(request.operation)\n" +
+                "  \(request.httpMethod) \(request.url)\n" +
+                "  Headers: \(self.getHeadersOutput(request.httpHeaders))\n" +
+                "  Body: \(self.getBodyOutput(request.body))"
+        )
+        return request
+    }
+
+    /// output response
+    public func chain(response: AWSResponse, context: AWSMiddlewareContext) throws -> AWSResponse {
+        self.log(
+            "Response:\n" +
+                "  Status : \(response.status.code)\n" +
+                "  Headers: \(self.getHeadersOutput(HTTPHeaders(response.headers.map { ($0, "\($1)") })))\n" +
+                "  Body: \(self.getBodyOutput(response.body))"
+        )
+        return response
+    }
+
+    #if compiler(>=5.6)
+    let log: @Sendable (@autoclosure () -> String) -> Void
+    #else
+    let log: (@autoclosure () -> String) -> Void
+    #endif
+}

Sources/SotoTestUtils/TestServer.swift

@@ -471,8 +471,11 @@ extension AWSTestServer {
             }
         }
         var requestHeaders: [String: String] = [:]
-        for (key, value) in head.headers {
-            requestHeaders[key.description] = value
+        for (key, _) in head.headers {
+            if requestHeaders[key] == nil {
+                requestHeaders[key] = head.headers[key].joined(separator: ",")
+            }
+            // requestHeaders[key.description] = value
         }
         return Request(method: head.method, uri: head.uri, headers: requestHeaders, body: byteBuffer)
     }

Tests/SotoCoreTests/AWSClientTests.swift

@@ -866,4 +866,39 @@ class AWSClientTests: XCTestCase {
 
         XCTAssertNoThrow(try response.wait())
     }
+
+    func testEditHeaderMiddleware() {
+        let awsServer = AWSTestServer(serviceProtocol: .json)
+        let client = createAWSClient(credentialProvider: .empty)
+        defer {
+            XCTAssertNoThrow(try client.syncShutdown())
+            XCTAssertNoThrow(try awsServer.stop())
+        }
+
+        // Test add header
+        let middleware = AWSEditHeadersMiddleware(
+            .add(name: "testAdd", value: "testValue"),
+            .add(name: "user-agent", value: "testEditHeaderMiddleware")
+        )
+        let config = createServiceConfig(endpoint: awsServer.address).with(patch: .init(middlewares: [middleware]))
+        let response = client.execute(operation: "test", path: "/", httpMethod: .POST, serviceConfig: config, logger: TestEnvironment.logger)
+        XCTAssertNoThrow(try awsServer.processRaw { request in
+            XCTAssertEqual(request.headers["testAdd"], "testValue")
+            XCTAssertEqual(request.headers["user-agent"], "Soto/6.0,testEditHeaderMiddleware")
+            return .result(AWSTestServer.Response.ok)
+        })
+        XCTAssertNoThrow(try response.wait())
+
+        // Test replace header
+        let middleware2 = AWSEditHeadersMiddleware(
+            .replace(name: "user-agent", value: "testEditHeaderMiddleware")
+        )
+        let config2 = createServiceConfig(endpoint: awsServer.address).with(patch: .init(middlewares: [middleware2]))
+        let response2 = client.execute(operation: "test", path: "/", httpMethod: .POST, serviceConfig: config2, logger: TestEnvironment.logger)
+        XCTAssertNoThrow(try awsServer.processRaw { request in
+            XCTAssertEqual(request.headers["user-agent"], "testEditHeaderMiddleware")
+            return .result(AWSTestServer.Response.ok)
+        })
+        XCTAssertNoThrow(try response2.wait())
+    }
 }