From 5a5aa2f3d6b95b42c5a73dbee59b1a340a83c932 Mon Sep 17 00:00:00 2001
From: Harry Barber <106155934+hlbarber@users.noreply.github.com>
Date: Wed, 29 Jun 2022 11:40:41 +0100
Subject: [PATCH] Set Content-Length header in server responses (#1423)

* Add `Content-Length` header to non-streaming server responses

* Add a `Content-Length` checks to the existing `httpResponseTests`
---
 codegen-server-test/model/simple.smithy       |  3 +++
 .../ServerHttpBoundProtocolGenerator.kt       | 26 +++++++++++++++++--
 codegen-test/model/misc.smithy                |  8 +++++-
 3 files changed, 34 insertions(+), 3 deletions(-)

diff --git a/codegen-server-test/model/simple.smithy b/codegen-server-test/model/simple.smithy
index d95ec84fd..6e094abe1 100644
--- a/codegen-server-test/model/simple.smithy
+++ b/codegen-server-test/model/simple.smithy
@@ -66,6 +66,9 @@ resource Service {
         params: { id: "1", name: "TestService" },
         body: "{\"id\":\"1\",\"name\":\"TestService\"}",
         code: 200,
+        headers: {
+            "Content-Length": "31"
+        }
     }
 ])
 operation RegisterService {
diff --git a/codegen-server/src/main/kotlin/software/amazon/smithy/rust/codegen/server/smithy/protocols/ServerHttpBoundProtocolGenerator.kt b/codegen-server/src/main/kotlin/software/amazon/smithy/rust/codegen/server/smithy/protocols/ServerHttpBoundProtocolGenerator.kt
index 5af80140a..56359d59c 100644
--- a/codegen-server/src/main/kotlin/software/amazon/smithy/rust/codegen/server/smithy/protocols/ServerHttpBoundProtocolGenerator.kt
+++ b/codegen-server/src/main/kotlin/software/amazon/smithy/rust/codegen/server/smithy/protocols/ServerHttpBoundProtocolGenerator.kt
@@ -473,6 +473,9 @@ private class ServerHttpBoundProtocolTraitImplGenerator(
                     val status =
                         variantShape.getTrait<HttpErrorTrait>()?.let { trait -> trait.code }
                             ?: errorTrait.defaultHttpStatusCode
+
+                    serverRenderContentLengthHeader()
+
                     rustTemplate(
                         """
                         builder.status($status).body(#{SmithyHttpServer}::body::to_boxed(payload))?
@@ -508,19 +511,22 @@ private class ServerHttpBoundProtocolTraitImplGenerator(
             val memberName = symbolProvider.toMemberName(it)
             rustTemplate(
                 """
-                let body = #{SmithyHttpServer}::body::to_boxed(#{SmithyHttpServer}::body::Body::wrap_stream(output.$memberName));
+                let payload = #{SmithyHttpServer}::body::Body::wrap_stream(output.$memberName);
                 """,
                 *codegenScope,
             )
         } ?: run {
             val payloadGenerator = HttpBoundProtocolPayloadGenerator(codegenContext, protocol, httpMessageType = HttpMessageType.RESPONSE)
-            withBlockTemplate("let body = #{SmithyHttpServer}::body::to_boxed(", ");", *codegenScope) {
+            withBlockTemplate("let payload = ", ";") {
                 payloadGenerator.generatePayload(this, "output", operationShape)
             }
+
+            serverRenderContentLengthHeader()
         }
 
         rustTemplate(
             """
+            let body = #{SmithyHttpServer}::body::to_boxed(payload);
             builder.body(body)?
             """,
             *codegenScope,
@@ -585,6 +591,22 @@ private class ServerHttpBoundProtocolTraitImplGenerator(
         }
     }
 
+    /**
+     * Adds the `Content-Length` header.
+     *
+     * Unlike the headers added in `serverRenderResponseHeaders` the `Content-Length` depends on
+     * the payload post-serialization.
+     */
+    private fun RustWriter.serverRenderContentLengthHeader() {
+        rustTemplate(
+            """
+            let content_length = payload.len();
+            builder = #{header_util}::set_response_header_if_absent(builder, #{http}::header::CONTENT_LENGTH, content_length);
+            """,
+            *codegenScope
+        )
+    }
+
     private fun serverRenderHttpResponseCode(
         defaultCode: Int
     ): Writable {
diff --git a/codegen-test/model/misc.smithy b/codegen-test/model/misc.smithy
index f32b892a1..d29744bf7 100644
--- a/codegen-test/model/misc.smithy
+++ b/codegen-test/model/misc.smithy
@@ -161,6 +161,9 @@ operation ResponseCodeDefaultOperation {
         id: "ResponseCodeHttpFallbackOperation",
         protocol: "aws.protocols#restJson1",
         code: 418,
+        headers: {
+            "Content-Length": "2"
+        }
     }
 ])
 @http(method: "GET", uri: "/responseCodeHttpFallbackOperation", code: 418)
@@ -178,7 +181,10 @@ structure EmptyStructure {}
         id: "ResponseCodeRequiredOperation",
         protocol: "aws.protocols#restJson1",
         code: 201,
-        params: {"responseCode": 201}
+        params: {"responseCode": 201},
+        headers: {
+            "Content-Length": "2"
+        }
     }
 ])
 @http(method: "GET", uri: "/responseCodeRequiredOperation", code: 200)
-- 
GitLab