counters/internal/server/middleware.go

package server

import (
	"crypto/subtle"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log/slog"
	"mime"
	"net"
	"net/http"
	"net/netip"
	"runtime/debug"
	"strconv"
	"strings"
	"time"

	"migadu_counters/internal/api"
)

// maxBody bounds every request body (the biggest legitimate body is a few
// hundred bytes).
const maxBody = 64 << 10

// httpError is a rejection by the validation layer.
type httpError struct {
	status int
	msg    string
}

func (e *httpError) Error() string { return e.msg }

func writeJSON(w http.ResponseWriter, status int, v any) {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(status)
	_ = json.NewEncoder(w).Encode(v)
}

func writeError(w http.ResponseWriter, status int, msg string) {
	writeJSON(w, status, api.ErrorResponse{Error: msg})
}

// reject answers an httpError and logs it: a burst of 4xx from mailqueuer
// means contract drift or a token mismatch.
func (s *Server) reject(w http.ResponseWriter, r *http.Request, route string, e *httpError) {
	writeError(w, e.status, e.msg)
	s.d.Logger.Warn("request rejected", "route", route, "remote", r.RemoteAddr, "status", e.status, "reason", e.msg)
}

// statusWriter captures the status code for metrics.
type statusWriter struct {
	http.ResponseWriter
	status int
}

func (w *statusWriter) WriteHeader(code int) {
	if w.status == 0 {
		w.status = code
	}
	w.ResponseWriter.WriteHeader(code)
}

func (w *statusWriter) Write(p []byte) (int, error) {
	if w.status == 0 {
		w.status = http.StatusOK
	}
	return w.ResponseWriter.Write(p)
}

// recoverer keeps the daemon alive when a handler panics. It sits outside
// http.TimeoutHandler, which re-panics in the caller's goroutine.
func (s *Server) recoverer(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		defer func() {
			if p := recover(); p != nil {
				s.d.Logger.Error("panic in handler", "panic", fmt.Sprint(p), "path", r.URL.Path, "stack", string(debug.Stack()))
				writeError(w, http.StatusInternalServerError, "internal error")
			}
		}()
		next.ServeHTTP(w, r)
	})
}

// observe records request metrics per route and status.
func (s *Server) observe(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		start := time.Now()
		sw := &statusWriter{ResponseWriter: w}
		next.ServeHTTP(sw, r)
		if sw.status == 0 {
			sw.status = http.StatusOK
		}
		route := r.URL.Path
		switch route {
		case "/in", "/out", "/counters":
		default:
			route = "other"
		}
		if s.d.Metrics != nil {
			s.d.Metrics.RecordRequest(route, strconv.Itoa(sw.status), time.Since(start).Seconds())
		}
	})
}

// route applies the validation steps common to every endpoint, in order:
// method, declared body size, bearer token, source IP, content type (POST).
func (s *Server) route(path, method string, h http.HandlerFunc) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.Method != method {
			w.Header().Set("Allow", method)
			s.reject(w, r, path, &httpError{http.StatusMethodNotAllowed, "method not allowed"})
			return
		}
		if r.ContentLength > maxBody {
			s.reject(w, r, path, &httpError{http.StatusRequestEntityTooLarge, "request body too large"})
			return
		}
		if e := s.checkAuth(r); e != nil {
			w.Header().Set("WWW-Authenticate", "Bearer")
			s.reject(w, r, path, e)
			return
		}
		if e := s.checkSource(r); e != nil {
			s.reject(w, r, path, e)
			return
		}
		if method == http.MethodPost {
			if e := checkContentType(r); e != nil {
				s.reject(w, r, path, e)
				return
			}
			r.Body = http.MaxBytesReader(w, r.Body, maxBody)
		}
		h(w, r)
	})
}

func (s *Server) checkAuth(r *http.Request) *httpError {
	if s.d.AuthToken == "" {
		return nil // config.Warnings() flags this at startup
	}
	const prefix = "Bearer "
	h := r.Header.Get("Authorization")
	if !strings.HasPrefix(h, prefix) {
		return &httpError{http.StatusUnauthorized, "unauthorized"}
	}
	if subtle.ConstantTimeCompare([]byte(strings.TrimSpace(h[len(prefix):])), []byte(s.d.AuthToken)) != 1 {
		return &httpError{http.StatusUnauthorized, "unauthorized"}
	}
	return nil
}

// checkSource matches RemoteAddr (never X-Forwarded-For: nothing proxies
// mailqueuer's calls) against allowed_cidrs when configured.
func (s *Server) checkSource(r *http.Request) *httpError {
	if len(s.d.AllowedCIDRs) == 0 {
		return nil
	}
	host, _, err := net.SplitHostPort(r.RemoteAddr)
	if err != nil {
		host = r.RemoteAddr
	}
	addr, err := netip.ParseAddr(host)
	if err != nil {
		return &httpError{http.StatusForbidden, "source address not allowed"}
	}
	addr = addr.Unmap()
	for _, p := range s.d.AllowedCIDRs {
		if p.Contains(addr) {
			return nil
		}
	}
	return &httpError{http.StatusForbidden, "source address not allowed"}
}

func checkContentType(r *http.Request) *httpError {
	mt, _, err := mime.ParseMediaType(r.Header.Get("Content-Type"))
	if err != nil || mt != "application/json" {
		return &httpError{http.StatusUnsupportedMediaType, "Content-Type must be application/json"}
	}
	return nil
}

// decodeJSON reads exactly one JSON object into dst. Unknown fields are
// ignored so mailqueuer can add fields without breaking counting.
func decodeJSON(r *http.Request, dst any) *httpError {
	dec := json.NewDecoder(r.Body)
	if err := dec.Decode(dst); err != nil {
		var tooBig *http.MaxBytesError
		if errors.As(err, &tooBig) {
			return &httpError{http.StatusRequestEntityTooLarge, "request body too large"}
		}
		return &httpError{http.StatusBadRequest, "invalid json: " + err.Error()}
	}
	if err := dec.Decode(&struct{}{}); err != io.EOF {
		var tooBig *http.MaxBytesError
		if errors.As(err, &tooBig) {
			return &httpError{http.StatusRequestEntityTooLarge, "request body too large"}
		}
		return &httpError{http.StatusBadRequest, "body must contain exactly one JSON object"}
	}
	return nil
}

// logger returns the configured logger or slog's default.
func (d Deps) logger() *slog.Logger {
	if d.Logger != nil {
		return d.Logger
	}
	return slog.Default()
}