BlueArchiveToolkit/internal/api/middleware.go

package api

import (
	"crypto/subtle"
	"net"
	"net/http"
	"net/url"
	"strings"
	"sync"
	"time"
)

func (s *Server) wrapHandler(next http.Handler) http.Handler {
	handler := next
	handler = s.authMiddleware(handler)
	handler = s.rateLimitMiddleware(handler)
	handler = s.securityHeadersMiddleware(handler)
	handler = s.accessLogMiddleware(handler)
	// Check before ServeMux can clean dot segments and route an escaped path
	// to a different host-shaped endpoint.
	handler = rejectDotSegmentsMiddleware(handler)
	return handler
}

func rejectDotSegmentsMiddleware(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if hasDotPathSegment(r.URL.Path) || (r.URL.RawPath != "" && hasDotPathSegment(r.URL.RawPath)) {
			http.NotFound(w, r)
			return
		}
		next.ServeHTTP(w, r)
	})
}

func hasDotPathSegment(rawPath string) bool {
	decoded, err := url.PathUnescape(rawPath)
	if err != nil {
		return true
	}
	for _, segment := range strings.Split(decoded, "/") {
		if segment == "." || segment == ".." {
			return true
		}
	}
	return false
}

func (s *Server) securityHeadersMiddleware(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", "no-referrer")
		w.Header().Set("Permissions-Policy", "geolocation=(), microphone=(), camera=()")
		w.Header().Set("Content-Security-Policy", "default-src 'none'; frame-ancestors 'none'")
		next.ServeHTTP(w, r)
	})
}

func (s *Server) authMiddleware(next http.Handler) http.Handler {
	if s.cfg.AuthToken == "" {
		return next
	}
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if s.authExempt(r.URL.Path) {
			next.ServeHTTP(w, r)
			return
		}
		if !constantTimeTokenEqual(s.requestToken(r), s.cfg.AuthToken) {
			w.Header().Set("WWW-Authenticate", `Bearer realm="bat-api"`)
			writeErrorJSON(w, http.StatusUnauthorized, "unauthorized", "missing or invalid access token")
			return
		}
		next.ServeHTTP(w, r)
	})
}

func (s *Server) authExempt(path string) bool {
	for _, exempt := range s.cfg.AuthExemptPaths {
		if exempt == path {
			return true
		}
		if strings.HasSuffix(exempt, "/") && strings.HasPrefix(path, exempt) {
			return true
		}
	}
	return false
}

func (s *Server) requestToken(r *http.Request) string {
	if auth := r.Header.Get("Authorization"); auth != "" {
		const prefix = "Bearer "
		if strings.HasPrefix(auth, prefix) {
			return strings.TrimSpace(strings.TrimPrefix(auth, prefix))
		}
	}
	for _, header := range []string{"X-BAT-Token", "X-BAT-API-Key"} {
		if token := strings.TrimSpace(r.Header.Get(header)); token != "" {
			return token
		}
	}
	if s.cfg.AuthQueryParam != "" {
		return r.URL.Query().Get(s.cfg.AuthQueryParam)
	}
	return ""
}

func constantTimeTokenEqual(got, want string) bool {
	if got == "" || want == "" {
		return false
	}
	gotBytes := []byte(got)
	wantBytes := []byte(want)
	if len(gotBytes) != len(wantBytes) {
		subtle.ConstantTimeCompare(wantBytes, wantBytes)
		return false
	}
	return subtle.ConstantTimeCompare(gotBytes, wantBytes) == 1
}

func (s *Server) rateLimitMiddleware(next http.Handler) http.Handler {
	if s.limiter == nil {
		return next
	}
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		client := s.clientIdentity(r)
		if !s.limiter.allow(client, time.Now()) {
			w.Header().Set("Retry-After", "1")
			writeErrorJSON(w, http.StatusTooManyRequests, "rate_limited", "too many requests")
			return
		}
		next.ServeHTTP(w, r)
	})
}

