AIRouter/internal/middleware/middleware.go

package middleware

import (
	"compress/gzip"
	"context"
	"fmt"
	"net/http"
	"runtime/debug"
	"strings"
	"time"

	"github.com/edgeai/gateway/internal/observability"
	"github.com/google/uuid"
)

type contextKey string

const (
	RequestIDKey contextKey = "request_id"
)

// CORS middleware 添加跨域响应头，支持浏览器客户端直接调用 API。
// allowedOrigins 为允许的源列表，"*" 表示允许所有源。
func CORS(allowedOrigins []string) func(http.Handler) http.Handler {
	allowed := map[string]bool{}
	for _, o := range allowedOrigins {
		allowed[o] = true
	}
	allowAll := allowed["*"]

	return func(next http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			origin := r.Header.Get("Origin")
			if origin != "" {
				if allowAll || allowed[origin] {
					w.Header().Set("Access-Control-Allow-Origin", origin)
					w.Header().Set("Access-Control-Allow-Methods", "GET, POST, PUT, PATCH, DELETE, OPTIONS")
					w.Header().Set("Access-Control-Allow-Headers", "Authorization, Content-Type, X-Request-ID")
					w.Header().Set("Access-Control-Expose-Headers", "X-Request-ID, X-Trace-Id, X-RateLimit-Limit, X-RateLimit-Remaining, Retry-After, X-Backpressure-Level, X-Timing-Queue-Ms, X-Timing-Inference-Ms, X-Timing-Total-Ms")
					w.Header().Set("Access-Control-Max-Age", "3600")
				}
			}
			// 处理预检请求
			if r.Method == http.MethodOptions {
				w.WriteHeader(http.StatusNoContent)
				return
			}
			next.ServeHTTP(w, r)
		})
	}
}

// RequestID middleware generates a unique request ID and sets it in context and response header.
// 同时设置 X-Trace-Id 响应头，便于客户端和分布式追踪系统关联请求。
func RequestID(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		requestID := r.Header.Get("X-Request-ID")
		if requestID == "" {
			requestID = uuid.New().String()
		}
		w.Header().Set("X-Request-ID", requestID)
		w.Header().Set("X-Trace-Id", requestID)
		ctx := context.WithValue(r.Context(), RequestIDKey, requestID)
		next.ServeHTTP(w, r.WithContext(ctx))
	})
}

// BodyLimit middleware rejects requests with bodies exceeding the given size.
func BodyLimit(maxMB int) func(http.Handler) http.Handler {
	maxBytes := int64(maxMB) * 1024 * 1024
	return func(next http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			if r.ContentLength > maxBytes {
				w.Header().Set("Content-Type", "application/json")
				w.WriteHeader(http.StatusBadRequest)
				fmt.Fprintf(w, `{"error":{"code":"INVALID_REQUEST","message":"request body exceeds %dMB limit","request_id":"%s"}}`, maxMB, r.Header.Get("X-Request-ID"))
				return
			}
			r.Body = http.MaxBytesReader(w, r.Body, maxBytes)
			next.ServeHTTP(w, r)
		})
	}
}

// Recovery middleware catches panics and returns 500.
func Recovery(logger *observability.Logger) func(http.Handler) http.Handler {
	return func(next http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			defer func() {
				if rec := recover(); rec != nil {
					logger.Error("panic recovered",
						observability.F().Event("panic").
							RequestID(r.Header.Get("X-Request-ID")).
							Reason(fmt.Sprintf("%v\n%s", rec, debug.Stack())))
					w.Header().Set("Content-Type", "application/json")
					w.WriteHeader(http.StatusInternalServerError)
					fmt.Fprintf(w, `{"error":{"code":"INTERNAL_ERROR","message":"internal server error","request_id":"%s"}}`, r.Header.Get("X-Request-ID"))
				}
			}()
			next.ServeHTTP(w, r)
		})
	}
}

// Logging middleware logs request method, path, status, and duration.
func Logging(logger *observability.Logger) func(http.Handler) http.Handler {
	return func(next http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			start := time.Now()
			rw := &responseWriter{ResponseWriter: w, status: 200}
			next.ServeHTTP(rw, r)
			fields := observability.F().
				Event("http_request").
				RequestID(r.Header.Get("X-Request-ID")).
				Set("method", r.Method).
				Set("path", r.URL.Path).
				Set("status", rw.status).
				Set("duration_ms", time.Since(start).Milliseconds()).
				Set("response_bytes", rw.bytesWritten)
			if origin := r.Header.Get("Origin"); origin != "" {
				fields.Set("origin", origin)
			}
			logger.Info("http request", fields)
		})
	}
}

