Add prevalent middlewares

.golangci.yml

@@ -1,5 +1,5 @@
 run:
-  deadline: 60s
+  timeout: 1m
 
 linters:
   enable:
@@ -16,14 +16,11 @@ issues:
 
 linters-settings:
   gci:
-    no-inline-comments: true
-    no-prefix-comments: false
     sections:
       - standard
       - default
       - prefix(go.artefactual.dev/tools)
-    section-separators:
-      - newLine
+    custom-order: true
   importas:
     no-unaliased: true
     no-extra-aliases: false

go.mod

@@ -10,6 +10,7 @@ require (
 	github.com/go-logr/logr v1.4.1
 	github.com/go-logr/zapr v1.3.0
 	github.com/google/go-cmp v0.6.0
+	github.com/gorilla/mux v1.8.1
 	go.temporal.io/api v1.29.2
 	go.temporal.io/sdk v1.26.0
 	go.uber.org/mock v0.4.0

go.sum

@@ -123,6 +123,8 @@ github.com/googleapis/enterprise-certificate-proxy v0.3.2 h1:Vie5ybvEvT75RniqhfF
 github.com/googleapis/enterprise-certificate-proxy v0.3.2/go.mod h1:VLSiSSBs/ksPL8kq3OBOQ6WRI2QnaFynd1DCjZ62+V0=
 github.com/googleapis/gax-go/v2 v2.12.2 h1:mhN09QQW1jEWeMF74zGR81R30z4VJzjZsfkUhuHF+DA=
 github.com/googleapis/gax-go/v2 v2.12.2/go.mod h1:61M8vcyyXR2kqKFxKrfA22jaA8JGF7Dc8App1U3H6jc=
+github.com/gorilla/mux v1.8.1 h1:TuBL49tXwgrFYWhqrNgrUNEY92u81SPhu7sTdzQEiWY=
+github.com/gorilla/mux v1.8.1/go.mod h1:AKf9I4AEqPTmMytcMc0KkNouC66V3BtZ4qD5fmWSiMQ=
 github.com/grpc-ecosystem/go-grpc-middleware v1.4.0 h1:UH//fgunKIs4JdUbpDl1VZCDaL56wXCB/5+wF6uHfaI=
 github.com/grpc-ecosystem/go-grpc-middleware v1.4.0/go.mod h1:g5qyo/la0ALbONm6Vbp88Yd8NsDy6rZz+RcrMPxvld8=
 github.com/grpc-ecosystem/grpc-gateway/v2 v2.19.1 h1:/c3QmbOGMGTOumP2iT/rCwB7b0QDGLKzqOmktBjT+Is=

middleware/middleware_test.go

@@ -0,0 +1,41 @@
+package middleware_test
+
+import (
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"testing"
+
+	"github.com/go-logr/logr/funcr"
+	"github.com/gorilla/mux"
+	"gotest.tools/v3/assert"
+
+	"go.artefactual.dev/tools/middleware"
+)
+
+func TestMiddlewares(t *testing.T) {
+	t.Parallel()
+
+	logged := ""
+	logger := funcr.New(func(prefix, args string) { logged += prefix + args }, funcr.Options{})
+
+	panicker := func(w http.ResponseWriter, r *http.Request) { panic("opsie") }
+	router := mux.NewRouter()
+	router.HandleFunc("/", panicker).Methods("GET")
+	router.Use(
+		middleware.Recover(logger),
+		middleware.WriteTimeout(0),
+		middleware.VersionHeader("", "v1.2.3"),
+	)
+
+	rw := httptest.NewRecorder()
+	req := httptest.NewRequest("GET", "/", nil)
+
+	router.ServeHTTP(rw, req)
+
+	// Recover logs the panic error and handles the response.
+	assert.Assert(t, strings.Contains(logged, "Panic error recovered."))
+
+	// VersionHeader injects a header into the response.
+	assert.Equal(t, rw.Header().Get("X-Version"), "v1.2.3")
+}

middleware/recover.go

@@ -0,0 +1,60 @@
+package middleware
+
+import (
+	"bytes"
+	"errors"
+	"net/http"
+	"runtime"
+	"strings"
+
+	"github.com/go-logr/logr"
+)
+
+// Recover from panics and logs the error.
+func Recover(logger logr.Logger) func(http.Handler) http.Handler {
+	return func(h http.Handler) http.Handler {
+		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			defer func() {
+				if rec := recover(); rec != nil {
+					// Don't recover if the request is aborted, as this would
+					// prevent the client from detecting the error.
+					if rec == http.ErrAbortHandler {
+						panic(rec)
+					}
+
+					// Prepare the error message and log it.
+					b := strings.Builder{}
+
+					switch x := rec.(type) {
+					case string:
+						b.WriteString("panic: ")
+						b.WriteString(x)
+					case error:
+						b.WriteString("panic: ")
+						b.WriteString(x.Error())
+					default:
+						b.WriteString("unknown panic")
+					}
+
+					const size = 64 << 10 // 64KB
+					buf := make([]byte, size)
+					n := runtime.Stack(buf, false)
+					lines := bytes.Split(buf[:n], []byte{'\n'})
+					b.WriteByte('\n')
+					for _, line := range lines {
+						b.Write(line)
+						b.WriteByte('\n')
+					}
+
+					logger.Error(errors.New(b.String()), "Panic error recovered.")
+
+					// Skip write header on upgrade connection.
+					if r.Header.Get("Connection") != "Upgrade" {
+						w.WriteHeader(http.StatusInternalServerError)
+					}
+				}
+			}()
+			h.ServeHTTP(w, r)
+		})
+	}
+}

