story/internal/httpapi/middleware.go

package httpapi

import (
	"context"
	"crypto/rand"
	"encoding/hex"
	"log/slog"
	"net"
	"net/http"
	"runtime/debug"
	"strconv"
	"strings"
	"time"

	"data-terminal-cloud/internal/deviceauth"
	"data-terminal-cloud/internal/identity"

	"go.opentelemetry.io/otel/trace"
)

type contextKey string

const requestIDKey contextKey = "request-id"

type RateLimitResult struct {
	Allowed    bool
	Remaining  int
	RetryAfter time.Duration
}

type RateLimiter interface {
	Allow(context.Context, string, string) (RateLimitResult, error)
}

func RequestID(ctx context.Context) string {
	value, _ := ctx.Value(requestIDKey).(string)
	return value
}

func Middleware(logger *slog.Logger, serviceName string, maxBodyBytes int64, limiter RateLimiter, next http.Handler) http.Handler {
	return requestIDMiddleware(
		securityHeadersMiddleware(
			recoveryMiddleware(logger,
				accessLogMiddleware(logger,
					metricsMiddleware(serviceName,
						rateLimitMiddleware(logger, limiter,
							bodyLimitMiddleware(maxBodyBytes, next),
						),
					),
				),
			),
		),
	)
}

func rateLimitMiddleware(logger *slog.Logger, limiter RateLimiter, next http.Handler) http.Handler {
	if limiter == nil {
		return next
	}
	return http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
		if request.URL.Path == "/metrics" || strings.HasPrefix(request.URL.Path, "/health/") {
			next.ServeHTTP(writer, request)
			return
		}
		identifier := request.RemoteAddr
		if host, _, err := net.SplitHostPort(request.RemoteAddr); err == nil {
			identifier = host
		}
		result, err := limiter.Allow(request.Context(), "network", identifier)
		if err != nil {
			logger.Warn("rate limiter unavailable; request allowed", "request_id", RequestID(request.Context()), "error", err)
			next.ServeHTTP(writer, request)
			return
		}
		writer.Header().Set("X-RateLimit-Remaining", strconv.Itoa(result.Remaining))
		if !result.Allowed {
			retrySeconds := int64((result.RetryAfter + time.Second - 1) / time.Second)
			if retrySeconds < 1 {
				retrySeconds = 1
			}
			writer.Header().Set("Retry-After", strconv.FormatInt(retrySeconds, 10))
			WriteError(writer, request, http.StatusTooManyRequests, "RATE_LIMITED", "请求过于频繁，请稍后重试")
			return
		}
		next.ServeHTTP(writer, request)
	})
}

// UserRateLimit applies independently keyed limits after the bearer token has
// been verified. Tenant limiting is also applied when the API request carries
// the tenant selection header used by tenant-scoped endpoints.
func UserRateLimit(logger *slog.Logger, limiter RateLimiter, next http.Handler) http.Handler {
	if limiter == nil {
		return next
	}
	return http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
		principal, ok := identity.PrincipalFromContext(request.Context())
		if !ok {
			next.ServeHTTP(writer, request)
			return
		}
		scopes := [][2]string{{"user", principal.Issuer + "\x00" + principal.Subject}}
		if tenantID := strings.TrimSpace(request.Header.Get("X-Tenant-Id")); tenantID != "" {
			scopes = append(scopes, [2]string{"tenant", tenantID})
		}
		if !allowScopes(writer, request, logger, limiter, scopes) {
			return
		}
		next.ServeHTTP(writer, request)
	})
}

// DeviceRateLimit applies device and tenant buckets only after the signed
// device request has been verified, so neither identifier can be spoofed to
// evade the network-address bucket.
func DeviceRateLimit(logger *slog.Logger, limiter RateLimiter, next http.Handler) http.Handler {
	if limiter == nil {
		return next
	}
	return http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
		deviceIdentity, ok := deviceauth.FromContext(request.Context())
		if !ok {
			next.ServeHTTP(writer, request)
			return
		}
		if !allowScopes(writer, request, logger, limiter, [][2]string{
			{"device", deviceIdentity.ID},
			{"tenant", deviceIdentity.TenantID},
		}) {
			return
		}
		next.ServeHTTP(writer, request)
	})
}