type responseWriter struct {
	http.ResponseWriter
	status       int
	bytesWritten int
}

func (rw *responseWriter) WriteHeader(code int) {
	rw.status = code
	rw.ResponseWriter.WriteHeader(code)
}

func (rw *responseWriter) Write(b []byte) (int, error) {
	n, err := rw.ResponseWriter.Write(b)
	rw.bytesWritten += n
	return n, err
}

func (rw *responseWriter) Flush() {
	if f, ok := rw.ResponseWriter.(http.Flusher); ok {
		f.Flush()
	}
}

// GetRequestID extracts the request ID from context.
func GetRequestID(ctx context.Context) string {
	if v, ok := ctx.Value(RequestIDKey).(string); ok {
		return v
	}
	return ""
}

// RateLimitFn 是限流检查函数的类型。
// 返回 allowed, limit, remaining。
type RateLimitFn func(appID string) (allowed bool, limit int, remaining int)

// RateLimit middleware 对所有认证请求执行 per-app 限流，并设置 X-RateLimit-* 响应头。
// skipPaths 中的路径不限流（如 /health、/ready）。
func RateLimit(checkFn RateLimitFn, getIdentity func(*http.Request) string, skipPaths map[string]bool) func(http.Handler) http.Handler {
	return func(next http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			if skipPaths[r.URL.Path] {
				next.ServeHTTP(w, r)
				return
			}
			appID := getIdentity(r)
			if appID == "" {
				next.ServeHTTP(w, r)
				return
			}
			allowed, limit, remaining := checkFn(appID)
			w.Header().Set("X-RateLimit-Limit", fmt.Sprintf("%d", limit))
			w.Header().Set("X-RateLimit-Remaining", fmt.Sprintf("%d", remaining))
			if !allowed {
				w.Header().Set("Retry-After", "1")
				w.Header().Set("Content-Type", "application/json")
				w.WriteHeader(http.StatusTooManyRequests)
				fmt.Fprintf(w, `{"error":{"code":"RATE_LIMITED","message":"rate limit exceeded for this application","request_id":"%s"}}`, r.Header.Get("X-Request-ID"))
				return
			}
			next.ServeHTTP(w, r)
		})
	}
}

// gzipResponseWriter 包装 ResponseWriter，对响应体进行 gzip 压缩。
// 仅当客户端发送 Accept-Encoding: gzip 且响应 Content-Type 为 JSON 时启用。
type gzipResponseWriter struct {
	http.ResponseWriter
	gz          *gzip.Writer
	contentType string
	gzStarted   bool
}

func (g *gzipResponseWriter) WriteHeader(code int) {
	ct := g.ResponseWriter.Header().Get("Content-Type")
	g.contentType = ct
	// SSE 流式响应不压缩
	if strings.Contains(ct, "text/event-stream") {
		g.ResponseWriter.WriteHeader(code)
		return
	}
	// 对 JSON 等可压缩内容启用 gzip
	if shouldCompress(ct) {
		g.ResponseWriter.Header().Set("Content-Encoding", "gzip")
		g.ResponseWriter.Header().Del("Content-Length")
		g.gzStarted = true
	}
	g.ResponseWriter.WriteHeader(code)
}

func (g *gzipResponseWriter) Write(b []byte) (int, error) {
	if g.gzStarted {
		return g.gz.Write(b)
	}
	return g.ResponseWriter.Write(b)
}

func (g *gzipResponseWriter) Flush() {
	if g.gzStarted {
		g.gz.Flush()
	}
	if f, ok := g.ResponseWriter.(http.Flusher); ok {
		f.Flush()
	}
}

// shouldCompress 判断 Content-Type 是否值得压缩。
func shouldCompress(ct string) bool {
	return strings.Contains(ct, "json") ||
		strings.Contains(ct, "text") ||
		strings.Contains(ct, "javascript") ||
		strings.Contains(ct, "xml")
}

// Gzip 中间件对响应体进行 gzip 压缩，减少网络传输量。
// 仅当客户端支持 gzip 且响应内容类型可压缩时启用。
func Gzip(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if !strings.Contains(r.Header.Get("Accept-Encoding"), "gzip") {
			next.ServeHTTP(w, r)
			return
		}

		gz := gzip.NewWriter(w)
		defer gz.Close()

		gw := &gzipResponseWriter{ResponseWriter: w, gz: gz}
		next.ServeHTTP(gw, r)
	})
}