pure-note/internal/middleware/middleware.go

// Package middleware 自写中间件链（§7.2）：
// SecurityHeaders → 日志 → 全局限流 → Origin 校验 →（路由级）MaxBytes → Auth → CSRF。
package middleware

import (
	"context"
	"log/slog"
	"net"
	"net/http"
	"strings"
	"time"
)

// statusWriter 捕获响应状态码供日志使用。
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(b []byte) (int, error) {
	if w.status == 0 {
		w.status = http.StatusOK
	}
	return w.ResponseWriter.Write(b)
}

// SecurityHeaders 下发安全 HTTP 头（§9.2）。
func SecurityHeaders(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		h := w.Header()
		h.Set("Content-Security-Policy",
			"default-src 'self'; "+
				"script-src 'self'; "+
				// 'unsafe-inline' 仅因 CodeMirror 经 style-mod 运行时注入 <style>（§9.2）。
				"style-src 'self' 'unsafe-inline'; "+
				"img-src 'self' data:; "+
				"font-src 'self'; "+
				"connect-src 'self'; "+
				"object-src 'none'; base-uri 'none'; "+
				"frame-ancestors 'none'; form-action 'self'")
		// 站点仅经 HTTPS 反代对外服务（§9.1-T11）
		h.Set("Strict-Transport-Security", "max-age=31536000; includeSubDomains")
		h.Set("Referrer-Policy", "strict-origin-when-cross-origin")
		h.Set("X-Content-Type-Options", "nosniff")
		h.Set("X-Frame-Options", "DENY")
		h.Set("Permissions-Policy", "camera=(), microphone=(), geolocation=()")
		next.ServeHTTP(w, r)
	})
}

// RequestLogger 请求日志：仅记录错误（status ≥ 400）。
func RequestLogger(log *slog.Logger) func(http.Handler) http.Handler {
	return func(next http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			sw := &statusWriter{ResponseWriter: w}
			start := time.Now()
			next.ServeHTTP(sw, r)
			if sw.status >= 400 {
				log.Warn("http",
					"method", r.Method,
					"path", r.URL.Path,
					"status", sw.status,
					"ip", ClientIP(r, trueBehindProxy(r)),
					"duration_ms", time.Since(start).Milliseconds(),
				)
			}
		})
	}
}

// behindProxyKey 由 Server 注入到请求上下文，供日志取 IP 用。
type ctxKey string

const behindProxyKey ctxKey = "behind_proxy"

// BehindProxy 中间件：把「位于可信反代之后」标记注入请求上下文，
// 供日志与限流取客户端 IP 使用。
func BehindProxy(v bool) func(http.Handler) http.Handler {
	return func(next http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			next.ServeHTTP(w, WithBehindProxy(r, v))
		})
	}
}

// WithBehindProxy 标记请求位于可信反代之后。
func WithBehindProxy(r *http.Request, v bool) *http.Request {
	return r.WithContext(context.WithValue(r.Context(), behindProxyKey, v))
}

func trueBehindProxy(r *http.Request) bool {
	v, _ := r.Context().Value(behindProxyKey).(bool)
	return v
}

// ClientIP 提取客户端 IP：behindProxy 时取 X-Forwarded-For 最右条目
// （Caddy 追加语义，§14），否则取 RemoteAddr。
func ClientIP(r *http.Request, behindProxy bool) string {
	if behindProxy {
		if xff := r.Header.Get("X-Forwarded-For"); xff != "" {
			parts := strings.Split(xff, ",")
			return strings.TrimSpace(parts[len(parts)-1])
		}
	}
	host, _, err := net.SplitHostPort(r.RemoteAddr)
	if err != nil {
		return r.RemoteAddr
	}
	return host
}

// MaxBytes 路由级请求体上限中间件（§7.2）。
func MaxBytes(n int64) func(http.Handler) http.Handler {
	return func(next http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			r.Body = http.MaxBytesReader(w, r.Body, n)
			next.ServeHTTP(w, r)
		})
	}
}

// NoStore 为 /api/admin/* 与 /api/auth/* 响应统一附加
// Cache-Control: no-store（防登出后 bfcache 回看，§9.2）。
func NoStore(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Cache-Control", "no-store")
		next.ServeHTTP(w, r)
	})
}