Merge pull request #370 from 3rabiii/feat/request-id-tracking

feat: implement Request ID middleware for log correlation

Author: aaronbrethorst

.gitignore

@@ -14,3 +14,4 @@ config.docker.json
 tmp/
 *.log
 .DS_Store
+/api

cmd/api/app.go

@@ -13,7 +13,6 @@ import (
 	"time"
 
 	"github.com/prometheus/client_golang/prometheus/promhttp"
-
 	"maglev.onebusaway.org/internal/app"
 	"maglev.onebusaway.org/internal/appconf"
 	"maglev.onebusaway.org/internal/clock"
@@ -118,7 +117,8 @@ func CreateServer(coreApp *app.Application, cfg appconf.Config) (*http.Server, *
 	// Add request logging middleware (outermost)
 	requestLogger := logging.NewStructuredLogger(os.Stdout, slog.LevelInfo)
 	requestLogMiddleware := restapi.NewRequestLoggingMiddleware(requestLogger)
-	handler := requestLogMiddleware(metricsHandler)
+
+	handler := restapi.RequestIDMiddleware(requestLogMiddleware(metricsHandler))
 
 	srv := &http.Server{
 		Addr:         fmt.Sprintf(":%d", cfg.Port),

go.mod

@@ -5,6 +5,7 @@ go 1.24.2
 require (
 	github.com/OneBusAway/go-gtfs v1.1.0
 	github.com/davecgh/go-spew v1.1.1
+	github.com/google/uuid v1.6.0
 	github.com/klauspost/compress v1.18.0
 	github.com/mattn/go-sqlite3 v1.14.24
 	github.com/prometheus/client_golang v1.23.2
@@ -25,7 +26,6 @@ require (
 	github.com/fatih/structtag v1.2.0 // indirect
 	github.com/go-sql-driver/mysql v1.9.3 // indirect
 	github.com/google/cel-go v0.26.1 // indirect
-	github.com/google/uuid v1.6.0 // indirect
 	github.com/inconshreveable/mousetrap v1.1.0 // indirect
 	github.com/jackc/pgpassfile v1.0.0 // indirect
 	github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 // indirect

internal/restapi/request_id_middleware.go

@@ -0,0 +1,31 @@
+package restapi
+
+import (
+	"context"
+	"net/http"
+	"regexp"
+
+	"github.com/google/uuid"
+)
+
+type contextKey string
+
+const RequestIDKey contextKey = "request_id"
+
+var validRequestIDRegex = regexp.MustCompile(`^[a-zA-Z0-9-._:]+$`)
+
+func RequestIDMiddleware(next http.Handler) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		reqID := r.Header.Get("X-Request-ID")
+
+		if reqID == "" || len(reqID) > 128 || !validRequestIDRegex.MatchString(reqID) {
+			reqID = uuid.NewString()
+		}
+
+		w.Header().Set("X-Request-ID", reqID)
+
+		ctx := context.WithValue(r.Context(), RequestIDKey, reqID)
+
+		next.ServeHTTP(w, r.WithContext(ctx))
+	})
+}

internal/restapi/request_id_middleware_test.go

@@ -0,0 +1,86 @@
+package restapi
+
+import (
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestRequestIDMiddleware(t *testing.T) {
+	t.Run("should generate request ID if missing", func(t *testing.T) {
+		nextHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			reqID, ok := r.Context().Value(RequestIDKey).(string)
+			assert.True(t, ok, "Context should contain request ID")
+			assert.NotEmpty(t, reqID, "Request ID should not be empty")
+		})
+
+		handlerToTest := RequestIDMiddleware(nextHandler)
+
+		req := httptest.NewRequest("GET", "http://example.com/foo", nil)
+		rec := httptest.NewRecorder()
+
+		handlerToTest.ServeHTTP(rec, req)
+
+		respID := rec.Header().Get("X-Request-ID")
+		assert.NotEmpty(t, respID, "Response header should contain X-Request-ID")
+		assert.Regexp(t, `^[0-9a-f-]{36}$`, respID)
+	})
+
+	t.Run("should preserve existing valid request ID", func(t *testing.T) {
+		existingID := "my-custom-trace-id-123"
+
+		nextHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			reqID, ok := r.Context().Value(RequestIDKey).(string)
+			assert.True(t, ok)
+			assert.Equal(t, existingID, reqID)
+		})
+
+		handlerToTest := RequestIDMiddleware(nextHandler)
+
+		req := httptest.NewRequest("GET", "http://example.com/foo", nil)
+		req.Header.Set("X-Request-ID", existingID)
+		rec := httptest.NewRecorder()
+
+		handlerToTest.ServeHTTP(rec, req)
+
+		assert.Equal(t, existingID, rec.Header().Get("X-Request-ID"))
+	})
+
+	t.Run("should replace invalid request ID", func(t *testing.T) {
+		testCases := []struct {
+			name      string
+			invalidID string
+		}{
+			{
+				name:      "ID too long (>128 chars)",
+				invalidID: strings.Repeat("a", 129),
+			},
+			{
+				name:      "ID contains invalid characters",
+				invalidID: "bad-id-<script>",
+			},
+		}
+
+		for _, tc := range testCases {
+			t.Run(tc.name, func(t *testing.T) {
+				nextHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+					reqID, ok := r.Context().Value(RequestIDKey).(string)
+					assert.True(t, ok)
+					assert.NotEqual(t, tc.invalidID, reqID)
+					assert.Regexp(t, `^[0-9a-f-]{36}$`, reqID)
+				})
+
+				handlerToTest := RequestIDMiddleware(nextHandler)
+
+				req := httptest.NewRequest("GET", "http://example.com/foo", nil)
+				req.Header.Set("X-Request-ID", tc.invalidID)
+				rec := httptest.NewRecorder()
+
+				handlerToTest.ServeHTTP(rec, req)
+			})
+		}
+	})
+}

internal/restapi/request_logging_middleware.go

@@ -41,11 +41,14 @@ func NewRequestLoggingMiddleware(logger *slog.Logger) func(http.Handler) http.Ha
 			// Log the request
 			duration := time.Since(start)
 
+			reqID, _ := r.Context().Value(RequestIDKey).(string)
+
 			logging.LogHTTPRequest(logger,
 				r.Method,
-				r.URL.Path, // Path without query parameters
+				r.URL.Path,
 				wrapped.statusCode,
-				float64(duration.Nanoseconds())/1e6, // Convert to milliseconds
+				float64(duration.Nanoseconds())/1e6,
+				slog.String("request_id", reqID),
 				slog.String("user_agent", r.Header.Get("User-Agent")),
 				slog.String("component", "http_server"))
 		})