feat: add panic recovery middleware (#3687)

coderd/httpapi/httpapi.go

@@ -59,6 +59,18 @@ func Forbidden(rw http.ResponseWriter) {
 	})
 }
 
+func InternalServerError(rw http.ResponseWriter, err error) {
+	var details string
+	if err != nil {
+		details = err.Error()
+	}
+
+	Write(rw, http.StatusInternalServerError, codersdk.Response{
+		Message: "An internal server error occurred.",
+		Detail:  details,
+	})
+}
+
 // Write outputs a standardized format to an HTTP response body.
 func Write(rw http.ResponseWriter, status int, response interface{}) {
 	buf := &bytes.Buffer{}

coderd/httpapi/httpapi_test.go

@@ -10,11 +10,46 @@ import (
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+	"golang.org/x/xerrors"
 
 	"github.com/coder/coder/coderd/httpapi"
 	"github.com/coder/coder/codersdk"
 )
 
+func TestInternalServerError(t *testing.T) {
+	t.Parallel()
+
+	t.Run("NoError", func(t *testing.T) {
+		t.Parallel()
+		w := httptest.NewRecorder()
+		httpapi.InternalServerError(w, nil)
+
+		var resp codersdk.Response
+		err := json.NewDecoder(w.Body).Decode(&resp)
+		require.NoError(t, err)
+		require.Equal(t, http.StatusInternalServerError, w.Code)
+		require.NotEmpty(t, resp.Message)
+		require.Empty(t, resp.Detail)
+	})
+
+	t.Run("WithError", func(t *testing.T) {
+		t.Parallel()
+		var (
+			w       = httptest.NewRecorder()
+			httpErr = xerrors.New("error!")
+		)
+
+		httpapi.InternalServerError(w, httpErr)
+
+		var resp codersdk.Response
+		err := json.NewDecoder(w.Body).Decode(&resp)
+		require.NoError(t, err)
+		require.Equal(t, http.StatusInternalServerError, w.Code)
+		require.NotEmpty(t, resp.Message)
+		require.Equal(t, httpErr.Error(), resp.Detail)
+	})
+}
+
 func TestWrite(t *testing.T) {
 	t.Parallel()
 	t.Run("NoErrors", func(t *testing.T) {

coderd/httpapi/request.go

@@ -0,0 +1,30 @@
+package httpapi
+
+import "net/http"
+
+const (
+	// XForwardedHostHeader is a header used by proxies to indicate the
+	// original host of the request.
+	XForwardedHostHeader = "X-Forwarded-Host"
+)
+
+// RequestHost returns the name of the host from the request.  It prioritizes
+// 'X-Forwarded-Host' over r.Host since most requests are being proxied.
+func RequestHost(r *http.Request) string {
+	host := r.Header.Get(XForwardedHostHeader)
+	if host != "" {
+		return host
+	}
+
+	return r.Host
+}
+
+func IsWebsocketUpgrade(r *http.Request) bool {
+	vs := r.Header.Values("Upgrade")
+	for _, v := range vs {
+		if v == "websocket" {
+			return true
+		}
+	}
+	return false
+}

coderd/httpapi/status_writer.go

@@ -0,0 +1,74 @@
+package httpapi
+
+import (
+	"bufio"
+	"net"
+	"net/http"
+
+	"golang.org/x/xerrors"
+)
+
+var _ http.ResponseWriter = (*StatusWriter)(nil)
+var _ 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
+}
+
+func (w *StatusWriter) WriteHeader(status int) {
+	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
+}
+
+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
+}

coderd/httpapi/status_writer_test.go

@@ -0,0 +1,129 @@
+package httpapi_test
+
+import (
+	"bufio"
+	"crypto/rand"
+	"net"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+	"golang.org/x/xerrors"
+
+	"github.com/coder/coder/coderd/httpapi"
+)
+
+func TestStatusWriter(t *testing.T) {
+	t.Parallel()
+
+	t.Run("WriteHeader", func(t *testing.T) {
+		t.Parallel()
+
+		var (
+			rec = httptest.NewRecorder()
+			w   = &httpapi.StatusWriter{ResponseWriter: rec}
+		)
+
+		w.WriteHeader(http.StatusOK)
+		require.Equal(t, http.StatusOK, w.Status)
+		// Validate that the code is written to the underlying Response.
+		require.Equal(t, http.StatusOK, rec.Code)
+	})
+
+	t.Run("WriteHeaderTwice", func(t *testing.T) {
+		t.Parallel()
+
+		var (
+			rec  = httptest.NewRecorder()
+			w    = &httpapi.StatusWriter{ResponseWriter: rec}
+			code = http.StatusNotFound
+		)
+
+		w.WriteHeader(code)
+		w.WriteHeader(http.StatusOK)
+		// Validate that we only record the first status code.
+		require.Equal(t, code, w.Status)
+		// Validate that the code is written to the underlying Response.
+		require.Equal(t, code, rec.Code)
+	})
+
+	t.Run("WriteNoHeader", func(t *testing.T) {
+		t.Parallel()
+		var (
+			rec  = httptest.NewRecorder()
+			w    = &httpapi.StatusWriter{ResponseWriter: rec}
+			body = []byte("hello")
+		)
+
+		_, err := w.Write(body)
+		require.NoError(t, err)
+
+		// Should set the status to OK.
+		require.Equal(t, http.StatusOK, w.Status)
+		// We don't record the body for codes <400.
+		require.Equal(t, []byte(nil), w.ResponseBody())
+		require.Equal(t, body, rec.Body.Bytes())
+	})
+
+	t.Run("WriteAfterHeader", func(t *testing.T) {
+		t.Parallel()
+		var (
+			rec  = httptest.NewRecorder()
+			w    = &httpapi.StatusWriter{ResponseWriter: rec}
+			body = []byte("hello")
+			code = http.StatusInternalServerError
+		)
+
+		w.WriteHeader(code)
+		_, err := w.Write(body)
+		require.NoError(t, err)
+
+		require.Equal(t, code, w.Status)
+		require.Equal(t, body, w.ResponseBody())
+		require.Equal(t, body, rec.Body.Bytes())
+	})
+
+	t.Run("WriteMaxBody", func(t *testing.T) {
+		t.Parallel()
+		var (
+			rec = httptest.NewRecorder()
+			w   = &httpapi.StatusWriter{ResponseWriter: rec}
+			// 8kb body.
+			body = make([]byte, 8<<10)
+			code = http.StatusInternalServerError
+		)
+
+		_, err := rand.Read(body)
+		require.NoError(t, err)
+
+		w.WriteHeader(code)
+		_, err = w.Write(body)
+		require.NoError(t, err)
+
+		require.Equal(t, code, w.Status)
+		require.Equal(t, body, rec.Body.Bytes())
+		require.Equal(t, body[:4096], w.ResponseBody())
+	})
+
+	t.Run("Hijack", func(t *testing.T) {
+		t.Parallel()
+		var (
+			rec = httptest.NewRecorder()
+		)
+
+		w := &httpapi.StatusWriter{ResponseWriter: hijacker{rec}}
+
+		_, _, err := w.Hijack()
+		require.Error(t, err)
+		require.Equal(t, "hijacked", err.Error())
+	})
+}
+
+type hijacker struct {
+	http.ResponseWriter
+}
+
+func (hijacker) Hijack() (net.Conn, *bufio.ReadWriter, error) {
+	return nil, nil, xerrors.New("hijacked")
+}