tiny-forge/internal/api/middleware.go

package api

import (
	"log/slog"
	"net/http"
	"runtime/debug"
	"sync"
	"time"
)

// logging is an HTTP middleware that logs every request with method, path,
// status code, and duration.
func logging(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		start := time.Now()
		wrapped := &statusRecorder{ResponseWriter: w, status: http.StatusOK}

		next.ServeHTTP(wrapped, r)

		slog.Info("http request",
			"method", r.Method,
			"path", r.URL.Path,
			"status", wrapped.status,
			"duration", time.Since(start).String(),
		)
	})
}

// recovery is an HTTP middleware that catches panics and returns a 500 response.
func recovery(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		defer func() {
			if err := recover(); err != nil {
				slog.Error("panic recovered", "error", err, "stack", string(debug.Stack()))
				respondError(w, http.StatusInternalServerError, "internal server error")
			}
		}()
		next.ServeHTTP(w, r)
	})
}

// securityHeaders sets standard security headers on all responses.
func securityHeaders(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("X-Content-Type-Options", "nosniff")
		w.Header().Set("X-Frame-Options", "DENY")
		w.Header().Set("Referrer-Policy", "strict-origin-when-cross-origin")
		w.Header().Set("Content-Security-Policy", "default-src 'self'; script-src 'self'; style-src 'self' 'unsafe-inline'; img-src 'self' data:; connect-src 'self'; font-src 'self'")
		next.ServeHTTP(w, r)
	})
}

// cors is an HTTP middleware that handles CORS for same-origin requests.
// The frontend is served from the same origin, so cross-origin requests are not expected.
func cors(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		// The frontend is served from the same origin, so cross-origin
		// requests are not expected. We do NOT reflect the Origin header
		// back, as that would allow any website to make credentialed requests.
		// If cross-origin support is needed in the future, maintain an
		// explicit allowlist of trusted origins here.

		if r.Method == http.MethodOptions {
			w.WriteHeader(http.StatusNoContent)
			return
		}

		next.ServeHTTP(w, r)
	})
}

// maxBodySize limits request body sizes to prevent memory exhaustion.
const maxBodySize = 1 << 20 // 1 MB

func limitBody(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		r.Body = http.MaxBytesReader(w, r.Body, maxBodySize)
		next.ServeHTTP(w, r)
	})
}

// rateLimiter provides per-IP rate limiting for login endpoints.
type rateLimiter struct {
	mu          sync.Mutex
	attempts    map[string][]time.Time
	lastCleanup time.Time
}

func newRateLimiter() *rateLimiter {
	return &rateLimiter{
		attempts:    make(map[string][]time.Time),
		lastCleanup: time.Now(),
	}
}

// allow checks if the IP is allowed to make another request.
// Returns false if the IP has exceeded the limit (10 requests per minute).
func (rl *rateLimiter) allow(ip string) bool {
	rl.mu.Lock()
	defer rl.mu.Unlock()

	now := time.Now()
	window := now.Add(-1 * time.Minute)

	// Periodically clean all stale IPs to prevent memory leak.
	if now.Sub(rl.lastCleanup) > 5*time.Minute {
		for k, times := range rl.attempts {
			filtered := times[:0]
			for _, t := range times {
				if t.After(window) {
					filtered = append(filtered, t)
				}
			}
			if len(filtered) == 0 {
				delete(rl.attempts, k)
			} else {
				rl.attempts[k] = filtered
			}
		}
		rl.lastCleanup = now
	}

	// Clean old entries for this IP.
	filtered := rl.attempts[ip][:0]
	for _, t := range rl.attempts[ip] {
		if t.After(window) {
			filtered = append(filtered, t)
		}
	}
	rl.attempts[ip] = filtered

	if len(filtered) >= 10 {
		return false
	}

	rl.attempts[ip] = append(rl.attempts[ip], now)
	return true
}

// jsonContentType is an HTTP middleware that sets the default Content-Type to JSON.
func jsonContentType(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		next.ServeHTTP(w, r)
	})
}

// statusRecorder wraps http.ResponseWriter to capture the status code.
type statusRecorder struct {
	http.ResponseWriter
	status int
}

func (r *statusRecorder) WriteHeader(code int) {
	r.status = code
	r.ResponseWriter.WriteHeader(code)
}

// Flush delegates to the underlying ResponseWriter if it supports http.Flusher (needed for SSE).
func (r *statusRecorder) Flush() {
	if f, ok := r.ResponseWriter.(http.Flusher); ok {
		f.Flush()
	}
}