routewatch/internal/server/middleware.go

package server

import (
	"bytes"
	"context"
	"crypto/rand"
	"encoding/json"
	"errors"
	"fmt"
	"log/slog"
	"net/http"
	"net/netip"
	"slices"
	"strings"
	"sync"
	"time"

	"github.com/go-chi/chi/v5"
)

// 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]any

		err := json.Unmarshal(rw.body.Bytes(), &originalResponse)
		if 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]any{
			"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"

// internalErrorMessage is the error message sent with 500 Internal Server
// Error, whatever went wrong; the handler logs the detail.
const internalErrorMessage = "Internal server error"

// deadlinePassed reports whether ctx ended because its deadline passed.
//
// A handler calls it right after each database read or host name lookup,
// before it looks at what came back, error or not, and answers 504 if it
// reports true. A read cut off by the time limit can return an error, or part
// of its rows or none of them without one, so only this check tells a cut-off
// read from a complete one.
func deadlinePassed(ctx context.Context) bool {
	return errors.Is(ctx.Err(), context.DeadlineExceeded)
}

// 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". A 4xx
// keeps its status, with the status's standard text as the message. 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 any

			err := json.Unmarshal(rw.body.Bytes(), &testParse)
			if 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 deadlinePassed(r.Context()) {
			statusCode = http.StatusGatewayTimeout
		}

		w.WriteHeader(statusCode)

		errorMsg := internalErrorMessage

		switch {
		case statusCode == http.StatusGatewayTimeout:
			errorMsg = requestTimeoutMessage
		case statusCode >= http.StatusBadRequest &&
			statusCode < http.StatusInternalServerError:
			errorMsg = http.StatusText(statusCode)
		}

		response := map[string]any{
			jsonKeyStatus: jsonStatusError,
			jsonKeyError: map[string]any{
				jsonKeyMsg:  errorMsg,
				jsonKeyCode: 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

			switch {
			case sw.statusCode >= http.StatusInternalServerError:
				logLevel = slog.LevelError
			case sw.statusCode >= http.StatusBadRequest:
				logLevel = slog.LevelWarn
			case 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,
			)
		})
	}
}

// DecodedPathMiddleware makes the router match routes against the decoded
// path, so that the handlers read their URL parameters decoded. Without it,
// chi matches against the escaped path whenever the request has one and hands
// the handlers their parameters still escaped: /ip/2001%3Adb8%3A%3A1 would
// look up "2001%3Adb8%3A%3A1".
func DecodedPathMiddleware(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		chi.RouteContext(r.Context()).RoutePath = r.URL.Path

		next.ServeHTTP(w, r)
	})
}

// The values of the headers SecurityHeadersMiddleware sets.
const (
	// strictTransportSecurity tells browsers to reach this host and its
	// subdomains only over HTTPS for a year. The server listens on plain
	// HTTP behind a TLS-terminating proxy, which passes the header on.
	strictTransportSecurity = "max-age=31536000; includeSubDomains"

	// contentSecurityPolicy lets a page load content only from this server,
	// and run or apply an inline script or style block only when the block
	// carries the nonce made for that answer, which replaces %[1]s. No other
	// site may frame the pages.
	contentSecurityPolicy = "default-src 'self'; " +
		"script-src 'nonce-%[1]s'; style-src 'nonce-%[1]s'; " +
		"frame-ancestors 'none'; base-uri 'none'; form-action 'self'"

	frameOptions       = "DENY"
	contentTypeOptions = "nosniff"
	referrerPolicy     = "strict-origin-when-cross-origin"

	// permissionsPolicy turns off browser features the pages do not use.
	permissionsPolicy = "accelerometer=(), camera=(), geolocation=(), " +
		"gyroscope=(), magnetometer=(), microphone=(), payment=(), usb=()"
)

// cspNonceKey is the request context key of the nonce SecurityHeadersMiddleware
// made for the answer.
type cspNonceKey struct{}

// SecurityHeadersMiddleware sets the security headers REPO_POLICIES.md requires
// on every answer. Each answer's Content-Security-Policy names a nonce of its
// own; a page reads it with cspNonce and writes it on its script and style
// blocks.
func SecurityHeadersMiddleware(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		nonce := rand.Text()

		header := w.Header()
		header.Set("Strict-Transport-Security", strictTransportSecurity)
		header.Set("Content-Security-Policy",
			fmt.Sprintf(contentSecurityPolicy, nonce))
		header.Set("X-Frame-Options", frameOptions)
		header.Set("X-Content-Type-Options", contentTypeOptions)
		header.Set("Referrer-Policy", referrerPolicy)
		header.Set("Permissions-Policy", permissionsPolicy)

		ctx := context.WithValue(r.Context(), cspNonceKey{}, nonce)
		next.ServeHTTP(w, r.WithContext(ctx))
	})
}

// cspNonce returns the nonce SecurityHeadersMiddleware made for the answer to
// the request with ctx.
func cspNonce(ctx context.Context) string {
	nonce, _ := ctx.Value(cspNonceKey{}).(string)

	return nonce
}

// RealIPMiddleware gives a request from one of trustedProxies the address of
// the client the proxies forwarded it for as its RemoteAddr: the last
// X-Forwarded-For entry that is not a trusted proxy. Each proxy appends the
// address it took the request from, so the entries left of that one are
// whatever the client sent. A request from any other peer keeps the peer's
// address, whatever headers it carries.
func RealIPMiddleware(trustedProxies []netip.Prefix) func(http.Handler) http.Handler {
	isTrusted := func(addr netip.Addr) bool {
		return slices.ContainsFunc(trustedProxies, func(proxy netip.Prefix) bool {
			return proxy.Contains(addr.Unmap())
		})
	}

	return func(next http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			peer, err := netip.ParseAddrPort(r.RemoteAddr)
			if err != nil || !isTrusted(peer.Addr()) {
				next.ServeHTTP(w, r)

				return
			}

			// A header sent on several lines is one list.
			entries := strings.Split(
				strings.Join(r.Header.Values("X-Forwarded-For"), ","), ",")

			for _, entry := range slices.Backward(entries) {
				addr, err := netip.ParseAddr(strings.TrimSpace(entry))
				if err != nil {
					break
				}

				r.RemoteAddr = addr.String()

				if !isTrusted(addr) {
					break
				}
			}

			next.ServeHTTP(w, r)
		})
	}
}