coder/coderd/tracing/status_writer.go

package tracing

import (
	"bufio"
	"flag"
	"fmt"
	"log"
	"net"
	"net/http"
	"runtime"
	"strings"

	"golang.org/x/xerrors"

	"github.com/coder/coder/v2/buildinfo"
)

var (
	_ http.ResponseWriter = (*StatusWriter)(nil)
	_ http.Hijacker       = (*StatusWriter)(nil)
)

// StatusWriter intercepts the status of the request and the response body up
// to maxBodySize if Status >= 400. It is guaranteed to be the ResponseWriter
// directly downstream from Middleware.
type StatusWriter struct {
	http.ResponseWriter
	Status       int
	Hijacked     bool
	responseBody []byte

	wroteHeader      bool
	wroteHeaderStack string
}

func StatusWriterMiddleware(next http.Handler) http.Handler {
	return http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) {
		sw := &StatusWriter{ResponseWriter: rw}
		next.ServeHTTP(sw, r)
	})
}

func (w *StatusWriter) WriteHeader(status int) {
	if buildinfo.IsDev() || flag.Lookup("test.v") != nil {
		if w.wroteHeader {
			stack := getStackString(2)
			wroteHeaderStack := w.wroteHeaderStack
			if wroteHeaderStack == "" {
				wroteHeaderStack = "unknown"
			}
			// It's fine that this logs to stdlib logger since it only happens
			// in dev builds and tests.
			log.Printf("duplicate call to (*StatusWriter.).WriteHeader(%d):\n\nstack: %s\n\nheader written at: %s", status, stack, wroteHeaderStack)
		} else {
			w.wroteHeaderStack = getStackString(2)
		}
	}
	if !w.wroteHeader {
		w.Status = status
		w.wroteHeader = true
	}
	w.ResponseWriter.WriteHeader(status)
}

func (w *StatusWriter) Write(b []byte) (int, error) {
	const maxBodySize = 4096

	if !w.wroteHeader {
		w.Status = http.StatusOK
		w.wroteHeader = true
	}

	if w.Status >= http.StatusBadRequest {
		// This is technically wrong as multiple calls to write
		// will simply overwrite w.ResponseBody but given that
		// we typically only write to the response body once
		// and this field is only used for logging I'm leaving
		// this as-is.
		w.responseBody = make([]byte, minInt(len(b), maxBodySize))
		copy(w.responseBody, b)
	}

	return w.ResponseWriter.Write(b)
}

func minInt(a, b int) int {
	if a < b {
		return a
	}
	return b
}

// Unwrap returns the underlying ResponseWriter, allowing
// http.ResponseController to reach it for SetWriteDeadline, etc.
func (w *StatusWriter) Unwrap() http.ResponseWriter {
	return w.ResponseWriter
}

func (w *StatusWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
	hijacker, ok := w.ResponseWriter.(http.Hijacker)
	if !ok {
		return nil, nil, xerrors.Errorf("%T is not a http.Hijacker", w.ResponseWriter)
	}
	w.Hijacked = true

	return hijacker.Hijack()
}

func (w *StatusWriter) ResponseBody() []byte {
	return w.responseBody
}

func (w *StatusWriter) Flush() {
	f, ok := w.ResponseWriter.(http.Flusher)
	if !ok {
		panic("http.ResponseWriter is not http.Flusher")
	}
	f.Flush()
}

func getStackString(skip int) string {
	// Get up to 5 callers, skipping this one and the skip count.
	pcs := make([]uintptr, 5)
	got := runtime.Callers(skip+1, pcs)
	frames := runtime.CallersFrames(pcs[:got])

	callers := []string{}
	for {
		frame, more := frames.Next()
		callers = append(callers, fmt.Sprintf("%s:%v", frame.File, frame.Line))
		if !more {
			break
		}
	}
	return strings.Join(callers, " -> ")
}