Warp/server/internal/httputil/middleware.go

package httputil

import (
	"log/slog"
	"net/http"
	"slices"
	"strconv"
	"strings"
	"time"
)

// Middleware wraps an http.Handler.
type Middleware func(http.Handler) http.Handler

// Chain applies middleware so that the first element is the outermost.
func Chain(h http.Handler, mw ...Middleware) http.Handler {
	for i := len(mw) - 1; i >= 0; i-- {
		h = mw[i](h)
	}
	return h
}

// ResponseRecorder captures status and size for logging and metrics.
type ResponseRecorder struct {
	http.ResponseWriter
	Status int
	Bytes  int
}

// WriteHeader records the status before passing it on.
func (r *ResponseRecorder) WriteHeader(status int) {
	if r.Status == 0 {
		r.Status = status
	}
	r.ResponseWriter.WriteHeader(status)
}

func (r *ResponseRecorder) Write(b []byte) (int, error) {
	if r.Status == 0 {
		r.Status = http.StatusOK
	}
	n, err := r.ResponseWriter.Write(b)
	r.Bytes += n
	return n, err
}

// Unwrap lets http.ResponseController reach the underlying writer.
func (r *ResponseRecorder) Unwrap() http.ResponseWriter { return r.ResponseWriter }

// CORSOptions configures the CORS middleware.
type CORSOptions struct {
	// AllowedOrigins is a list of origins or the single element "*".
	AllowedOrigins []string
}

// corsHeaders are the headers the Matrix spec tells a homeserver to allow.
const (
	corsMethods = "GET, POST, PUT, DELETE, OPTIONS"
	corsHeaders = "X-Requested-With, Content-Type, Authorization, Accept, Origin, If-Match, If-None-Match"
	corsExpose  = "ETag, Location, Retry-After"
	corsMaxAge  = "86400"
)

// CORS answers preflight requests and adds the access-control headers to
// every response. The Matrix spec expects homeservers to allow all origins;
// a narrower list is honoured for deployments that only serve their own
// first-party clients.
func CORS(opts CORSOptions) Middleware {
	allowAll := len(opts.AllowedOrigins) == 0 || slices.Contains(opts.AllowedOrigins, "*")

	return func(next http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			origin := r.Header.Get("Origin")
			switch {
			case allowAll:
				w.Header().Set("Access-Control-Allow-Origin", "*")
			case origin != "" && slices.Contains(opts.AllowedOrigins, origin):
				w.Header().Set("Access-Control-Allow-Origin", origin)
				w.Header().Add("Vary", "Origin")
			default:
				w.Header().Add("Vary", "Origin")
			}
			w.Header().Set("Access-Control-Allow-Methods", corsMethods)
			w.Header().Set("Access-Control-Allow-Headers", corsHeaders)
			w.Header().Set("Access-Control-Expose-Headers", corsExpose)

			if r.Method == http.MethodOptions {
				w.Header().Set("Access-Control-Max-Age", corsMaxAge)
				w.WriteHeader(http.StatusOK)
				return
			}
			next.ServeHTTP(w, r)
		})
	}
}

// RequestLogger logs one line per request and puts a request-scoped logger
// on the context.
func RequestLogger(base *slog.Logger) Middleware {
	return func(next http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			start := time.Now()
			rec := &ResponseRecorder{ResponseWriter: w}

			l := base.With("method", r.Method, "path", r.URL.Path)
			next.ServeHTTP(rec, r.WithContext(WithLogger(r.Context(), l)))

			status := rec.Status
			if status == 0 {
				status = http.StatusOK
			}
			level := slog.LevelInfo
			switch {
			case status >= 500:
				level = slog.LevelError
			case status >= 400:
				level = slog.LevelWarn
			case isNoise(r.URL.Path):
				level = slog.LevelDebug
			}
			l.LogAttrs(r.Context(), level, "request",
				slog.Int("status", status),
				slog.Int("bytes", rec.Bytes),
				slog.Duration("duration", time.Since(start)),
				slog.String("remote", remoteIP(r)),
			)
		})
	}
}

// isNoise keeps health checks and metrics scrapes out of the info log.
func isNoise(path string) bool {
	return path == "/healthz" || path == "/metrics"
}

func remoteIP(r *http.Request) string {
	host := r.RemoteAddr
	if i := strings.LastIndex(host, ":"); i > 0 {
		host = host[:i]
	}
	return strings.Trim(host, "[]")
}

// Recoverer turns a panic in a handler into a logged M_UNKNOWN response so
// one bad request cannot take the process down.
func Recoverer(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		defer func() {
			if v := recover(); v != nil {
				if v == http.ErrAbortHandler {
					panic(v)
				}
				LoggerFrom(r.Context()).Error("panic in handler",
					"panic", v, "method", r.Method, "path", r.URL.Path)
				WriteJSON(w, http.StatusInternalServerError, InternalError())
			}
		}()
		next.ServeHTTP(w, r)
	})
}

// NoCache marks a response as uncacheable; used for /_matrix/client/versions
// and the well-known documents during development.
func NoCache(w http.ResponseWriter) {
	w.Header().Set("Cache-Control", "no-store")
}

// CacheFor sets a public max-age, used for the well-known documents.
func CacheFor(w http.ResponseWriter, d time.Duration) {
	w.Header().Set("Cache-Control", "public, max-age="+strconv.Itoa(int(d.Seconds())))
}

// NotFoundHandler answers anything unrouted with M_UNRECOGNIZED, which is
// what the Matrix spec requires for unimplemented endpoints.
func NotFoundHandler() http.Handler {
	return Wrap(func(http.ResponseWriter, *http.Request) error {
		return Unrecognized()
	})
}