httputil: add response recorder

This implements a new-style response recording `http.ResponseWriter`.

Signed-off-by: Hank Donnay <[email protected]>

Author: hdonnay

internal/httputil/responserecorder.go

@@ -0,0 +1,50 @@
+package httputil
+
+import "net/http"
+
+// ResponseRecorder returns a ResponseWriter that records the HTTP status and
+// body length into the provided pointers, and returns another response writer
+// that understand the go 1.20 http `Unwrap` scheme.
+func ResponseRecorder(status *int, length *int64, w http.ResponseWriter) http.ResponseWriter {
+	// Handle nils being passed, just to be nice.
+	if length == nil {
+		length = new(int64)
+	}
+	if status == nil {
+		status = new(int)
+	}
+	return &responseRecord{
+		ResponseWriter: w,
+		status:         status,
+		length:         length,
+	}
+}
+
+var _ http.ResponseWriter = (*responseRecord)(nil)
+
+type responseRecord struct {
+	http.ResponseWriter
+	status    *int
+	length    *int64
+	writecall bool
+}
+
+func (r *responseRecord) Unwrap() http.ResponseWriter {
+	return r.ResponseWriter
+}
+
+func (r *responseRecord) WriteHeader(c int) {
+	if r.writecall {
+		return
+	}
+	*r.status = c
+	r.ResponseWriter.WriteHeader(c)
+	r.writecall = true
+}
+
+func (r *responseRecord) Write(b []byte) (int, error) {
+	r.WriteHeader(http.StatusOK)
+	n, err := r.ResponseWriter.Write(b)
+	*r.length += int64(n)
+	return n, err
+}

internal/httputil/responserecorder_test.go

@@ -0,0 +1,27 @@
+package httputil
+
+import (
+	"net/http"
+	"net/http/httptest"
+	"testing"
+)
+
+func TestResponseRecorder(t *testing.T) {
+	var status int
+	var length int64
+
+	rec := httptest.NewRecorder()
+	w := ResponseRecorder(&status, &length, rec)
+
+	sz := 512
+	if n, err := w.Write(make([]byte, sz)); err != nil || n != sz {
+		t.Errorf("unexpected Write return: (%v, %v)", n, err)
+	}
+	t.Logf("wrote %d bytes, status %q", length, http.StatusText(status))
+	if got, want := status, http.StatusOK; got != want {
+		t.Errorf("bad status; got: %d, want: %d", got, want)
+	}
+	if got, want := length, int64(sz); got != want {
+		t.Errorf("bad length; got: %d, want: %d", got, want)
+	}
+}