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>
This commit is contained in:
@@ -0,0 +1,235 @@
|
||||
package shadow
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"hash/fnv"
|
||||
"io"
|
||||
"log/slog"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"sync/atomic"
|
||||
|
||||
"aigateway.local/core/internal/gateway"
|
||||
"aigateway.local/core/internal/platform/config"
|
||||
)
|
||||
|
||||
type Metrics struct{ eligible, sent, matched, mismatched, errors, skipped atomic.Uint64 }
|
||||
|
||||
type Middleware struct {
|
||||
config config.Shadow
|
||||
logger *slog.Logger
|
||||
client *http.Client
|
||||
limit chan struct{}
|
||||
metrics Metrics
|
||||
}
|
||||
|
||||
func New(cfg config.Shadow, logger *slog.Logger) *Middleware {
|
||||
return &Middleware{config: cfg, logger: logger, client: &http.Client{Timeout: cfg.Timeout, CheckRedirect: func(*http.Request, []*http.Request) error { return fmt.Errorf("shadow redirects are disabled") }}, limit: make(chan struct{}, cfg.MaxConcurrent)}
|
||||
}
|
||||
|
||||
func (m *Middleware) Wrap(next http.Handler) http.Handler {
|
||||
if m == nil || m.config.BaseURL == "" || m.config.SampleRate <= 0 {
|
||||
return next
|
||||
}
|
||||
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.Method != http.MethodPost || !m.sample(gateway.RequestID(r.Context())) {
|
||||
next.ServeHTTP(w, r)
|
||||
return
|
||||
}
|
||||
body, eligible := m.captureRequest(r)
|
||||
if !eligible {
|
||||
next.ServeHTTP(w, r)
|
||||
return
|
||||
}
|
||||
m.metrics.eligible.Add(1)
|
||||
select {
|
||||
case m.limit <- struct{}{}:
|
||||
default:
|
||||
m.metrics.skipped.Add(1)
|
||||
next.ServeHTTP(w, r)
|
||||
return
|
||||
}
|
||||
capture := &captureWriter{ResponseWriter: w, limit: m.config.MaxBodyBytes}
|
||||
next.ServeHTTP(capture, r)
|
||||
primaryStatus := capture.Status()
|
||||
primaryBody := append([]byte(nil), capture.body...)
|
||||
headers := shadowHeaders(r.Header)
|
||||
path := r.URL.RequestURI()
|
||||
go func() {
|
||||
defer func() { <-m.limit }()
|
||||
m.compare(context.WithoutCancel(r.Context()), path, headers, body, primaryStatus, primaryBody)
|
||||
}()
|
||||
})
|
||||
}
|
||||
|
||||
func (m *Middleware) captureRequest(r *http.Request) ([]byte, bool) {
|
||||
if r.Method != http.MethodPost || r.Body == nil || r.Body == http.NoBody {
|
||||
return nil, false
|
||||
}
|
||||
original := r.Body
|
||||
raw, err := io.ReadAll(io.LimitReader(original, m.config.MaxBodyBytes+1))
|
||||
if err != nil || int64(len(raw)) > m.config.MaxBodyBytes {
|
||||
r.Body = &multiReadCloser{Reader: io.MultiReader(bytes.NewReader(raw), original), Closer: original}
|
||||
return nil, false
|
||||
}
|
||||
_ = original.Close()
|
||||
r.Body = io.NopCloser(bytes.NewReader(raw))
|
||||
r.GetBody = func() (io.ReadCloser, error) { return io.NopCloser(bytes.NewReader(raw)), nil }
|
||||
var envelope struct {
|
||||
Stream bool `json:"stream"`
|
||||
}
|
||||
if json.Unmarshal(raw, &envelope) != nil || envelope.Stream {
|
||||
return nil, false
|
||||
}
|
||||
return raw, true
|
||||
}
|
||||
|
||||
type multiReadCloser struct {
|
||||
io.Reader
|
||||
io.Closer
|
||||
}
|
||||
|
||||
func (m *Middleware) sample(requestID string) bool {
|
||||
hash := fnv.New64a()
|
||||
_, _ = hash.Write([]byte(requestID))
|
||||
return float64(hash.Sum64()%1_000_000)/1_000_000 < m.config.SampleRate
|
||||
}
|
||||
|
||||
func (m *Middleware) compare(ctx context.Context, path string, headers http.Header, body []byte, primaryStatus int, primaryBody []byte) {
|
||||
ctx, cancel := context.WithTimeout(ctx, m.config.Timeout)
|
||||
defer cancel()
|
||||
target, err := url.Parse(m.config.BaseURL + path)
|
||||
if err != nil {
|
||||
m.fail("invalid shadow URL", err)
|
||||
return
|
||||
}
|
||||
request, err := http.NewRequestWithContext(ctx, http.MethodPost, target.String(), bytes.NewReader(body))
|
||||
if err != nil {
|
||||
m.fail("create shadow request", err)
|
||||
return
|
||||
}
|
||||
request.Header = headers
|
||||
request.Header.Set("Authorization", "Bearer "+m.config.APIKey)
|
||||
request.Header.Del("X-Gateway-API-Key")
|
||||
request.Header.Set("X-Gateway-Shadow", "true")
|
||||
response, err := m.client.Do(request)
|
||||
if err != nil {
|
||||
m.fail("shadow request failed", err)
|
||||
return
|
||||
}
|
||||
defer response.Body.Close()
|
||||
raw, err := io.ReadAll(io.LimitReader(response.Body, m.config.MaxBodyBytes+1))
|
||||
if err != nil || int64(len(raw)) > m.config.MaxBodyBytes {
|
||||
m.fail("shadow response invalid", err)
|
||||
return
|
||||
}
|
||||
m.metrics.sent.Add(1)
|
||||
primarySignature := contractSignature(primaryStatus, primaryBody)
|
||||
shadowSignature := contractSignature(response.StatusCode, raw)
|
||||
if primarySignature == shadowSignature {
|
||||
m.metrics.matched.Add(1)
|
||||
return
|
||||
}
|
||||
m.metrics.mismatched.Add(1)
|
||||
if m.logger != nil {
|
||||
m.logger.Warn("shadow contract mismatch", "path", path, "primary_status", primaryStatus, "shadow_status", response.StatusCode, "primary_signature", primarySignature, "shadow_signature", shadowSignature)
|
||||
}
|
||||
}
|
||||
func (m *Middleware) fail(message string, err error) {
|
||||
m.metrics.errors.Add(1)
|
||||
if m.logger != nil {
|
||||
m.logger.Warn(message, "error", err)
|
||||
}
|
||||
}
|
||||
|
||||
func contractSignature(status int, raw []byte) string {
|
||||
var value any
|
||||
if json.Unmarshal(raw, &value) != nil {
|
||||
return fmt.Sprintf("%d:non-json", status)
|
||||
}
|
||||
shape, _ := json.Marshal(jsonShape(value))
|
||||
return fmt.Sprintf("%d:%s", status, shape)
|
||||
}
|
||||
func jsonShape(value any) any {
|
||||
switch item := value.(type) {
|
||||
case map[string]any:
|
||||
result := make(map[string]any, len(item))
|
||||
for key, child := range item {
|
||||
switch key {
|
||||
case "id", "created", "created_at", "system_fingerprint":
|
||||
continue
|
||||
}
|
||||
result[key] = jsonShape(child)
|
||||
}
|
||||
return result
|
||||
case []any:
|
||||
if len(item) == 0 {
|
||||
return []any{}
|
||||
}
|
||||
return []any{jsonShape(item[0])}
|
||||
case string:
|
||||
return "string"
|
||||
case float64:
|
||||
return "number"
|
||||
case bool:
|
||||
return "boolean"
|
||||
case nil:
|
||||
return nil
|
||||
default:
|
||||
return fmt.Sprintf("%T", item)
|
||||
}
|
||||
}
|
||||
func shadowHeaders(source http.Header) http.Header {
|
||||
result := make(http.Header)
|
||||
for _, key := range []string{"Content-Type", "Accept", "Idempotency-Key", "X-Gateway-Provider"} {
|
||||
if value := source.Values(key); len(value) > 0 {
|
||||
result[key] = append([]string(nil), value...)
|
||||
}
|
||||
}
|
||||
return result
|
||||
}
|
||||
func (m *Middleware) Prometheus() string {
|
||||
if m == nil {
|
||||
return ""
|
||||
}
|
||||
return fmt.Sprintf("gateway_shadow_eligible_total %d\ngateway_shadow_sent_total %d\ngateway_shadow_contract_matches_total %d\ngateway_shadow_contract_mismatches_total %d\ngateway_shadow_errors_total %d\ngateway_shadow_skipped_total %d\n", m.metrics.eligible.Load(), m.metrics.sent.Load(), m.metrics.matched.Load(), m.metrics.mismatched.Load(), m.metrics.errors.Load(), m.metrics.skipped.Load())
|
||||
}
|
||||
|
||||
type captureWriter struct {
|
||||
http.ResponseWriter
|
||||
status int
|
||||
body []byte
|
||||
limit int64
|
||||
}
|
||||
|
||||
func (w *captureWriter) WriteHeader(status int) {
|
||||
if w.status == 0 {
|
||||
w.status = status
|
||||
}
|
||||
w.ResponseWriter.WriteHeader(status)
|
||||
}
|
||||
func (w *captureWriter) Write(value []byte) (int, error) {
|
||||
if w.status == 0 {
|
||||
w.status = http.StatusOK
|
||||
}
|
||||
if int64(len(w.body)) < w.limit {
|
||||
remaining := int(w.limit - int64(len(w.body)))
|
||||
w.body = append(w.body, value[:min(len(value), remaining)]...)
|
||||
}
|
||||
return w.ResponseWriter.Write(value)
|
||||
}
|
||||
func (w *captureWriter) Status() int {
|
||||
if w.status == 0 {
|
||||
return http.StatusOK
|
||||
}
|
||||
return w.status
|
||||
}
|
||||
func (w *captureWriter) Unwrap() http.ResponseWriter { return w.ResponseWriter }
|
||||
func (w *captureWriter) Flush() {
|
||||
if flusher, ok := w.ResponseWriter.(http.Flusher); ok {
|
||||
flusher.Flush()
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,83 @@
|
||||
package shadow
|
||||
|
||||
import (
|
||||
"io"
|
||||
"log/slog"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"aigateway.local/core/internal/gateway"
|
||||
"aigateway.local/core/internal/platform/config"
|
||||
)
|
||||
|
||||
func TestContractSignatureComparesShapeNotGeneratedValues(t *testing.T) {
|
||||
left := contractSignature(200, []byte(`{"id":"one","choices":[{"message":{"content":"answer one"}}]}`))
|
||||
right := contractSignature(200, []byte(`{"id":"two","choices":[{"message":{"content":"answer two"}}]}`))
|
||||
if left != right {
|
||||
t.Fatalf("expected equal shapes: %s != %s", left, right)
|
||||
}
|
||||
if left == contractSignature(500, []byte(`{"choices":[]}`)) {
|
||||
t.Fatal("status and schema mismatch must differ")
|
||||
}
|
||||
}
|
||||
|
||||
func TestMiddlewareUsesDedicatedCredentialAndDoesNotDelayPrimary(t *testing.T) {
|
||||
delivered := make(chan string, 1)
|
||||
legacy := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
raw, _ := io.ReadAll(r.Body)
|
||||
delivered <- r.Header.Get("Authorization") + ":" + string(raw)
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_, _ = w.Write([]byte(`{"id":"legacy","choices":[{"message":{"content":"different"}}]}`))
|
||||
}))
|
||||
defer legacy.Close()
|
||||
middleware := New(config.Shadow{BaseURL: legacy.URL, APIKey: "shadow-secret", SampleRate: 1, Timeout: time.Second, MaxBodyBytes: 1 << 20, MaxConcurrent: 2}, slog.Default())
|
||||
primary := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_, _ = w.Write([]byte(`{"id":"go","choices":[{"message":{"content":"primary"}}]}`))
|
||||
})
|
||||
request := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{"model":"test","stream":false}`))
|
||||
request.Header.Set("Authorization", "Bearer production-secret")
|
||||
request = request.WithContext(gateway.WithRequestID(request.Context(), "shadow-test"))
|
||||
response := httptest.NewRecorder()
|
||||
middleware.Wrap(primary).ServeHTTP(response, request)
|
||||
if response.Code != 200 {
|
||||
t.Fatalf("primary status %d", response.Code)
|
||||
}
|
||||
select {
|
||||
case value := <-delivered:
|
||||
if !strings.HasPrefix(value, "Bearer shadow-secret:") {
|
||||
t.Fatalf("shadow credential not replaced: %s", value)
|
||||
}
|
||||
case <-time.After(2 * time.Second):
|
||||
t.Fatal("shadow request was not delivered")
|
||||
}
|
||||
deadline := time.Now().Add(time.Second)
|
||||
for middleware.metrics.matched.Load()+middleware.metrics.mismatched.Load() == 0 && time.Now().Before(deadline) {
|
||||
time.Sleep(time.Millisecond)
|
||||
}
|
||||
if middleware.metrics.matched.Load() != 1 {
|
||||
t.Fatalf("expected shape match, metrics=%s", middleware.Prometheus())
|
||||
}
|
||||
}
|
||||
|
||||
func TestOversizedShadowBodyIsPreservedForPrimary(t *testing.T) {
|
||||
middleware := New(config.Shadow{BaseURL: "http://127.0.0.1:1", APIKey: "shadow", SampleRate: 1, Timeout: time.Second, MaxBodyBytes: 16, MaxConcurrent: 1}, slog.Default())
|
||||
payload := `{"stream":false,"prompt":"this body is larger than the shadow limit"}`
|
||||
primary := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
raw, _ := io.ReadAll(r.Body)
|
||||
if string(raw) != payload {
|
||||
t.Fatalf("primary body changed: %q", raw)
|
||||
}
|
||||
w.WriteHeader(http.StatusNoContent)
|
||||
})
|
||||
request := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(payload))
|
||||
request = request.WithContext(gateway.WithRequestID(request.Context(), "large-shadow"))
|
||||
response := httptest.NewRecorder()
|
||||
middleware.Wrap(primary).ServeHTTP(response, request)
|
||||
if response.Code != http.StatusNoContent {
|
||||
t.Fatalf("status %d", response.Code)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user