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,272 @@
+package contentpolicy
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"io"
+	"log/slog"
+	"net/http"
+	"regexp"
+	"slices"
+	"strings"
+	"sync/atomic"
+	"time"
+
+	"github.com/jackc/pgx/v5/pgxpool"
+)
+
+var ErrBodyTooLarge = errors.New("request body is too large")
+
+type Rule struct {
+	Name        string `json:"name"`
+	Pattern     string `json:"pattern"`
+	Replacement string `json:"replacement"`
+}
+
+type Policy struct {
+	ID          string    `json:"id"`
+	Name        string    `json:"name"`
+	Description string    `json:"description"`
+	Action      string    `json:"action"`
+	Priority    int       `json:"priority"`
+	Paths       []string  `json:"paths"`
+	Models      []string  `json:"models"`
+	APIKeyIDs   []string  `json:"api_key_ids"`
+	Rules       []Rule    `json:"rules"`
+	Enabled     bool      `json:"enabled"`
+	Revision    int64     `json:"revision"`
+	CreatedAt   time.Time `json:"created_at"`
+	UpdatedAt   time.Time `json:"updated_at"`
+}
+
+type compiledRule struct {
+	name, replacement string
+	expression        *regexp.Regexp
+}
+type compiledPolicy struct {
+	Policy
+	rules []compiledRule
+}
+type snapshot struct{ policies []compiledPolicy }
+
+type Engine struct {
+	pool    *pgxpool.Pool
+	logger  *slog.Logger
+	refresh time.Duration
+	current atomic.Pointer[snapshot]
+}
+
+type Match struct {
+	PolicyID   string   `json:"policy_id"`
+	PolicyName string   `json:"policy_name"`
+	Action     string   `json:"action"`
+	Rules      []string `json:"rules"`
+}
+
+type Result struct {
+	Matches           []Match
+	Blocked, Redacted bool
+}
+
+func NewEngine(pool *pgxpool.Pool, refresh time.Duration, logger *slog.Logger) *Engine {
+	if refresh <= 0 {
+		refresh = 30 * time.Second
+	}
+	e := &Engine{pool: pool, refresh: refresh, logger: logger}
+	e.current.Store(&snapshot{})
+	return e
+}
+
+func (e *Engine) Run(ctx context.Context) {
+	_ = e.Reload(ctx)
+	ticker := time.NewTicker(e.refresh)
+	defer ticker.Stop()
+	for {
+		select {
+		case <-ticker.C:
+			if err := e.Reload(ctx); err != nil && e.logger != nil {
+				e.logger.Warn("content policy refresh failed", "error", err)
+			}
+		case <-ctx.Done():
+			return
+		}
+	}
+}
+
+func (e *Engine) Reload(ctx context.Context) error {
+	if e == nil || e.pool == nil {
+		return errors.New("content policy store unavailable")
+	}
+	rows, err := e.pool.Query(ctx, `SELECT id::text,name,description,action,priority,paths,models,api_key_ids::text[],rules,enabled,revision,created_at,updated_at FROM gateway.content_policies WHERE enabled ORDER BY priority DESC,name,id`)
+	if err != nil {
+		return fmt.Errorf("load content policies: %w", err)
+	}
+	defer rows.Close()
+	loaded := make([]compiledPolicy, 0)
+	for rows.Next() {
+		var policy Policy
+		var raw []byte
+		if err := rows.Scan(&policy.ID, &policy.Name, &policy.Description, &policy.Action, &policy.Priority, &policy.Paths, &policy.Models, &policy.APIKeyIDs, &raw, &policy.Enabled, &policy.Revision, &policy.CreatedAt, &policy.UpdatedAt); err != nil {
+			return err
+		}
+		if err := json.Unmarshal(raw, &policy.Rules); err != nil {
+			return fmt.Errorf("decode policy %s: %w", policy.ID, err)
+		}
+		compiled, err := compilePolicy(policy)
+		if err != nil {
+			return fmt.Errorf("compile policy %s: %w", policy.ID, err)
+		}
+		loaded = append(loaded, compiled)
+	}
+	if err := rows.Err(); err != nil {
+		return err
+	}
+	e.current.Store(&snapshot{policies: loaded})
+	return nil
+}
+
+func compilePolicy(policy Policy) (compiledPolicy, error) {
+	if policy.Action != "audit" && policy.Action != "block" && policy.Action != "redact" {
+		return compiledPolicy{}, errors.New("unsupported action")
+	}
+	if len(policy.Rules) == 0 || len(policy.Rules) > 20 {
+		return compiledPolicy{}, errors.New("rules must contain 1 to 20 entries")
+	}
+	result := compiledPolicy{Policy: policy, rules: make([]compiledRule, 0, len(policy.Rules))}
+	for _, rule := range policy.Rules {
+		if strings.TrimSpace(rule.Name) == "" || len(rule.Pattern) == 0 || len(rule.Pattern) > 512 {
+			return compiledPolicy{}, errors.New("invalid rule name or pattern length")
+		}
+		expression, err := regexp.Compile(rule.Pattern)
+		if err != nil {
+			return compiledPolicy{}, err
+		}
+		replacement := rule.Replacement
+		if replacement == "" {
+			replacement = "[REDACTED]"
+		}
+		result.rules = append(result.rules, compiledRule{name: rule.Name, replacement: replacement, expression: expression})
+	}
+	return result, nil
+}
+
+func Validate(policy Policy) error { _, err := compilePolicy(policy); return err }
+
+func (e *Engine) Apply(request *http.Request, maxBody int64, apiKeyID string) (Result, error) {
+	if e == nil || request.Body == nil || request.Body == http.NoBody || request.Method == http.MethodGet {
+		return Result{}, nil
+	}
+	body, err := io.ReadAll(io.LimitReader(request.Body, maxBody+1))
+	if err != nil {
+		return Result{}, err
+	}
+	_ = request.Body.Close()
+	if int64(len(body)) > maxBody {
+		return Result{}, ErrBodyTooLarge
+	}
+	restoreBody(request, body)
+	var document any
+	decoder := json.NewDecoder(bytes.NewReader(body))
+	decoder.UseNumber()
+	if decoder.Decode(&document) != nil {
+		return Result{}, nil
+	}
+	model := findModel(document)
+	result := Result{}
+	for _, policy := range e.current.Load().policies {
+		if !policyApplies(policy, request.URL.Path, model, apiKeyID) {
+			continue
+		}
+		matched := make([]string, 0)
+		changed := false
+		walkText(&document, false, func(value string) string {
+			for _, rule := range policy.rules {
+				if rule.expression.MatchString(value) {
+					if !slices.Contains(matched, rule.name) {
+						matched = append(matched, rule.name)
+					}
+					if policy.Action == "redact" {
+						next := rule.expression.ReplaceAllString(value, rule.replacement)
+						changed = changed || next != value
+						value = next
+					}
+				}
+			}
+			return value
+		})
+		if len(matched) == 0 {
+			continue
+		}
+		result.Matches = append(result.Matches, Match{PolicyID: policy.ID, PolicyName: policy.Name, Action: policy.Action, Rules: matched})
+		if policy.Action == "block" {
+			result.Blocked = true
+			break
+		}
+		result.Redacted = result.Redacted || changed
+	}
+	if result.Redacted && !result.Blocked {
+		encoded, err := json.Marshal(document)
+		if err != nil {
+			return Result{}, err
+		}
+		restoreBody(request, encoded)
+	}
+	return result, nil
+}
+
+func policyApplies(policy compiledPolicy, path, model, apiKeyID string) bool {
+	return matches(policy.Paths, path) && matches(policy.Models, model) && matches(policy.APIKeyIDs, apiKeyID)
+}
+
+func matches(values []string, candidate string) bool {
+	if len(values) == 0 {
+		return true
+	}
+	for _, value := range values {
+		if value == "*" || value == candidate {
+			return true
+		}
+	}
+	return false
+}
+
+var textualFields = map[string]bool{"content": true, "text": true, "input": true, "prompt": true, "instructions": true}
+
+func walkText(value *any, selected bool, transform func(string) string) {
+	switch current := (*value).(type) {
+	case map[string]any:
+		for key, child := range current {
+			local := child
+			walkText(&local, selected || textualFields[strings.ToLower(key)], transform)
+			current[key] = local
+		}
+	case []any:
+		for index, child := range current {
+			local := child
+			walkText(&local, selected, transform)
+			current[index] = local
+		}
+	case string:
+		if selected {
+			*value = transform(current)
+		}
+	}
+}
+
+func findModel(document any) string {
+	if object, ok := document.(map[string]any); ok {
+		if model, ok := object["model"].(string); ok {
+			return model
+		}
+	}
+	return ""
+}
+
+func restoreBody(request *http.Request, body []byte) {
+	request.Body = io.NopCloser(bytes.NewReader(body))
+	request.ContentLength = int64(len(body))
+	request.GetBody = func() (io.ReadCloser, error) { return io.NopCloser(bytes.NewReader(body)), nil }
+}