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

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