Handle responses without content length header

Co-authored-by: Kevin Pollet <pollet.kevin@gmail.com>

pkg/proxy/fast/connpool.go

@@ -216,7 +216,9 @@ func (c *conn) handleResponse(r rwWithUpgrade) error {
 		return nil
 	}
 
-	if res.Header.ContentLength() == 0 {
+	hasContentLength := res.Header.Peek("Content-Length") != nil
+
+	if hasContentLength && res.Header.ContentLength() == 0 {
 		return nil
 	}
 
@@ -226,6 +228,20 @@ func (c *conn) handleResponse(r rwWithUpgrade) error {
 		return nil
 	}
 
+	if !hasContentLength {
+		b := c.bufferPool.Get()
+		if b == nil {
+			b = make([]byte, bufferSize)
+		}
+		defer c.bufferPool.Put(b)
+
+		if _, err := io.CopyBuffer(r.RW, c.br, b); err != nil {
+			return err
+		}
+
+		return nil
+	}
+
 	// Chunked response, Content-Length is set to -1 by FastProxy when "Transfer-Encoding: chunked" header is received.
 	if res.Header.ContentLength() == -1 {
 		cbr := httputil.NewChunkedReader(c.br)

pkg/proxy/fast/proxy_test.go

@@ -306,6 +306,46 @@ func TestHeadRequest(t *testing.T) {
 	assert.Equal(t, http.StatusOK, res.Code)
 }
 
+func TestNoContentLengthResponse(t *testing.T) {
+	backendListener, err := net.Listen("tcp", ":0")
+	require.NoError(t, err)
+
+	t.Cleanup(func() {
+		_ = backendListener.Close()
+	})
+
+	go func() {
+		t.Helper()
+
+		conn, err := backendListener.Accept()
+		require.NoError(t, err)
+
+		_, err = conn.Write([]byte("HTTP/1.1 200 OK\r\n\r\nfoo"))
+		require.NoError(t, err)
+
+		// CloseWrite the connection to signal the end of the response.
+		if v, ok := conn.(interface{ CloseWrite() error }); ok {
+			err = v.CloseWrite()
+			require.NoError(t, err)
+		}
+	}()
+
+	builder := NewProxyBuilder(&transportManagerMock{}, static.FastProxyConfig{})
+
+	serverURL := "http://" + backendListener.Addr().String()
+
+	proxyHandler, err := builder.Build("", testhelpers.MustParseURL(serverURL), true, true)
+	require.NoError(t, err)
+
+	req := httptest.NewRequest(http.MethodGet, "/", http.NoBody)
+	res := httptest.NewRecorder()
+
+	proxyHandler.ServeHTTP(res, req)
+
+	assert.Equal(t, http.StatusOK, res.Code)
+	assert.Equal(t, "foo", res.Body.String())
+}
+
 func newCertificate(t *testing.T, domain string) *tls.Certificate {
 	t.Helper()