func (s *Server) accessLogMiddleware(next http.Handler) http.Handler {
	if !s.cfg.AccessLog {
		return next
	}
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		start := time.Now()
		rec := &statusRecorder{ResponseWriter: w, status: http.StatusOK}
		next.ServeHTTP(rec, r)
		requestID := safeLogValue(r.Header.Get("X-Request-ID"))
		if requestID == "" {
			requestID = "-"
		}
		s.logger.Printf(
			"access method=%s path=%s status=%d bytes=%d duration_ms=%d client_ip=%s request_id=%s user_agent=%q",
			r.Method,
			r.URL.Path,
			rec.status,
			rec.bytes,
			time.Since(start).Milliseconds(),
			s.clientIdentity(r),
			requestID,
			r.UserAgent(),
		)
	})
}

func (s *Server) clientIdentity(r *http.Request) string {
	if s.cfg.TrustProxyHeaders {
		if forwarded := r.Header.Get("X-Forwarded-For"); forwarded != "" {
			for _, part := range strings.Split(forwarded, ",") {
				if ip := strings.TrimSpace(part); ip != "" {
					return ip
				}
			}
		}
		if realIP := strings.TrimSpace(r.Header.Get("X-Real-IP")); realIP != "" {
			return realIP
		}
	}
	host, _, err := net.SplitHostPort(r.RemoteAddr)
	if err == nil && host != "" {
		return host
	}
	return r.RemoteAddr
}

func safeLogValue(raw string) string {
	raw = strings.TrimSpace(raw)
	if raw == "" {
		return ""
	}
	raw = strings.Map(func(r rune) rune {
		if r < 32 || r == 127 {
			return -1
		}
		return r
	}, raw)
	if len(raw) > 128 {
		return raw[:128]
	}
	return raw
}

type statusRecorder struct {
	http.ResponseWriter
	status int
	bytes  int
}

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

func (r *statusRecorder) Write(data []byte) (int, error) {
	n, err := r.ResponseWriter.Write(data)
	r.bytes += n
	return n, err
}

func (r *statusRecorder) Unwrap() http.ResponseWriter {
	return r.ResponseWriter
}

type tokenBucketLimiter struct {
	mu          sync.Mutex
	rate        float64
	burst       float64
	buckets     map[string]*tokenBucket
	lastCleanup time.Time
}

type tokenBucket struct {
	tokens   float64
	last     time.Time
	lastSeen time.Time
}

func newTokenBucketLimiter(rps float64, burst int) *tokenBucketLimiter {
	return &tokenBucketLimiter{
		rate:    rps,
		burst:   float64(burst),
		buckets: map[string]*tokenBucket{},
	}
}

func (l *tokenBucketLimiter) allow(key string, now time.Time) bool {
	if key == "" {
		key = "unknown"
	}
	l.mu.Lock()
	defer l.mu.Unlock()
	if l.lastCleanup.IsZero() || now.Sub(l.lastCleanup) > time.Minute {
		l.cleanup(now)
	}
	bucket := l.buckets[key]
	if bucket == nil {
		bucket = &tokenBucket{tokens: l.burst, last: now, lastSeen: now}
		l.buckets[key] = bucket
	}
	elapsed := now.Sub(bucket.last).Seconds()
	bucket.tokens += elapsed * l.rate
	if bucket.tokens > l.burst {
		bucket.tokens = l.burst
	}
	bucket.last = now
	bucket.lastSeen = now
	if bucket.tokens < 1 {
		return false
	}
	bucket.tokens--
	return true
}

func (l *tokenBucketLimiter) cleanup(now time.Time) {
	for key, bucket := range l.buckets {
		if now.Sub(bucket.lastSeen) > 10*time.Minute {
			delete(l.buckets, key)
		}
	}
	l.lastCleanup = now
}

func writeErrorJSON(w http.ResponseWriter, status int, code string, message string) {
	writeNoStoreJSON(w, status, ErrorResponse{
		Error: ErrorBody{
			Code:    code,
			Message: message,
			Status:  status,
		},
	})
}