AI Gateway Go 0.10.0 源码快照 + 旗舰版需求规划报告

M0-M7 已完成:核心网关(身份/RBAC/TOTP/OIDC/SAML/Provider/配额/路由/内容策略/审计/定价)+ 资源市场(MCP/Skills/数字员工)。含 22 个 PostgreSQL 迁移、管理端/门户端前端源码、OpenAPI 契约、部署 compose。 Co-Authored-By: Claude <noreply@anthropic.com>
2026-08-12 11:45:54 +08:00
commit 5759c1862e
807 changed files with 114727 additions and 0 deletions
@@ -0,0 +1,430 @@
+package gateway
+
+import (
+	"context"
+	"crypto/subtle"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"log/slog"
+	"net"
+	"net/http"
+	"net/http/httputil"
+	"strconv"
+	"strings"
+	"sync"
+	"time"
+
+	"aigateway.local/core/internal/apikey"
+	"aigateway.local/core/internal/contentpolicy"
+	"aigateway.local/core/internal/pricing"
+	"aigateway.local/core/internal/provider"
+)
+
+type Proxy struct {
+	resolver   AdapterResolver
+	auth       apikey.KeyAuthenticator
+	maxBody    int64
+	logger     *slog.Logger
+	transport  *http.Transport
+	proxies    sync.Map
+	circuits   sync.Map
+	admission  AdmissionController
+	tokenQuota TokenQuotaController
+	resilience ResiliencePolicy
+	audit      AuditRecorder
+	policies   *contentpolicy.Engine
+	pricing    *pricing.Service
+}
+
+type cachedProxy struct {
+	key   string
+	proxy *httputil.ReverseProxy
+}
+
+var (
+	ErrProviderNotFound    = errors.New("requested provider is not available")
+	ErrProviderUnavailable = errors.New("provider configuration is unavailable")
+)
+
+type ResolvedAdapter struct {
+	Code         string
+	Revision     int64
+	Adapter      provider.Adapter
+	Capabilities map[provider.Capability]bool
+}
+
+type AdapterResolver interface {
+	Resolve(providerCode string) (ResolvedAdapter, error)
+}
+
+func NewProxy(adapter provider.Adapter, apiKey string, maxBody int64, logger *slog.Logger) *Proxy {
+	return NewProxyWithAuthenticator(adapter, staticKeyAuthenticator(apiKey), maxBody, logger)
+}
+
+func NewProxyWithAuthenticator(adapter provider.Adapter, authenticator apikey.KeyAuthenticator, maxBody int64, logger *slog.Logger) *Proxy {
+	return NewDynamicProxy(staticAdapterResolver{adapter: adapter}, authenticator, maxBody, logger)
+}
+
+func NewDynamicProxy(resolver AdapterResolver, authenticator apikey.KeyAuthenticator, maxBody int64, logger *slog.Logger) *Proxy {
+	transport := &http.Transport{
+		Proxy:                 http.ProxyFromEnvironment,
+		DialContext:           (&net.Dialer{Timeout: 5 * time.Second, KeepAlive: 30 * time.Second}).DialContext,
+		ForceAttemptHTTP2:     true,
+		MaxIdleConns:          512,
+		MaxIdleConnsPerHost:   256,
+		IdleConnTimeout:       90 * time.Second,
+		TLSHandshakeTimeout:   5 * time.Second,
+		ResponseHeaderTimeout: 60 * time.Second,
+		ExpectContinueTimeout: time.Second,
+	}
+	return &Proxy{resolver: resolver, auth: authenticator, maxBody: maxBody, logger: logger, transport: transport, resilience: DefaultResiliencePolicy()}
+}
+
+func (p *Proxy) SetAdmissionController(controller AdmissionController) {
+	p.admission = controller
+}
+
+func (p *Proxy) SetTokenQuotaController(controller TokenQuotaController) {
+	p.tokenQuota = controller
+}
+
+func (p *Proxy) SetResiliencePolicy(policy ResiliencePolicy) {
+	p.resilience = policy
+	p.transport.ResponseHeaderTimeout = policy.ResponseHeaderTimeout
+}
+
+func (p *Proxy) SetAuditRecorder(recorder AuditRecorder)             { p.audit = recorder }
+func (p *Proxy) SetContentPolicyEngine(engine *contentpolicy.Engine) { p.policies = engine }
+func (p *Proxy) SetPricingService(service *pricing.Service)          { p.pricing = service }
+
+func (p *Proxy) ServeHTTP(writer http.ResponseWriter, request *http.Request) {
+	started := time.Now()
+	if !isSupportedPath(request.URL.Path, request.Method) {
+		writeOpenAIError(writer, http.StatusNotFound, "invalid_request_error", "unsupported gateway endpoint")
+		return
+	}
+	principal, err := p.authorized(request)
+	if err != nil {
+		if errors.Is(err, apikey.ErrStore) {
+			writeOpenAIError(writer, http.StatusServiceUnavailable, "authentication_unavailable", "API key service is unavailable")
+			return
+		}
+		writer.Header().Set("WWW-Authenticate", "Bearer")
+		writeOpenAIError(writer, http.StatusUnauthorized, "authentication_error", "invalid API key")
+		return
+	}
+	span := newAuditSpan(p.audit, principal, request, started)
+	if span != nil {
+		span.pricing = p.pricing
+	}
+	if span != nil {
+		statusWriter := &statusResponseWriter{ResponseWriter: writer}
+		writer = statusWriter
+		defer func() { span.finish(statusWriter.Status()) }()
+	}
+	if p.admission != nil && principal.APIKeyID != "" {
+		decision, err := p.admission.Allow(request.Context(), principal, time.Now())
+		if err != nil {
+			writeOpenAIError(writer, http.StatusServiceUnavailable, "rate_limit_unavailable", "rate limit service is unavailable")
+			return
+		}
+		writeAdmissionHeaders(writer.Header(), decision)
+		if !decision.Allowed {
+			writer.Header().Set("Retry-After", strconv.FormatInt(max(int64(decision.RetryAfter.Seconds()), 1), 10))
+			if decision.Reason == AdmissionMonthlyQuota {
+				writeOpenAIError(writer, http.StatusTooManyRequests, "insufficient_quota", "monthly request quota exceeded")
+			} else {
+				writeOpenAIError(writer, http.StatusTooManyRequests, "rate_limit_error", "request rate limit exceeded")
+			}
+			return
+		}
+	}
+	if request.ContentLength > p.maxBody {
+		writeOpenAIError(writer, http.StatusRequestEntityTooLarge, "invalid_request_error", "request body is too large")
+		return
+	}
+	if p.policies != nil {
+		policyResult, policyErr := p.policies.Apply(request, p.maxBody, principal.APIKeyID)
+		if errors.Is(policyErr, contentpolicy.ErrBodyTooLarge) {
+			writeOpenAIError(writer, http.StatusRequestEntityTooLarge, "invalid_request_error", "request body is too large")
+			return
+		}
+		if policyErr != nil {
+			writeOpenAIError(writer, http.StatusBadRequest, "invalid_request_error", "request body could not be inspected")
+			return
+		}
+		span.setContentPolicy(policyResult)
+		if policyResult.Blocked {
+			writeOpenAIError(writer, http.StatusUnprocessableEntity, "content_policy_violation", "request was blocked by content policy")
+			return
+		}
+		if policyResult.Redacted {
+			writer.Header().Set("X-Gateway-Content-Redacted", "true")
+		}
+	}
+	routeResolver, routingEnabled := p.resolver.(ModelRouteResolver)
+	routingEnabled = routingEnabled && routeResolver.ModelRoutingEnabled()
+	var modelPayload modelRequest
+	if routingEnabled {
+		modelPayload, err = readModelRequest(request, p.maxBody)
+		if errors.Is(err, errRequestBodyTooLarge) {
+			writeOpenAIError(writer, http.StatusRequestEntityTooLarge, "invalid_request_error", "request body is too large")
+			return
+		}
+		if err != nil {
+			writeOpenAIError(writer, http.StatusBadRequest, "invalid_request_error", "request body could not be read")
+			return
+		}
+	}
+	var usage *usageSession
+	if p.tokenQuota != nil && principal.APIKeyID != "" {
+		estimate := int64(0)
+		if principal.MonthlyTokenQuota > 0 {
+			estimate, err = prepareTokenBudget(request, p.maxBody)
+			if errors.Is(err, errRequestBodyTooLarge) {
+				writeOpenAIError(writer, http.StatusRequestEntityTooLarge, "invalid_request_error", "request body is too large")
+				return
+			}
+			if err != nil {
+				writeOpenAIError(writer, http.StatusBadRequest, "invalid_request_error", "request body could not be read")
+				return
+			}
+		}
+		reservation, reserveErr := p.tokenQuota.Reserve(request.Context(), principal, estimate, time.Now())
+		if reserveErr != nil && principal.MonthlyTokenQuota > 0 {
+			writeOpenAIError(writer, http.StatusServiceUnavailable, "token_quota_unavailable", "token quota service is unavailable")
+			return
+		}
+		if reserveErr == nil && !reservation.Allowed {
+			writeTokenQuotaHeaders(writer.Header(), reservation)
+			writer.Header().Set("Retry-After", strconv.FormatInt(max(int64(time.Until(reservation.ResetAt).Seconds()), 1), 10))
+			writeOpenAIError(writer, http.StatusTooManyRequests, "insufficient_quota", "monthly token quota exceeded")
+			return
+		}
+		if reserveErr == nil && reservation.CounterKey != "" {
+			writeTokenQuotaHeaders(writer.Header(), reservation)
+			usage = &usageSession{controller: p.tokenQuota, reservation: reservation, log: p.logger}
+			request = withUsageSession(request, usage)
+			defer usage.finish(TokenUsage{})
+		}
+	}
+	if span != nil {
+		if usage == nil {
+			usage = &usageSession{log: p.logger}
+			request = withUsageSession(request, usage)
+			defer usage.finish(TokenUsage{})
+		}
+		usage.onFinish = span.setUsage
+	}
+	if request.Body != nil && request.Method != http.MethodGet {
+		if request.Header.Get("Idempotency-Key") != "" && request.GetBody == nil {
+			if _, err := prepareTokenBudget(request, p.maxBody); err != nil {
+				if errors.Is(err, errRequestBodyTooLarge) {
+					writeOpenAIError(writer, http.StatusRequestEntityTooLarge, "invalid_request_error", "request body is too large")
+				} else {
+					writeOpenAIError(writer, http.StatusBadRequest, "invalid_request_error", "request body could not be read")
+				}
+				return
+			}
+		}
+		request.Body = http.MaxBytesReader(writer, request.Body, p.maxBody)
+	}
+	providerCode := strings.ToLower(strings.TrimSpace(request.Header.Get("X-Gateway-Provider")))
+	var resolved ResolvedAdapter
+	if routingEnabled && modelPayload.model != "" {
+		span.setModel(modelPayload.model)
+		tenantID := ""
+		if principal.TenantID != nil {
+			tenantID = *principal.TenantID
+		}
+		route, routeErr := routeResolver.ResolveModelRoute(ModelRouteQuery{
+			ProviderCode: providerCode, Model: modelPayload.model, Endpoint: request.URL.Path,
+			APIKeyID: principal.APIKeyID, TenantID: tenantID, Seed: RequestID(request.Context()),
+		})
+		if routeErr != nil {
+			writeOpenAIError(writer, http.StatusServiceUnavailable, "provider_unavailable", "model routing configuration is unavailable")
+			return
+		}
+		if route.Known && !route.Matched {
+			writeOpenAIError(writer, http.StatusBadRequest, "invalid_request_error", "model route is not available for this request")
+			return
+		}
+		if route.Matched {
+			resolved = route.ResolvedAdapter
+			if err := modelPayload.rewrite(request, route.TargetModel); err != nil {
+				writeOpenAIError(writer, http.StatusBadRequest, "invalid_request_error", "request model could not be rewritten")
+				return
+			}
+			writer.Header().Set("X-Gateway-Model", route.TargetModel)
+		}
+	}
+	if resolved.Adapter == nil {
+		resolved, err = p.resolver.Resolve(providerCode)
+	}
+	if err != nil {
+		if errors.Is(err, ErrProviderNotFound) {
+			writeOpenAIError(writer, http.StatusBadRequest, "invalid_request_error", "requested provider is not available")
+			return
+		}
+		writeOpenAIError(writer, http.StatusServiceUnavailable, "provider_unavailable", "provider configuration is unavailable")
+		return
+	}
+	capability := capabilityForPath(request.URL.Path)
+	if capability != "" && !resolved.Capabilities[capability] {
+		writeOpenAIError(writer, http.StatusBadRequest, "invalid_request_error", "provider does not support this endpoint")
+		return
+	}
+	request.Header.Del("X-Gateway-Provider")
+	writer.Header().Set("X-Gateway-Provider", resolved.Code)
+	span.setRoute(resolved.Code, writer.Header().Get("X-Gateway-Model"))
+	request.Body = span.captureBody(request.Body)
+	p.proxyFor(resolved).ServeHTTP(writer, request)
+}
+
+func (p *Proxy) proxyFor(resolved ResolvedAdapter) *httputil.ReverseProxy {
+	target := resolved.Adapter.Target()
+	key := fmt.Sprintf("%s:%d:%s", resolved.Code, resolved.Revision, target.String())
+	if cached, ok := p.proxies.Load(resolved.Code); ok {
+		entry := cached.(cachedProxy)
+		if entry.key == key {
+			return entry.proxy
+		}
+	}
+	reverseProxy := httputil.NewSingleHostReverseProxy(target)
+	originalDirector := reverseProxy.Director
+	reverseProxy.Director = func(request *http.Request) {
+		originalDirector(request)
+		request.Host = target.Host
+		resolved.Adapter.Prepare(request)
+	}
+	reverseProxy.FlushInterval = -1
+	circuitValue, _ := p.circuits.LoadOrStore(resolved.Code, newCircuitBreaker(p.resilience))
+	reverseProxy.Transport = &resilientTransport{
+		base: p.transport, circuit: circuitValue.(*circuitBreaker), maxRetries: p.resilience.MaxRetries, backoff: p.resilience.RetryBackoff,
+	}
+	reverseProxy.ModifyResponse = func(response *http.Response) error {
+		if session := usageSessionFrom(response.Request); session != nil {
+			if response.StatusCode >= http.StatusOK && response.StatusCode < http.StatusMultipleChoices {
+				session.fallback = session.reservation.Reserved
+				response.Body = newUsageReadCloser(response.Body, response.Header.Get("Content-Type"), session)
+			} else {
+				session.fallback = 0
+				session.finish(TokenUsage{})
+			}
+		}
+		return nil
+	}
+	reverseProxy.ErrorHandler = func(writer http.ResponseWriter, request *http.Request, err error) {
+		if session := usageSessionFrom(request); session != nil {
+			session.fallback = 0
+			session.finish(TokenUsage{})
+		}
+		p.logger.Error("upstream request failed", "request_id", RequestID(request.Context()), "provider", resolved.Code, "error", err)
+		if errors.Is(err, ErrCircuitOpen) {
+			writer.Header().Set("Retry-After", "30")
+			writeOpenAIError(writer, http.StatusServiceUnavailable, "provider_unavailable", "provider circuit is temporarily open")
+			return
+		}
+		writeOpenAIError(writer, http.StatusBadGateway, "upstream_error", "upstream service is unavailable")
+	}
+	p.proxies.Store(resolved.Code, cachedProxy{key: key, proxy: reverseProxy})
+	return reverseProxy
+}
+
+func (p *Proxy) authorized(request *http.Request) (apikey.Principal, error) {
+	presented := strings.TrimSpace(request.Header.Get("X-Gateway-API-Key"))
+	if presented == "" {
+		authorization := strings.TrimSpace(request.Header.Get("Authorization"))
+		if len(authorization) > len("Bearer ") && strings.EqualFold(authorization[:len("Bearer ")], "Bearer ") {
+			presented = strings.TrimSpace(authorization[len("Bearer "):])
+		}
+	}
+	if p.auth == nil {
+		return apikey.Principal{}, apikey.ErrInvalid
+	}
+	if authenticator, ok := p.auth.(apikey.PrincipalAuthenticator); ok {
+		return authenticator.AuthenticatePrincipal(request.Context(), presented)
+	}
+	return apikey.Principal{}, p.auth.Authenticate(request.Context(), presented)
+}
+
+func writeAdmissionHeaders(header http.Header, decision AdmissionDecision) {
+	if decision.Limit <= 0 || decision.ResetAt.IsZero() {
+		return
+	}
+	header.Set("X-RateLimit-Limit", strconv.FormatInt(decision.Limit, 10))
+	header.Set("X-RateLimit-Remaining", strconv.FormatInt(max(decision.Remaining, 0), 10))
+	header.Set("X-RateLimit-Reset", strconv.FormatInt(decision.ResetAt.Unix(), 10))
+}
+
+type staticKeyAuthenticator string
+
+func (a staticKeyAuthenticator) Authenticate(_ context.Context, presented string) error {
+	key := string(a)
+	if key == "" {
+		return nil
+	}
+	if len(presented) != len(key) || subtle.ConstantTimeCompare([]byte(presented), []byte(key)) != 1 {
+		return apikey.ErrInvalid
+	}
+	return nil
+}
+
+type staticAdapterResolver struct{ adapter provider.Adapter }
+
+func (r staticAdapterResolver) Resolve(code string) (ResolvedAdapter, error) {
+	if r.adapter == nil || code != "" && code != "environment" {
+		return ResolvedAdapter{}, ErrProviderNotFound
+	}
+	capabilities := make(map[provider.Capability]bool)
+	for _, capability := range r.adapter.Capabilities() {
+		capabilities[capability] = true
+	}
+	return ResolvedAdapter{Code: "environment", Adapter: r.adapter, Capabilities: capabilities}, nil
+}
+
+func capabilityForPath(path string) provider.Capability {
+	switch path {
+	case "/v1/models":
+		return provider.CapabilityModels
+	case "/v1/chat/completions":
+		return provider.CapabilityChat
+	case "/v1/responses":
+		return provider.CapabilityResponses
+	case "/v1/embeddings":
+		return provider.CapabilityEmbeddings
+	case "/v1/messages":
+		return provider.CapabilityMessages
+	default:
+		return ""
+	}
+}
+
+func isSupportedPath(path, method string) bool {
+	switch path {
+	case "/v1/models":
+		return method == http.MethodGet
+	case "/v1/chat/completions", "/v1/responses", "/v1/embeddings", "/v1/messages":
+		return method == http.MethodPost
+	default:
+		return false
+	}
+}
+
+func writeOpenAIError(writer http.ResponseWriter, status int, errorType, message string) {
+	writer.Header().Set("Content-Type", "application/json")
+	writer.WriteHeader(status)
+	_ = json.NewEncoder(writer).Encode(map[string]any{
+		"error": map[string]any{"message": message, "type": errorType, "param": nil, "code": nil},
+	})
+}
+
+var errInvalidAdapter = errors.New("invalid provider adapter")
+
+func ValidateAdapter(adapter provider.Adapter) error {
+	if adapter == nil || adapter.Target() == nil || adapter.Name() == "" {
+		return errInvalidAdapter
+	}
+	return nil
+}