middleware/recover_test.go

@@ -0,0 +1,34 @@
+package middleware_test
+
+import (
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/go-logr/logr/funcr"
+	"gotest.tools/v3/assert"
+	"gotest.tools/v3/assert/cmp"
+
+	"go.artefactual.dev/tools/middleware"
+)
+
+func TestRescoverMiddleware(t *testing.T) {
+	t.Parallel()
+
+	req := httptest.NewRequest("GET", "http://example.com/foo", nil)
+	w := httptest.NewRecorder()
+
+	var logged string
+	logger := funcr.New(
+		func(prefix, args string) { logged = args },
+		funcr.Options{},
+	)
+
+	handler := http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) { panic("opsie") })
+	mw := middleware.Recover(logger)
+
+	mw(handler).ServeHTTP(w, req)
+
+	assert.Assert(t, cmp.Contains(logged, "\"msg\"=\"Panic error recovered.\""))
+	assert.Assert(t, cmp.Contains(logged, "\"error\"=\"panic: opsie"))
+}

middleware/timeout.go

@@ -0,0 +1,22 @@
+package middleware
+
+import (
+	"net/http"
+	"time"
+)
+
+// WriteTimeout sets the write deadline for writing the response. Use 0 to set
+// an unlimited timeout.
+func WriteTimeout(timeout time.Duration) func(http.Handler) http.Handler {
+	return func(h http.Handler) http.Handler {
+		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			rc := http.NewResponseController(w)
+			var deadline time.Time
+			if timeout != 0 {
+				deadline = time.Now().Add(timeout)
+			}
+			_ = rc.SetWriteDeadline(deadline)
+			h.ServeHTTP(w, r)
+		})
+	}
+}

middleware/timeout_test.go

@@ -0,0 +1,46 @@
+package middleware_test
+
+import (
+	"io"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+	"time"
+
+	"gotest.tools/v3/assert"
+
+	"go.artefactual.dev/tools/middleware"
+)
+
+func TestWriteTimeout(t *testing.T) {
+	t.Parallel()
+
+	h := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		time.Sleep(time.Microsecond * 100)
+		w.Write([]byte("Hi there!"))
+	})
+
+	t.Run("Sets a write timeout", func(t *testing.T) {
+		t.Parallel()
+
+		ts := httptest.NewServer(middleware.WriteTimeout(time.Microsecond)(h))
+		defer ts.Close()
+
+		_, err := ts.Client().Get(ts.URL)
+		assert.ErrorIs(t, err, io.EOF)
+	})
+
+	t.Run("Sets an unlimited write timeout", func(t *testing.T) {
+		t.Parallel()
+
+		ts := httptest.NewServer(middleware.WriteTimeout(0)(h))
+		defer ts.Close()
+
+		resp, err := ts.Client().Get(ts.URL)
+		assert.NilError(t, err)
+
+		blob, err := io.ReadAll(resp.Body)
+		assert.NilError(t, err)
+		assert.Equal(t, string(blob), "Hi there!")
+	})
+}

middleware/version.go

@@ -0,0 +1,17 @@
+package middleware
+
+import "net/http"
+
+// VersionHeader sets a version header on the response. If name is empty, it
+// defaults to "X-Version".
+func VersionHeader(name, version string) func(http.Handler) http.Handler {
+	if name == "" {
+		name = "X-Version"
+	}
+	return func(h http.Handler) http.Handler {
+		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			w.Header().Set(name, version)
+			h.ServeHTTP(w, r)
+		})
+	}
+}

middleware/version_test.go

@@ -0,0 +1,28 @@
+package middleware_test
+
+import (
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"gotest.tools/v3/assert"
+
+	"go.artefactual.dev/tools/middleware"
+)
+
+func TestVersionHeaderMiddleware(t *testing.T) {
+	t.Parallel()
+
+	req := httptest.NewRequest("GET", "http://example.com/foo", nil)
+	w := httptest.NewRecorder()
+
+	var continued bool
+	handler := http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) { continued = true })
+	mw := middleware.VersionHeader("", "v1.2.3")
+
+	mw(handler).ServeHTTP(w, req)
+	resp := w.Result()
+
+	assert.Equal(t, resp.Header.Get("X-Version"), "v1.2.3")
+	assert.Equal(t, continued, true)
+}