diff --git a/internal/martian/head.go b/internal/martian/head.go new file mode 100644 index 00000000..7186f01e --- /dev/null +++ b/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 +} diff --git a/internal/martian/proxy.go b/internal/martian/proxy.go index 2ac1ed3b..bfca80b9 100644 --- a/internal/martian/proxy.go +++ b/internal/martian/proxy.go @@ -699,13 +699,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 { log.Errorf("martian: got error while writing response back to client: %v", err)