func allowScopes(writer http.ResponseWriter, request *http.Request, logger *slog.Logger, limiter RateLimiter, scopes [][2]string) bool {
	for _, item := range scopes {
		result, err := limiter.Allow(request.Context(), item[0], item[1])
		if err != nil {
			logger.Warn("rate limiter unavailable; scope allowed", "request_id", RequestID(request.Context()), "scope", item[0], "error", err)
			continue
		}
		writer.Header().Set("X-RateLimit-Remaining-"+strings.Title(item[0]), strconv.Itoa(result.Remaining)) //nolint:staticcheck // fixed protocol header casing
		if result.Allowed {
			continue
		}
		retrySeconds := int64((result.RetryAfter + time.Second - 1) / time.Second)
		if retrySeconds < 1 {
			retrySeconds = 1
		}
		writer.Header().Set("Retry-After", strconv.FormatInt(retrySeconds, 10))
		writer.Header().Set("X-RateLimit-Scope", item[0])
		WriteError(writer, request, http.StatusTooManyRequests, "RATE_LIMITED", "请求过于频繁，请稍后重试")
		return false
	}
	return true
}

func requestIDMiddleware(next http.Handler) http.Handler {
	return http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
		requestID := strings.TrimSpace(request.Header.Get("X-Request-Id"))
		if requestID == "" || len(requestID) > 128 {
			requestID = newRequestID()
		}
		writer.Header().Set("X-Request-Id", requestID)
		ctx := context.WithValue(request.Context(), requestIDKey, requestID)
		next.ServeHTTP(writer, request.WithContext(ctx))
	})
}

func securityHeadersMiddleware(next http.Handler) http.Handler {
	return http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
		writer.Header().Set("X-Content-Type-Options", "nosniff")
		writer.Header().Set("X-Frame-Options", "DENY")
		writer.Header().Set("Referrer-Policy", "no-referrer")
		next.ServeHTTP(writer, request)
	})
}

func recoveryMiddleware(logger *slog.Logger, next http.Handler) http.Handler {
	return http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
		defer func() {
			if recovered := recover(); recovered != nil {
				logger.Error("request panic", "request_id", RequestID(request.Context()), "panic", recovered, "stack", string(debug.Stack()))
				WriteError(writer, request, http.StatusInternalServerError, "INTERNAL_ERROR", "服务暂时不可用")
			}
		}()
		next.ServeHTTP(writer, request)
	})
}

func accessLogMiddleware(logger *slog.Logger, next http.Handler) http.Handler {
	return http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
		startedAt := time.Now()
		defer func() {
			spanContext := trace.SpanContextFromContext(request.Context())
			attributes := []any{
				"request_id", RequestID(request.Context()),
				"method", request.Method,
				"route", request.Pattern,
				"duration_ms", time.Since(startedAt).Milliseconds(),
			}
			if spanContext.IsValid() {
				attributes = append(attributes,
					"trace_id", spanContext.TraceID().String(),
					"span_id", spanContext.SpanID().String(),
				)
			}
			logger.Info("http request", attributes...)
		}()
		next.ServeHTTP(writer, request)
	})
}

func bodyLimitMiddleware(maxBodyBytes int64, next http.Handler) http.Handler {
	return http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
		if request.ContentLength > maxBodyBytes {
			WriteError(writer, request, http.StatusRequestEntityTooLarge, "REQUEST_BODY_TOO_LARGE", "请求正文超出允许大小")
			return
		}
		request.Body = http.MaxBytesReader(writer, request.Body, maxBodyBytes)
		next.ServeHTTP(writer, request)
	})
}

func newRequestID() string {
	bytes := make([]byte, 16)
	if _, err := rand.Read(bytes); err != nil {
		return time.Now().UTC().Format("20060102T150405.000000000")
	}
	return hex.EncodeToString(bytes)
}