From 8599b2451ca17de39a35ad4c0e02610f7ae4fa7e Mon Sep 17 00:00:00 2001
From: Axel Gustav <dev@axelgustav.de>
Date: Mon, 18 Nov 2024 16:09:24 -0400
Subject: [PATCH] Adds middleware to insert X-Request-Start unix millis
 timestamp

---
 internal/server/request_start_middleware.go   | 28 ++++++++++++++
 .../server/request_start_middleware_test.go   | 38 +++++++++++++++++++
 internal/server/server.go                     |  2 +
 3 files changed, 68 insertions(+)
 create mode 100644 internal/server/request_start_middleware.go
 create mode 100644 internal/server/request_start_middleware_test.go

diff --git a/internal/server/request_start_middleware.go b/internal/server/request_start_middleware.go
new file mode 100644
index 0000000..2469671
--- /dev/null
+++ b/internal/server/request_start_middleware.go
@@ -0,0 +1,28 @@
+package server
+
+import (
+	"net/http"
+	"strconv"
+	"time"
+)
+
+const (
+	requestStartHeader = "X-Request-Start"
+)
+
+type RequestStartMiddleware struct {
+	next http.Handler
+}
+
+func WithRequestStartMiddleware(next http.Handler) http.Handler {
+	return &RequestStartMiddleware{
+		next: next,
+	}
+}
+
+func (h *RequestStartMiddleware) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+	if r.Header.Get(requestStartHeader) == "" {
+		r.Header.Set(requestStartHeader, strconv.FormatInt(time.Now().UnixMilli(), 10))
+	}
+	h.next.ServeHTTP(w, r)
+}
diff --git a/internal/server/request_start_middleware_test.go b/internal/server/request_start_middleware_test.go
new file mode 100644
index 0000000..f1b7db4
--- /dev/null
+++ b/internal/server/request_start_middleware_test.go
@@ -0,0 +1,38 @@
+package server
+
+import (
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestRequestStartMiddleware_AddsUnixMilliWhenNotPresent(t *testing.T) {
+	handler := WithRequestStartMiddleware(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		timestamp := r.Header.Get(requestStartHeader)
+		assert.NotEmpty(t, timestamp)
+	}))
+
+	r := httptest.NewRequest("GET", "/", nil)
+	w := httptest.NewRecorder()
+	handler.ServeHTTP(w, r)
+
+	assert.Equal(t, http.StatusOK, w.Code)
+}
+
+func TestRequestStartMiddleware_PreservesExistingHeaderWhenPresent(t *testing.T) {
+	next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		timestamp := r.Header.Get(requestStartHeader)
+		assert.Equal(t, "1234", timestamp)
+	})
+	handler := WithRequestStartMiddleware(next)
+
+	r := httptest.NewRequest("GET", "/", nil)
+	r.Header.Set(requestStartHeader, "1234")
+
+	w := httptest.NewRecorder()
+	handler.ServeHTTP(w, r)
+
+	assert.Equal(t, http.StatusOK, w.Code)
+}
diff --git a/internal/server/server.go b/internal/server/server.go
index bcd31f9..361c28a 100644
--- a/internal/server/server.go
+++ b/internal/server/server.go
@@ -119,10 +119,12 @@ func (s *Server) startCommandHandler() error {
 func (s *Server) buildHandler() http.Handler {
 	var handler http.Handler
 
+	// Note: handlers are executed in the inverse order.
 	handler = s.router
 	handler, _ = WithErrorPageMiddleware(pages.DefaultErrorPages, true, handler)
 	handler = WithLoggingMiddleware(slog.Default(), s.config.HttpPort, s.config.HttpsPort, handler)
 	handler = WithRequestIDMiddleware(handler)
+	handler = WithRequestStartMiddleware(handler)
 
 	return handler
 }