ai-gateway-go/internal/contentpolicy/engine.go

package contentpolicy

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log/slog"
	"net/http"
	"regexp"
	"slices"
	"strconv"
	"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{}
	// 只收集"本请求命中"的 redact 策略的规则,供字节级重写使用;
	// 跨端点/模型/API Key 作用域的策略不得改写本请求。
	matchedRedactionRules := make([]compiledRule, 0)
	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 policy.Action == "redact" && changed {
			matchedRedactionRules = append(matchedRedactionRules, policy.rules...)
		}
	}
	if result.Redacted && !result.Blocked {
		// 在原始请求字节上做字符串字面量级替换,而不是解码后重新
		// json.Marshal:重编码会改变键序、数字格式与 HTML 转义,破坏上游
		// 的请求签名/哈希与字节级契约,且对大 body 是双倍编解码开销。
		encoded, changed := redactBytes(body, matchedRedactionRules)
		if changed {
			restoreBody(request, encoded)
		} else {
			// 字节级替换未生效(如规则只匹配解码后文本但替换失败):
			// 不得谎报已脱敏,否则 X-Gateway-Content-Redacted 与审计
			// 都声称敏感信息已被移除,而实际请求原样发往上游。
			result.Redacted = false
		}
	}
	return result, nil
}

// redactBytes 逐字节扫描 JSON,仅对"位于 textual 字段(content/text/input/
// prompt/instructions)下的值字符串"应用脱敏规则,其余字节(键名、空白、
// 数字、布尔、结构符)原样保留。与解码后替换相比:
//   - 字节流与原始请求一致,除被替换的匹配段外零改动,不破坏上游签名/
//     哈希与字节级契约,也不做大 body 的双倍编解码;
//   - 字符串内的 JSON 转义按原文匹配(如 \uXXXX),secret 模式通常不含
//     需要转义的字符,实际影响可忽略。
//
// 返回替换后的字节与是否发生过替换。
func redactBytes(raw []byte, rules []compiledRule) ([]byte, bool) {
	if len(rules) == 0 {
		return raw, false
	}
	type frame struct {
		inObject bool
		selected bool
	}
	changed := false
	out := make([]byte, 0, len(raw)+64)
	curSelected := false
	keySelected := false
	stack := make([]frame, 0, 8)
	for i := 0; i < len(raw); {
		ch := raw[i]
		switch ch {
		case '{', '[':
			stack = append(stack, frame{inObject: ch == '{', selected: curSelected})
			out = append(out, ch)
			i++
			continue
		case '}', ']':
			if len(stack) > 0 {
				curSelected = stack[len(stack)-1].selected
				stack = stack[:len(stack)-1]
			}
			out = append(out, ch)
			i++
			continue
		case ',':
			// 对象内逗号后是键:selected 由下一个键决定;数组内逗号后是
			// 元素,继承当前 selected。
			if len(stack) > 0 && stack[len(stack)-1].inObject {
				curSelected = false
			}
			out = append(out, ch)
			i++
			continue
		case ':':
			// 键后冒号:值字符串的 selected 由该键决定。
			curSelected = keySelected
			out = append(out, ch)
			i++
			continue
		case '"':
			// 定位字符串结束(处理反斜杠转义)。
			j := i + 1
			escaped := false
			for j < len(raw) {
				if escaped {
					escaped = false
					j++
					continue
				}
				if raw[j] == '\\' {
					escaped = true
					j++
					continue
				}
				if raw[j] == '"' {
					break
				}
				j++
			}
			if j >= len(raw) {
				// 截断/畸形 JSON:剩余字节原样保留。
				out = append(out, raw[i:]...)
				break
			}
			content := raw[i+1 : j]
			// 键还是值:字符串后第一个非空白字符是 ':' 即为对象键。
			k := j + 1
			for k < len(raw) && (raw[k] == ' ' || raw[k] == '\t' || raw[k] == '\n' || raw[k] == '\r') {
				k++
			}
			isKey := k < len(raw) && raw[k] == ':'
			if isKey {
				// textual 字段名传播到其值:父级 selected 或键名命中。
				parentSelected := false
				if len(stack) > 0 {
					parentSelected = stack[len(stack)-1].selected
				}
				keySelected = parentSelected || textualFields[strings.ToLower(string(content))]
			} else if curSelected {
				if replaced, hit := applyRules(content, rules); hit {
					changed = true
					out = append(out, '"')
					out = append(out, replaced...)
					out = append(out, '"')
					i = j + 1
					continue
				}
			}
			out = append(out, raw[i:j+1]...)
			i = j + 1
			continue
		default:
			out = append(out, ch)
			i++
		}
	}
	return out, changed
}

func applyRules(value []byte, rules []compiledRule) ([]byte, bool) {
	changed := false
	text := string(value)
	for _, rule := range rules {
		if rule.expression.MatchString(text) {
			next := rule.expression.ReplaceAllString(text, jsonEscapeReplacement(rule.replacement))
			if next != text {
				text = next
				changed = true
			}
			continue
		}
		// 原文(含 JSON 转义)未命中,但解码后的文本可能命中
		// (Go 的 json.Marshal 会把 - & < > 等转义为 \u002d \u0026 ...):
		// 对解码文本应用替换后重新做 JSON 字符串转义,其余字节不变。
		var decoded string
		wrapped := append([]byte(`"`), value...)
		wrapped = append(wrapped, '"')
		if json.Unmarshal(wrapped, &decoded) == nil && decoded != text && rule.expression.MatchString(decoded) {
			next := rule.expression.ReplaceAllString(decoded, rule.replacement)
			quoted := strconv.Quote(next)
			text = quoted[1 : len(quoted)-1]
			changed = true
		}
	}
	return []byte(text), changed
}

// jsonEscapeReplacement 把替换文本转义成可安全嵌入 JSON 字符串字面量的
// 形式:替换文本含引号/反斜杠/控制字符时直接插入会破坏 JSON 结构。
func jsonEscapeReplacement(value string) string {
	if !strings.ContainsAny(value, "\"\\\n\r\t") && !strings.ContainsFunc(value, func(r rune) bool { return r < 0x20 }) {
		return value
	}
	quoted := strconv.Quote(value)
	return quoted[1 : len(quoted)-1]
}

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 }
}