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

@@ -18,6 +18,7 @@ import (
 	"context"
 	"crypto/tls"
 	"errors"
+	"io"
 	"net"
 	"net/http"
 	"net/url"
@@ -353,7 +354,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 {
@@ -377,3 +389,11 @@ func upgradeType(h http.Header) string {
 	}
 	return h.Get("Upgrade")
 }
+
+type panicReader struct{}
+
+func (panicReader) Read(p []byte) (int, error) {
+	panic("unexpected read")
+}
+
+var panicBody = io.NopCloser(panicReader{})

internal/martian/proxy_conn.go

@@ -26,10 +26,13 @@ import (
 	"io"
 	"net"
 	"net/http"
+	"strconv"
+	"strings"
 	"time"
 
 	"github.com/saucelabs/forwarder/internal/martian/log"
 	"github.com/saucelabs/forwarder/internal/martian/proxyutil"
+	"golang.org/x/exp/maps"
 )
 
 type proxyConn struct {
@@ -256,8 +259,7 @@ func (p *proxyConn) handleUpgradeResponse(res *http.Response) error {
 		log.Errorf(res.Request.Context(), "internal error: switching protocols response with non-writable body")
 		return errClose
 	}
-
-	res.Body = nil
+	res.Body = panicBody
 
 	if err := p.tunnel(resUpType, res, uconn); err != nil {
 		log.Errorf(res.Request.Context(), "%s tunnel: %v", resUpType, err)
@@ -427,11 +429,11 @@ 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.
-		err = writeHeadResponse(p.brw.Writer, res)
+		err = writeHeaderOnlyResponse(p.brw.Writer, res)
 	default:
 		// Add support for Server Sent Events - relay HTTP chunks and flush after each chunk.
 		// This is safe for events that are smaller than the buffer io.Copy uses (32KB).
@@ -471,3 +473,53 @@ func (p *proxyConn) writeResponse(res *http.Response) error {
 
 	return nil
 }
+
+// writeHeaderOnlyResponse writes the status line and header of r to w.
+func writeHeaderOnlyResponse(w io.Writer, res *http.Response) error {
+	// Status line
+	text := res.Status
+	if text == "" {
+		text = http.StatusText(res.StatusCode)
+		if text == "" {
+			text = "status code " + strconv.Itoa(res.StatusCode)
+		}
+	} else {
+		// Just to reduce stutter, if user set res.Status to "200 OK" and StatusCode to 200.
+		// Not important.
+		text = strings.TrimPrefix(text, strconv.Itoa(res.StatusCode)+" ")
+	}
+
+	if _, err := fmt.Fprintf(w, "HTTP/%d.%d %03d %s\r\n", res.ProtoMajor, res.ProtoMinor, res.StatusCode, text); err != nil {
+		return err
+	}
+
+	// Header
+	if err := res.Header.Write(w); err != nil {
+		return err
+	}
+
+	// Add Trailer header if needed
+	if len(res.Trailer) > 0 {
+		if _, err := io.WriteString(w, "Trailer: "); err != nil {
+			return err
+		}
+
+		for i, k := range maps.Keys(res.Trailer) {
+			if i > 0 {
+				if _, err := io.WriteString(w, ", "); err != nil {
+					return err
+				}
+			}
+			if _, err := io.WriteString(w, k); err != nil {
+				return err
+			}
+		}
+	}
+
+	// End-of-header
+	if _, err := io.WriteString(w, "\r\n"); err != nil {
+		return err
+	}
+
+	return nil
+}

internal/martian/proxy_handler.go

@@ -149,8 +149,7 @@ func (p proxyHandler) handleUpgradeResponse(rw http.ResponseWriter, req *http.Re
 		log.Errorf(ctx, "%s tunnel: internal error: switching protocols response with non-ReadWriteCloser body", resUpType)
 		panic(http.ErrAbortHandler)
 	}
-
-	res.Body = nil
+	res.Body = panicBody
 
 	if err := p.tunnel(resUpType, rw, req, res, uconn); err != nil {
 		log.Errorf(ctx, "%s tunnel: %v", resUpType, err)

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()