Publish ITWorx Pulse source

internal/correlation/correlation.go

@@ -0,0 +1,43 @@
+package correlation
+
+import (
+	"context"
+	"crypto/rand"
+	"encoding/hex"
+	"net/http"
+	"regexp"
+)
+
+const Header = "X-Correlation-ID"
+
+type contextKey struct{}
+
+var validID = regexp.MustCompile(`^[A-Za-z0-9._:-]{8,64}$`)
+
+func New() string {
+	bytes := make([]byte, 16)
+	if _, err := rand.Read(bytes); err != nil {
+		return "correlation-unavailable"
+	}
+	return hex.EncodeToString(bytes)
+}
+
+func FromContext(ctx context.Context) string {
+	value, _ := ctx.Value(contextKey{}).(string)
+	return value
+}
+
+func WithContext(ctx context.Context, id string) context.Context {
+	return context.WithValue(ctx, contextKey{}, id)
+}
+
+func Middleware(next http.Handler) http.Handler {
+	return http.HandlerFunc(func(response http.ResponseWriter, request *http.Request) {
+		id := request.Header.Get(Header)
+		if !validID.MatchString(id) {
+			id = New()
+		}
+		response.Header().Set(Header, id)
+		next.ServeHTTP(response, request.WithContext(WithContext(request.Context(), id)))
+	})
+}

internal/correlation/correlation_test.go

@@ -0,0 +1,39 @@
+package correlation
+
+import (
+	"net/http"
+	"net/http/httptest"
+	"testing"
+)
+
+func TestMiddlewarePreservesValidCorrelationID(t *testing.T) {
+	handler := Middleware(http.HandlerFunc(func(response http.ResponseWriter, request *http.Request) {
+		if got := FromContext(request.Context()); got != "request-123" {
+			t.Errorf("context correlation ID = %q", got)
+		}
+		response.WriteHeader(http.StatusNoContent)
+	}))
+	request := httptest.NewRequest(http.MethodGet, "/", nil)
+	request.Header.Set(Header, "request-123")
+	response := httptest.NewRecorder()
+	handler.ServeHTTP(response, request)
+	if response.Header().Get(Header) != "request-123" {
+		t.Fatalf("response correlation ID = %q", response.Header().Get(Header))
+	}
+}
+
+func TestMiddlewareReplacesInvalidCorrelationID(t *testing.T) {
+	handler := Middleware(http.HandlerFunc(func(response http.ResponseWriter, request *http.Request) {
+		if len(FromContext(request.Context())) < 8 {
+			t.Error("generated correlation ID is too short")
+		}
+		response.WriteHeader(http.StatusNoContent)
+	}))
+	request := httptest.NewRequest(http.MethodGet, "/", nil)
+	request.Header.Set(Header, "secret\nforged")
+	response := httptest.NewRecorder()
+	handler.ServeHTTP(response, request)
+	if response.Header().Get(Header) == "secret\nforged" || response.Header().Get(Header) == "" {
+		t.Fatalf("invalid correlation ID was not replaced: %q", response.Header().Get(Header))
+	}
+}