Keep body if request uses chunked TransferEncoding (awslabs#152) (#5)

Fixes awslabs#151

Co-authored-by: Luis Gerhorst <luisgerhorst@users.noreply.github.com>

handler/proxy_client.go

@@ -99,6 +99,25 @@ func copyHeaderWithoutOverwrite(dst, src http.Header) {
 	}
 }
 
+// RFC2616, Section 4.4: If a Transfer-Encoding header field (Section 14.41) is
+// present and has any value other than "identity", then the transfer-length is
+// defined by use of the "chunked" transfer-coding (Section 3.6). [...] If a
+// message is received with both a Transfer-Encoding header field and a
+// Content-Length header field, the latter MUST be ignored.
+//
+// RFC2616, Section 3.6: Whenever a transfer-coding is applied to a
+// message-body, the set of transfer-codings MUST include "chunked", unless the
+// message is terminated by closing the connection.
+func chunked(transferEncoding []string) bool {
+	for _, v := range transferEncoding {
+		// This interprets identity-only headers as no header.
+		if v != "identity" {
+			return true
+		}
+	}
+	return false
+}
+
 func (p *ProxyClient) Do(req *http.Request) (*http.Response, error) {
 	proxyURL := *req.URL
 	if p.HostOverride != "" {
@@ -124,7 +143,11 @@ func (p *ProxyClient) Do(req *http.Request) (*http.Response, error) {
 	if err != nil {
 		return nil, err
 	}
-	if req.ContentLength >= 0 {
+
+	var reqChunked = chunked(req.TransferEncoding);
+
+	// Ignore ContentLength if "chunked" transfer-coding is used.
+	if !reqChunked && req.ContentLength >= 0 {
 		proxyReq.ContentLength = req.ContentLength
 	}
 
@@ -145,12 +168,20 @@ func (p *ProxyClient) Do(req *http.Request) (*http.Response, error) {
 		return nil, err
 	}
 
-	// When ContentLength is 0 we also need to set the body to http.NoBody to avoid Go http client
-	// to magically set Transfer-Encoding: chunked. Service like S3 does not support chunk encoding.
-	// We need to manipulate the Body value after signv4 signing because the signing process wraps
-	// the original body into another struct, which will result in Transfer-Encoding: chunked being set.
-	if proxyReq.ContentLength == 0 {
-		proxyReq.Body = http.NoBody
+	// go Documentation net/http, func (*Request) Write: If Body is present,
+	// Content-Length is <= 0 and TransferEncoding hasn't been set to
+	// "identity", Write adds "Transfer-Encoding: chunked" to the header.
+	// Body is closed after it is sent.
+	//
+	// Service like S3 does not support chunk encoding. We need to manipulate
+	// the Body value after signv4 signing because the signing process wraps the
+	// original body into another struct, which will result in
+	// Transfer-Encoding: chunked being set.
+	if !reqChunked {
+		// Set to identity to prevent write() from setting it to chunked.
+		proxyReq.TransferEncoding = []string{"identity"}
+	} else {
+		proxyReq.TransferEncoding = req.TransferEncoding
 	}
 
 	// Remove any headers specified

handler/proxy_client_test.go

@@ -288,6 +288,83 @@ func TestProxyClient_Do(t *testing.T) {
 				},
 			},
 		},
+		{
+			name: "should not drop body for chunked requests",
+			request: &http.Request{
+				Method:           "POST",
+				URL:              &url.URL{},
+				Host:             "not.important.host",
+				TransferEncoding: []string{"chunked"},
+				Body:             io.NopCloser(strings.NewReader("5\r\nhello\r\n")),
+			},
+			proxyClient: &ProxyClient{
+				Signer:              v4.NewSigner(credentials.NewCredentials(&mockProvider{})),
+				SigningNameOverride: "ec2",
+				RegionOverride:      "us-west-2",
+				Client:              &mockHTTPClient{},
+			},
+			want: &want{
+				resp: &http.Response{},
+				err:  nil,
+				request: &http.Request{
+					Host: "not.important.host",
+					// The test callback func will check that new body is at
+					// least as large as the original body.
+				},
+			},
+		},
+		{
+			name: "should ignore content length for chunked requests",
+			request: &http.Request{
+				Method:           "POST",
+				URL:              &url.URL{},
+				Host:             "not.important.host",
+				TransferEncoding: []string{"chunked"},
+				ContentLength:    10000, // invalid, but should be ignored
+				Body:             io.NopCloser(strings.NewReader("5\r\nhello\r\n0\r\n\r\n")),
+			},
+			proxyClient: &ProxyClient{
+				Signer:              v4.NewSigner(credentials.NewCredentials(&mockProvider{})),
+				SigningNameOverride: "ec2",
+				RegionOverride:      "us-west-2",
+				Client:              &mockHTTPClient{},
+			},
+			want: &want{
+				resp: &http.Response{},
+				err:  nil,
+				request: &http.Request{
+					Host: "not.important.host",
+					// The test callback func will check that new body is at
+					// least as large as the original body.
+				},
+			},
+		},
+		{
+			name: "should propagate nested unknown transfer encodings",
+			request: &http.Request{
+				Method:           "POST",
+				URL:              &url.URL{},
+				Host:             "not.important.host",
+				TransferEncoding: []string{"chunked", "unregisteredBogus"},
+				ContentLength:    0,
+				Body:             io.NopCloser(strings.NewReader("5\r\nHELLO\r\n0\r\n\r\n")),
+			},
+			proxyClient: &ProxyClient{
+				Signer:              v4.NewSigner(credentials.NewCredentials(&mockProvider{})),
+				SigningNameOverride: "ec2",
+				RegionOverride:      "us-west-2",
+				Client:              &mockHTTPClient{},
+			},
+			want: &want{
+				resp: &http.Response{},
+				err:  nil,
+				request: &http.Request{
+					Host: "not.important.host",
+					// The test callback func will check that new body is at
+					// least as large as the original body.
+				},
+			},
+		},
 	}
 
 	for _, tt := range tests {
@@ -304,15 +381,32 @@ func TestProxyClient_Do(t *testing.T) {
 				return
 			}
 
-			// Ensure content length is propagated to the proxy request
-			if proxyRequest != nil {
+			// Ensure encoding is propagated to the proxy request.
+			assert.Equal(t, chunked(tt.request.TransferEncoding), chunked(proxyRequest.TransferEncoding));
+			if chunked(tt.request.TransferEncoding) {
+				assert.Equal(t, tt.request.TransferEncoding, proxyRequest.TransferEncoding);
+			} else {
+				// Ensure content length is propagated to the proxy request.
 				assert.Equal(t, tt.request.ContentLength, proxyRequest.ContentLength)
+
+				// If this assertion is not true, then Go http client will use
+				// TransferEncoding: chunked, which may not be supported by AWS
+				// services like S3.
+				if tt.request.ContentLength == 0 {
+					assert.Equal(t, []string{"identity"}, proxyRequest.TransferEncoding)
+				}
 			}
 
-			// If this assertion is not true, then Go http client will use TransferEncoding: chunked, which
-			// may not be supported by AWS services like S3.
-			if tt.request.ContentLength == 0 {
-				assert.Equal(t, http.NoBody, proxyRequest.Body)
+			// Proxied request bodies should be at least as large as the original.
+			if tt.request.Body != nil {
+				ttBody, ttErr := io.ReadAll(tt.request.Body)
+				if proxyRequest.Body != nil {
+					proxyBody, proxyErr := io.ReadAll(proxyRequest.Body)
+					assert.Equal(t, ttErr, proxyErr)
+					assert.True(t, len(proxyBody) >= len(ttBody))
+				} else {
+					assert.Equal(t, 0, len(ttBody))
+				}
 			}
 		})
 	}