From d1dcdb63c0af567f779d14c2d69a7af8e66b5dbf Mon Sep 17 00:00:00 2001
From: Colin Adler <colin1adler@gmail.com>
Date: Fri, 2 Sep 2022 12:21:11 -0500
Subject: [PATCH] feat: add unique ids to all HTTP requests

Helpful for identifying all logs, traces, and audit logs that correspond
to a single HTTP request.
---
 coderd/audit/request.go         |  2 ++
 coderd/coderd.go                |  1 +
 coderd/httpmw/requestid.go      | 34 +++++++++++++++++++++++++++++++++
 coderd/httpmw/requestid_test.go | 33 ++++++++++++++++++++++++++++++++
 4 files changed, 70 insertions(+)
 create mode 100644 coderd/httpmw/requestid.go
 create mode 100644 coderd/httpmw/requestid_test.go

diff --git a/coderd/audit/request.go b/coderd/audit/request.go
index d547a2d3f4c48..7fcd383287e10 100644
--- a/coderd/audit/request.go
+++ b/coderd/audit/request.go
@@ -12,6 +12,7 @@ import (
 	"cdr.dev/slog"
 	"github.com/coder/coder/coderd/database"
 	"github.com/coder/coder/coderd/httpapi"
+	"github.com/coder/coder/coderd/httpmw"
 )
 
 type RequestParams struct {
@@ -69,6 +70,7 @@ func InitRequest[T Auditable](w http.ResponseWriter, p *RequestParams) (*Request
 			Action:         p.Action,
 			Diff:           diffRaw,
 			StatusCode:     int32(sw.Status),
+			RequestID:      httpmw.RequestID(p.Request),
 		})
 		if err != nil {
 			p.Log.Error(ctx, "export audit log", slog.Error(err))
diff --git a/coderd/coderd.go b/coderd/coderd.go
index 52305ba256623..b4042197d1c2f 100644
--- a/coderd/coderd.go
+++ b/coderd/coderd.go
@@ -155,6 +155,7 @@ func New(options *Options) *API {
 	apiKeyMiddleware := httpmw.ExtractAPIKey(options.Database, oauthConfigs, false)
 
 	r.Use(
+		httpmw.AttachRequestID,
 		httpmw.Recover(api.Logger),
 		httpmw.Logger(api.Logger),
 		httpmw.Prometheus(options.PrometheusRegistry),
diff --git a/coderd/httpmw/requestid.go b/coderd/httpmw/requestid.go
new file mode 100644
index 0000000000000..f2d7b2234d178
--- /dev/null
+++ b/coderd/httpmw/requestid.go
@@ -0,0 +1,34 @@
+package httpmw
+
+import (
+	"context"
+	"net/http"
+
+	"github.com/google/uuid"
+
+	"cdr.dev/slog"
+)
+
+type requestIDContextKey struct{}
+
+// RequestID returns the ID of the request.
+func RequestID(r *http.Request) uuid.UUID {
+	rid, ok := r.Context().Value(requestIDContextKey{}).(uuid.UUID)
+	if !ok {
+		panic("developer error: request id middleware not provided")
+	}
+	return rid
+}
+
+// AttachRequestID adds a request ID to each HTTP request.
+func AttachRequestID(next http.Handler) http.Handler {
+	return http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) {
+		rid := uuid.New()
+
+		ctx := context.WithValue(r.Context(), requestIDContextKey{}, rid)
+		ctx = slog.With(ctx, slog.F("request_id", rid))
+
+		rw.Header().Set("X-Coder-Request-Id", rid.String())
+		next.ServeHTTP(rw, r.WithContext(ctx))
+	})
+}
diff --git a/coderd/httpmw/requestid_test.go b/coderd/httpmw/requestid_test.go
new file mode 100644
index 0000000000000..c632dbbde8c4e
--- /dev/null
+++ b/coderd/httpmw/requestid_test.go
@@ -0,0 +1,33 @@
+package httpmw_test
+
+import (
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/go-chi/chi/v5"
+	"github.com/stretchr/testify/require"
+
+	"github.com/coder/coder/coderd/httpmw"
+)
+
+func TestRequestID(t *testing.T) {
+	t.Parallel()
+
+	rtr := chi.NewRouter()
+	rtr.Use(httpmw.AttachRequestID)
+	rtr.Get("/", func(w http.ResponseWriter, r *http.Request) {
+		rid := httpmw.RequestID(r)
+		w.WriteHeader(http.StatusOK)
+		w.Write([]byte(rid.String()))
+	})
+	r := httptest.NewRequest("GET", "/", nil)
+	rw := httptest.NewRecorder()
+	rtr.ServeHTTP(rw, r)
+
+	res := rw.Result()
+	defer res.Body.Close()
+	require.Equal(t, http.StatusOK, res.StatusCode)
+	require.NotEmpty(t, res.Header.Get("X-Coder-Request-ID"))
+	require.NotEmpty(t, rw.Body.Bytes())
+}