fix Excess found: excess = 2 url = / (zero-length body) in HEAD response

.golangci.yml

@@ -50,7 +50,7 @@ linters-settings:
   errcheck:
     check-blank: true
   funlen:
-    lines: 140
+    lines: 160
     statements: 75
   gocognit:
     min-complexity: 50

internal/martian/handler.go

@@ -54,6 +54,10 @@ func addTrailerHeader(rw http.ResponseWriter, tr http.Header) int {
 }
 
 func copyBody(w io.Writer, body io.ReadCloser) error {
+	if body == http.NoBody {
+		return nil
+	}
+
 	bufp := copyBufPool.Get().(*[]byte) //nolint:forcetypeassert // It's *[]byte.
 	buf := *bufp
 	defer copyBufPool.Put(bufp)

internal/martian/head.go

@@ -0,0 +1,77 @@
+// Copyright 2023 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 martian
+
+import (
+	"fmt"
+	"io"
+	"net/http"
+	"strconv"
+	"strings"
+
+	"golang.org/x/exp/maps"
+)
+
+// writeHeadResponse writes the status line and header of r to w.
+func writeHeadResponse(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/httpspec/httpspec_test.go

@@ -40,6 +40,7 @@ func TestNewStack(t *testing.T) {
 	// Hop-by-hop header to be removed.
 	req.Header.Set("Hop-By-Hop", "true")
 	req.Header.Set("Connection", "Hop-By-Hop")
+	req.Header.Set("Transfer-Encoding", "chunked")
 
 	req.RemoteAddr = "10.0.0.1:5000"
 
@@ -50,6 +51,12 @@ func TestNewStack(t *testing.T) {
 	if got, want := req.Header.Get("Hop-By-Hop"), ""; got != want {
 		t.Errorf("req.Header.Get(%q): got %q, want %q", "Hop-By-Hop", got, want)
 	}
+	if got, want := req.Header.Get("Connection"), ""; got != want {
+		t.Errorf("req.Header.Get(%q): got %q, want %q", "Connection", got, want)
+	}
+	if got, want := req.Header.Get("Transfer-Encoding"), ""; got != want {
+		t.Errorf("req.Header.Get(%q): got %q, want %q", "Transfer-Encoding", got, want)
+	}
 	if got, want := req.Header.Get("X-Forwarded-For"), "10.0.0.1"; got != want {
 		t.Errorf("req.Header.Get(%q): got %q, want %q", "X-Forwarded-For", got, want)
 	}
@@ -62,6 +69,7 @@ func TestNewStack(t *testing.T) {
 	// Hop-by-hop header to be removed.
 	res.Header.Set("Hop-By-Hop", "true")
 	res.Header.Set("Connection", "Hop-By-Hop")
+	res.Header.Set("Transfer-Encoding", "chunked")
 
 	if err := stack.ModifyResponse(res); err != nil {
 		t.Fatalf("ModifyResponse(): got %v, want no error", err)
@@ -70,4 +78,10 @@ func TestNewStack(t *testing.T) {
 	if got, want := res.Header.Get("Hop-By-Hop"), ""; got != want {
 		t.Errorf("res.Header.Get(%q): got %q, want %q", "Hop-By-Hop", got, want)
 	}
+	if got, want := res.Header.Get("Connection"), ""; got != want {
+		t.Errorf("res.Header.Get(%q): got %q, want %q", "Connection", got, want)
+	}
+	if got, want := res.Header.Get("Transfer-Encoding"), ""; got != want {
+		t.Errorf("res.Header.Get(%q): got %q, want %q", "Transfer-Encoding", got, want)
+	}
 }

internal/martian/proxy.go

@@ -705,13 +705,20 @@ func (p *Proxy) handle(ctx *Context, conn net.Conn, brw *bufio.ReadWriter) error
 		}
 	}
 
-	// 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).
-	// If the event is larger than the buffer, the event will be split into multiple chunks.
-	if shouldFlush(res) {
-		err = res.Write(flushAfterChunkWriter{brw.Writer})
+	if req.Method == "HEAD" && res.Body == http.NoBody {
+		// 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(brw.Writer, res)
 	} else {
-		err = res.Write(brw)
+		// 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).
+		// If the event is larger than the buffer, the event will be split into multiple chunks.
+		if shouldFlush(res) {
+			err = res.Write(flushAfterChunkWriter{brw.Writer})
+		} else {
+			err = res.Write(brw)
+		}
 	}
 	if err != nil {
 		tracedLog.Errorf("martian: got error while writing response back to client: %v", err)