Do not close client connection after returning HTTP status code 304

internal/martian/flush.go

@@ -36,21 +36,23 @@ func shouldChunk(res *http.Response) bool {
 		return false
 	}
 
-	// Please read 3.3.2 and 3.3.3 of RFC 7230 for more details https://datatracker.ietf.org/doc/html/rfc7230#section-3.3.2.
-	if res.Request.Method == http.MethodHead {
-		return false
-	}
-	// The 204/304 response MUST NOT contain a
-	// message-body, and thus is always terminated by the first empty line
-	// after the header fields.
-	if res.StatusCode == http.StatusNoContent || res.StatusCode == http.StatusNotModified {
-		return false
-	}
-	if res.StatusCode < 200 {
-		return false
-	}
+	return !isHeaderOnlySpec(res)
+}
 
-	return true
+// isHeaderOnlySpec returns true iff the response should have only headers according to the HTTP spec (RFC 7230).
+//
+// Any response to a HEAD request and any response with a 1xx
+// (Informational), 204 (No Content), or 304 (Not Modified) status
+// code is always terminated by the first empty line after the
+// header fields, regardless of the header fields present in the
+// message, and thus cannot contain a message body.
+//
+// https://datatracker.ietf.org/doc/html/rfc7230#section-3.3.3
+func isHeaderOnlySpec(res *http.Response) bool {
+	return res.Request.Method == http.MethodHead ||
+		res.StatusCode/100 == 1 ||
+		res.StatusCode == http.StatusNoContent ||
+		res.StatusCode == http.StatusNotModified
 }
 
 func isTextEventStream(res *http.Response) bool {

internal/martian/log/std.go

@@ -0,0 +1,42 @@
+// Copyright 2022-2024 Sauce Labs Inc., all rights reserved.
+//
+// Copyright 2015 Google Inc. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package log
+
+import (
+	"context"
+	stdlog "log"
+)
+
+type testLogger struct{}
+
+var _ Logger = testLogger{}
+
+func (testLogger) Infof(_ context.Context, format string, args ...any) {
+	stdlog.Printf("INFO: "+format, args...)
+}
+
+func (testLogger) Debugf(_ context.Context, format string, args ...any) {
+	stdlog.Printf("DEBUG: "+format, args...)
+}
+
+func (testLogger) Errorf(_ context.Context, format string, args ...any) {
+	stdlog.Printf("ERROR: "+format, args...)
+}
+
+func SetTestLogger() {
+	SetLogger(testLogger{})
+}

internal/martian/proxy.go

@@ -353,7 +353,18 @@ func (p *Proxy) roundTrip(req *http.Request) (*http.Response, error) {
 		return proxyutil.NewResponse(200, http.NoBody, req), nil
 	}
 
-	return p.rt.RoundTrip(req)
+	res, err := p.rt.RoundTrip(req)
+	if err != nil {
+		return nil, err
+	}
+
+	if isHeaderOnlySpec(res) && res.StatusCode != http.StatusSwitchingProtocols && res.Body != http.NoBody {
+		log.Infof(req.Context(), "unexpected body in header-only response: %d, closing body", res.StatusCode)
+		res.Body.Close()
+		res.Body = http.NoBody
+	}
+
+	return res, err
 }
 
 func (p *Proxy) errorResponse(req *http.Request, err error) *http.Response {

internal/martian/proxy_conn.go

@@ -427,7 +427,7 @@ func (p *proxyConn) writeResponse(res *http.Response) error {
 	switch {
 	case req.Method == http.MethodConnect && res.StatusCode/100 == 2:
 		err = writeConnectOKResponse(p.brw.Writer)
-	case req.Method == http.MethodHead && res.Body == http.NoBody:
+	case isHeaderOnlySpec(res):
 		// The http package is misbehaving when writing a HEAD response.
 		// See https://github.com/golang/go/issues/62015 for details.
 		// This works around the issue by writing the response manually.

internal/martian/proxy_test.go

@@ -485,6 +485,77 @@ func TestIntegrationHTTP101SwitchingProtocols(t *testing.T) {
 	}
 }
 
+func TestIntegrationHTTP304NotModified(t *testing.T) {
+	t.Parallel()
+
+	tm := martiantest.NewModifier()
+	h := testHelper{
+		Proxy: func(p *Proxy) {
+			p.ReadTimeout = 200 * time.Millisecond
+			p.WriteTimeout = 200 * time.Millisecond
+			p.RequestModifier = tm
+			p.ResponseModifier = tm
+			p.AllowHTTP = true
+		},
+	}
+
+	sl, err := net.Listen("tcp", "localhost:0")
+	if err != nil {
+		t.Fatalf("net.Listen(): got %v, want no error", err)
+	}
+
+	go func() {
+		conn, err := sl.Accept()
+		if err != nil {
+			log.Errorf(context.TODO(), "proxy_test: failed to accept connection: %v", err)
+			return
+		}
+		defer conn.Close()
+
+		log.Infof(context.TODO(), "proxy_test: accepted connection: %s", conn.RemoteAddr())
+
+		req, err := http.ReadRequest(bufio.NewReader(conn))
+		if err != nil {
+			log.Errorf(context.TODO(), "proxy_test: failed to read request: %v", err)
+			return
+		}
+
+		res := proxyutil.NewResponse(304, http.NoBody, req)
+		res.Header.Set("Content-Length", "13")
+		res.Write(conn)
+		log.Infof(context.TODO(), "proxy_test: sent 304 response")
+
+		log.Infof(context.TODO(), "proxy_test: closed connection")
+	}()
+
+	conn, cancel := h.proxyConn(t)
+	defer cancel()
+	defer conn.Close()
+
+	host := sl.Addr().String()
+
+	req, err := http.NewRequest(http.MethodGet, "http://"+host, http.NoBody)
+	if err != nil {
+		t.Fatalf("http.NewRequest(): got %v, want no error", err)
+	}
+	if err := req.WriteProxy(conn); err != nil {
+		t.Fatalf("req.WriteProxy(): got %v, want no error", err)
+	}
+
+	res, err := http.ReadResponse(bufio.NewReader(conn), nil)
+	if err != nil {
+		t.Fatalf("http.ReadResponse(): got %v, want no error", err)
+	}
+	defer res.Body.Close()
+
+	if got, want := res.StatusCode, 304; got != want {
+		t.Fatalf("res.StatusCode: got %d, want %d", got, want)
+	}
+	if _, err := conn.Read(nil); err != nil {
+		t.Fatalf("conn.Read(): got %v, want no error", err)
+	}
+}
+
 func TestIntegrationUnexpectedUpstreamFailure(t *testing.T) {
 	t.Parallel()