From 38f87ed52e6bbabcaa182866f87591e97cc38ad5 Mon Sep 17 00:00:00 2001
From: Amit Kumar Lamba <amit.lamba@uber.com>
Date: Tue, 23 Feb 2021 16:27:45 +0530
Subject: [PATCH 1/2] skip writing content when status code is 204
 (StatusNoContent)

---
 runtime/server_http_response.go | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/runtime/server_http_response.go b/runtime/server_http_response.go
index bdf42e04e..1f22dff7e 100644
--- a/runtime/server_http_response.go
+++ b/runtime/server_http_response.go
@@ -24,6 +24,7 @@ import (
 	"context"
 	"fmt"
 	"net/http"
+	"strconv"
 	"strings"
 	"time"
 
@@ -291,7 +292,9 @@ func (res *ServerHTTPResponse) flush(ctx context.Context) {
 
 	res.flushed = true
 	res.writeHeader(res.pendingStatusCode)
-	res.writeBytes(res.pendingBodyBytes)
+	if res.pendingStatusCode != http.StatusNoContent {
+		res.writeBytes(res.pendingBodyBytes)
+	}
 	res.finish(ctx)
 }
 
@@ -308,6 +311,7 @@ func (res *ServerHTTPResponse) writeBytes(bytes []byte) {
 		res.contextLogger.Error(res.Request.Context(),
 			"Could not write string to resp body",
 			zap.Error(err),
+			zap.String("bytesLength", strconv.Itoa(len(bytes))),
 		)
 	}
 }

From 42d9a8d73050ef10646dd92735cf904cf4cf1302 Mon Sep 17 00:00:00 2001
From: Amit Kumar Lamba <amit.lamba@uber.com>
Date: Thu, 25 Feb 2021 10:58:45 +0530
Subject: [PATCH 2/2] TestPendingResponseBody204StatusNoContent test case

---
 runtime/server_http_response_test.go | 65 ++++++++++++++++++++++++++++
 1 file changed, 65 insertions(+)

diff --git a/runtime/server_http_response_test.go b/runtime/server_http_response_test.go
index ac2408557..d87a5b0ab 100644
--- a/runtime/server_http_response_test.go
+++ b/runtime/server_http_response_test.go
@@ -520,6 +520,71 @@ func TestPendingResponseBody(t *testing.T) {
 	)
 }
 
+func TestPendingResponseBody204StatusNoContent(t *testing.T) {
+	gateway, err := benchGateway.CreateGateway(
+		defaultTestConfig,
+		defaultTestOptions,
+		exampleGateway.CreateGateway,
+	)
+
+	if !assert.NoError(t, err) {
+		return
+	}
+	defer gateway.Close()
+
+	bgateway := gateway.(*benchGateway.BenchGateway)
+	deps := createDefaultDependencies(bgateway)
+	err = bgateway.ActualGateway.HTTPRouter.Handle(
+		"GET", "/foo", http.HandlerFunc(zanzibar.NewRouterEndpoint(
+			bgateway.ActualGateway.ContextExtractor,
+			deps,
+			"foo", "foo",
+			func(
+				ctx context.Context,
+				req *zanzibar.ServerHTTPRequest,
+				res *zanzibar.ServerHTTPResponse,
+			) {
+				obj := &MyBody{
+					Token: "myToken",
+					Client: MyBodyClient{
+						Token: "myClientToken",
+					},
+				}
+				bytes, err := json.Marshal(obj)
+				statusCode := 204
+				assert.NoError(t, err)
+				res.WriteJSON(statusCode, nil, obj)
+
+				pendingBytes, pendingStatusCode := res.GetPendingResponse()
+				assert.Equal(t, bytes, pendingBytes)
+				assert.Equal(t, statusCode, pendingStatusCode)
+
+				headers := res.Headers()
+				assert.NotNil(t, headers)
+			},
+		).HandleRequest),
+	)
+	assert.NoError(t, err)
+
+	resp, err := gateway.MakeRequest("GET", "/foo", nil, nil)
+	if !assert.NoError(t, err) {
+		return
+	}
+
+	assert.Equal(t, resp.StatusCode, 204)
+	bytes, err := ioutil.ReadAll(resp.Body)
+	if !assert.NoError(t, err) {
+		return
+	}
+
+	//The body would become blank
+	assert.Equal(
+		t,
+		"",
+		string(bytes),
+	)
+}
+
 func TestPendingResponseObject(t *testing.T) {
 	gateway, err := benchGateway.CreateGateway(
 		defaultTestConfig,