routewatch/internal/server/middleware.go

package server

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"log/slog"
	"net/http"
	"sync"
	"time"
)

// responseWriter wraps http.ResponseWriter to capture the response
type responseWriter struct {
	http.ResponseWriter
	body       *bytes.Buffer
	statusCode int
	written    bool
	mu         sync.Mutex
}

func (rw *responseWriter) Write(b []byte) (int, error) {
	rw.mu.Lock()
	defer rw.mu.Unlock()

	if !rw.written {
		rw.written = true
	}

	return rw.body.Write(b)
}

func (rw *responseWriter) WriteHeader(statusCode int) {
	rw.mu.Lock()
	defer rw.mu.Unlock()

	if !rw.written {
		rw.statusCode = statusCode
		rw.written = true
	}
}

func (rw *responseWriter) Header() http.Header {
	return rw.ResponseWriter.Header()
}

// JSONResponseMiddleware is an HTTP middleware that wraps all JSON responses
// with a @meta field containing execution metadata. The metadata includes the
// time zone (always UTC), API version, and request execution time in milliseconds.
//
// Endpoints "/" and "/status" are excluded from this processing and passed through
// unchanged. Non-JSON responses and empty responses are also passed through unchanged.
func JSONResponseMiddleware(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		// Skip non-JSON endpoints
		if r.URL.Path == "/" || r.URL.Path == "/status" {
			next.ServeHTTP(w, r)

			return
		}

		startTime := time.Now()

		// Create a custom response writer to capture the response
		rw := &responseWriter{
			ResponseWriter: w,
			body:           &bytes.Buffer{},
			statusCode:     http.StatusOK,
		}

		// Serve the request
		next.ServeHTTP(rw, r)

		// Calculate execution time
		execTime := time.Since(startTime)

		// Only process JSON responses
		contentType := rw.Header().Get("Content-Type")
		if contentType != "application/json" || rw.body.Len() == 0 {
			// Write the original response
			w.WriteHeader(rw.statusCode)
			_, _ = w.Write(rw.body.Bytes())

			return
		}

		// Parse the original response
		var originalResponse map[string]interface{}
		if err := json.Unmarshal(rw.body.Bytes(), &originalResponse); err != nil {
			// If we can't parse it, just write original
			w.WriteHeader(rw.statusCode)
			_, _ = w.Write(rw.body.Bytes())

			return
		}

		// Add @meta field
		originalResponse["@meta"] = map[string]interface{}{
			"time_zone":      "UTC",
			"api_version":    1,
			"execution_time": fmt.Sprintf("%d ms", execTime.Milliseconds()),
		}

		// Write the enhanced response
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(rw.statusCode)
		_ = json.NewEncoder(w).Encode(originalResponse)
	})
}

// requestTimeoutMessage is the error message sent with 504 Gateway Timeout.
const requestTimeoutMessage = "Request timeout"

// JSONValidationMiddleware is an HTTP middleware that validates JSON API responses.
// It ensures that responses with Content-Type "application/json" contain valid JSON.
//
// If a response is not valid JSON or is empty when JSON is expected, the middleware
// returns a properly formatted JSON error response. If the request's time limit
// has passed, it answers 504 with the error message "Request timeout". For other
// errors, it returns "Internal server error" with status 500 if the original
// status was 200.
//
// Non-JSON responses are passed through unchanged.
func JSONValidationMiddleware(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		// Create a custom response writer to capture the response
		rw := &responseWriter{
			ResponseWriter: w,
			body:           &bytes.Buffer{},
			statusCode:     http.StatusOK,
		}

		// Serve the request
		next.ServeHTTP(rw, r)

		// Check if it's meant to be a JSON response
		contentType := rw.Header().Get("Content-Type")
		isJSON := contentType == "application/json" || contentType == ""

		// If it's not JSON or has content, pass through
		if !isJSON && rw.body.Len() > 0 {
			w.WriteHeader(rw.statusCode)
			_, _ = w.Write(rw.body.Bytes())

			return
		}

		// For JSON responses, validate the JSON
		if rw.body.Len() > 0 {
			var testParse interface{}
			if err := json.Unmarshal(rw.body.Bytes(), &testParse); err == nil {
				// Valid JSON, write it out
				w.Header().Set("Content-Type", "application/json")
				w.WriteHeader(rw.statusCode)
				_, _ = w.Write(rw.body.Bytes())

				return
			}
		}

		// If we get here, either there's no body or invalid JSON
		// Write a proper error response
		w.Header().Set("Content-Type", "application/json")

		// Determine appropriate status code
		statusCode := rw.statusCode
		if statusCode == http.StatusOK {
			statusCode = http.StatusInternalServerError
		}
		if errors.Is(r.Context().Err(), context.DeadlineExceeded) {
			statusCode = http.StatusGatewayTimeout
		}
		w.WriteHeader(statusCode)

		errorMsg := "Internal server error"
		if statusCode == http.StatusGatewayTimeout {
			errorMsg = requestTimeoutMessage
		}

		response := map[string]interface{}{
			"status": "error",
			"error": map[string]interface{}{
				"msg":  errorMsg,
				"code": statusCode,
			},
		}
		_ = json.NewEncoder(w).Encode(response)
	})
}

// statusWriter wraps http.ResponseWriter to capture the status code
type statusWriter struct {
	http.ResponseWriter
	statusCode int
	written    bool
}

func (sw *statusWriter) WriteHeader(statusCode int) {
	if !sw.written {
		sw.statusCode = statusCode
		sw.written = true
	}
	sw.ResponseWriter.WriteHeader(statusCode)
}

func (sw *statusWriter) Write(b []byte) (int, error) {
	if !sw.written {
		sw.statusCode = http.StatusOK
		sw.written = true
	}

	return sw.ResponseWriter.Write(b)
}

// RequestLoggerMiddleware creates a structured logging middleware using slog.
func RequestLoggerMiddleware(logger *slog.Logger) func(http.Handler) http.Handler {
	return func(next http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			start := time.Now()

			// Wrap response writer to capture status
			sw := &statusWriter{ResponseWriter: w, statusCode: http.StatusOK}

			// Log request start
			logger.Debug("HTTP request started",
				"method", r.Method,
				"path", r.URL.Path,
				"remote_addr", r.RemoteAddr,
				"user_agent", r.UserAgent(),
			)

			// Serve the request
			next.ServeHTTP(sw, r)

			// Log request completion
			duration := time.Since(start)
			logLevel := slog.LevelInfo

			// Slow query threshold (1 second)
			const slowQueryThreshold = 1 * time.Second

			if sw.statusCode >= http.StatusInternalServerError {
				logLevel = slog.LevelError
			} else if sw.statusCode >= http.StatusBadRequest {
				logLevel = slog.LevelWarn
			} else if duration >= slowQueryThreshold {
				logLevel = slog.LevelWarn
			}

			logger.Log(r.Context(), logLevel, "HTTP request completed",
				"method", r.Method,
				"path", r.URL.Path,
				"status", sw.statusCode,
				"duration_ms", duration.Milliseconds(),
				"remote_addr", r.RemoteAddr,
				"slow", duration >= slowQueryThreshold,
			)
		})
	}
}