fix: content-length middleware should not error on event streams (#608)

smithy-lang · Oct 25, 2023 · 31d7652 · 31d7652
1 parent e0b4f25
commit 31d7652
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 7 deletions.
diff --git a/Sources/ClientRuntime/Networking/Http/Middlewares/ContentLengthMiddleware.swift b/Sources/ClientRuntime/Networking/Http/Middlewares/ContentLengthMiddleware.swift
@@ -6,11 +6,17 @@ public struct ContentLengthMiddleware<OperationStackOutput: HttpResponseBinding>
 
     private let contentLengthHeaderName = "Content-Length"
 
-    private var requiresLength: Bool = false
+    private var requiresLength: Bool?
 
-    private var unsignedPayload: Bool = false
+    private var unsignedPayload: Bool?
 
-    public init(requiresLength: Bool = false, unsignedPayload: Bool = false) {
+    /// Creates a new `ContentLengthMiddleware` with the supplied parameters
+    /// - Parameters:
+    ///   - requiresLength: Trait requires the length of a blob stream to be known. 
+    ///     When the request body is not a streaming blob, `nil` should be passed. Defaults to `nil`.
+    ///   - unsignedPayload: Trait signifies that the length of a stream in payload does not need to be known.
+    ///     When the request body is not a streaming blob, `nil` should be passed. Defaults to `nil`.
+    public init(requiresLength: Bool? = nil, unsignedPayload: Bool? = nil) {
         self.requiresLength = requiresLength
         self.unsignedPayload = unsignedPayload
     }
@@ -29,13 +35,16 @@ public struct ContentLengthMiddleware<OperationStackOutput: HttpResponseBinding>
         case .stream(let stream):
             if let length = stream.length {
                 input.headers.update(name: "Content-Length", value: String(length))
-            } else if !requiresLength && unsignedPayload {
-                // only for HTTP/1.1 requests, will be removed in all HTTP/2 requests
+            } else if (requiresLength == false && unsignedPayload == true) ||
+                        (requiresLength == nil && unsignedPayload == nil) {
+                // Transfer-Encoding can be sent on all Event Streams where length cannot be determined
+                // or on blob Data Streams where requiresLength is true and unsignedPayload is false
+                // Only for HTTP/1.1 requests, will be removed in all HTTP/2 requests
                 input.headers.update(name: "Transfer-Encoding", value: "Chunked")
             } else {
                 let operation = context.attributes.get(key: AttributeKey<String>(name: "Operation"))
                              ?? "Error getting operation name"
-                let errorMessage = unsignedPayload ?
+                let errorMessage = (unsignedPayload ?? false) ?
                     "Missing content-length for operation: \(operation)" :
                     "Missing content-length for SigV4 signing on operation: \(operation)"
                 throw StreamError.notSupported(errorMessage)

diff --git a/Tests/ClientRuntimeTests/NetworkingTests/ContentLengthMiddlewareTests.swift b/Tests/ClientRuntimeTests/NetworkingTests/ContentLengthMiddlewareTests.swift
@@ -27,6 +27,13 @@ class ContentLengthMiddlewareTests: XCTestCase {
         try await AssertHeadersArePresent(expectedHeaders: ["Transfer-Encoding": "Chunked"])
     }
 
+    func testTransferEncodingChunkedSetWithNilTraits() async throws {
+        // default constructor
+        addContentLengthMiddlewareWith(requiresLength: nil, unsignedPayload: nil)
+        forceEmptyStream()
+        try await AssertHeadersArePresent(expectedHeaders: ["Transfer-Encoding": "Chunked"])
+    }
+
     func testContentLengthSetWhenStreamLengthAvailableAndRequiresLengthSet() async throws {
         addContentLengthMiddlewareWith(requiresLength: true, unsignedPayload: false)
         try await AssertHeadersArePresent(expectedHeaders: ["Content-Length": "0"])
@@ -54,7 +61,7 @@ class ContentLengthMiddlewareTests: XCTestCase {
         }
     }
 
-    private func addContentLengthMiddlewareWith(requiresLength: Bool, unsignedPayload: Bool) {
+    private func addContentLengthMiddlewareWith(requiresLength: Bool?, unsignedPayload: Bool?) {
         stack.finalizeStep.intercept(
             position: .before,
             middleware: ContentLengthMiddleware(requiresLength: requiresLength, unsignedPayload: unsignedPayload)