rdev/internal/logging/middleware.go

package logging

import (
	"net/http"
	"time"

	"github.com/google/uuid"
)

// responseWriter wraps http.ResponseWriter to capture status code.
type responseWriter struct {
	http.ResponseWriter
	status       int
	wroteHeader  bool
	bytesWritten int
}

func (rw *responseWriter) WriteHeader(code int) {
	if rw.wroteHeader {
		return
	}
	rw.status = code
	rw.wroteHeader = true
	rw.ResponseWriter.WriteHeader(code)
}

func (rw *responseWriter) Write(b []byte) (int, error) {
	if !rw.wroteHeader {
		rw.WriteHeader(http.StatusOK)
	}
	n, err := rw.ResponseWriter.Write(b)
	rw.bytesWritten += n
	return n, err
}

// Flush implements http.Flusher, required for SSE streaming through middleware chains.
func (rw *responseWriter) Flush() {
	if f, ok := rw.ResponseWriter.(http.Flusher); ok {
		f.Flush()
	}
}

// Unwrap returns the original http.ResponseWriter, required for http.ResponseController.
func (rw *responseWriter) Unwrap() http.ResponseWriter {
	return rw.ResponseWriter
}

// RequestIDHeader is the header name for request ID propagation.
const RequestIDHeader = "X-Request-ID"

// MiddlewareConfig configures the logging middleware.
type MiddlewareConfig struct {
	// Logger is the logger to use. If nil, uses Default().
	Logger *Logger

	// SkipPaths are paths that should not be logged.
	SkipPaths map[string]bool

	// GenerateRequestID controls whether to generate a request ID if not present.
	GenerateRequestID bool
}

// DefaultMiddlewareConfig returns sensible defaults.
func DefaultMiddlewareConfig() MiddlewareConfig {
	return MiddlewareConfig{
		SkipPaths: map[string]bool{
			"/health": true,
			"/ready":  true,
		},
		GenerateRequestID: true,
	}
}

// Middleware returns an HTTP middleware that logs requests and enriches context.
// It:
//   - Generates or propagates request IDs
//   - Logs request start (debug) and completion (info/warn/error based on status)
//   - Stores an enriched logger in context for handlers to use
func Middleware(cfg MiddlewareConfig) func(http.Handler) http.Handler {
	logger := cfg.Logger
	if logger == nil {
		logger = Default()
	}

	return func(next http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			// Skip logging for configured paths
			if cfg.SkipPaths[r.URL.Path] {
				next.ServeHTTP(w, r)
				return
			}

			start := time.Now()

			// Get or generate request ID
			requestID := r.Header.Get(RequestIDHeader)
			if requestID == "" && cfg.GenerateRequestID {
				requestID = uuid.New().String()
			}

			// Set request ID in response header
			if requestID != "" {
				w.Header().Set(RequestIDHeader, requestID)
			}

			// Wrap response writer to capture status code and bytes
			wrapped := &responseWriter{
				ResponseWriter: w,
				status:         http.StatusOK,
			}

			// Create request-scoped logger
			reqLogger := logger.With(
				FieldRequestID, requestID,
				FieldHTTPMethod, r.Method,
				FieldHTTPPath, r.URL.Path,
				FieldHTTPRemoteAddr, r.RemoteAddr,
			)

			// Store logger in context for handlers to use
			ctx := WithContext(r.Context(), reqLogger)

			// Log request start at debug level
			reqLogger.Debug("request started",
				FieldHTTPUserAgent, r.UserAgent(),
				FieldHTTPHost, r.Host,
			)

			// Call next handler with enriched context
			next.ServeHTTP(wrapped, r.WithContext(ctx))

			// Calculate duration
			duration := time.Since(start)

			// Log completion with appropriate level based on status
			attrs := []any{
				FieldHTTPStatus, wrapped.status,
				FieldDuration, duration.Milliseconds(),
				"bytes", wrapped.bytesWritten,
			}

			switch {
			case wrapped.status >= 500:
				reqLogger.Error("request completed", attrs...)
			case wrapped.status >= 400:
				reqLogger.Warn("request completed", attrs...)
			default:
				reqLogger.Info("request completed", attrs...)
			}
		})
	}
}