llmgw/internal/guard/gate.go

package guard

import (
	"context"
	"encoding/json"
	"fmt"
	"log/slog"
	"maps"
	"regexp"
	"sort"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"gitea.mynodes.xyz/ilker/llmgw/internal/guard/detect"
)

// Attachment ties a policy to a part of the tenancy tree: scope is org, team,
// project, user or key, and ScopeID which one.
type Attachment struct {
	Scope, ScopeID, Policy string
}

// Attachments is where they are kept: the store.
type Attachments interface {
	GuardAttachments(ctx context.Context) ([]Attachment, error)
}

// Scopes, from the widest.
var Scopes = []string{"org", "team", "project", "user", "key"}

// Gate is the Guard the gateway runs: the policy library, the attachments,
// and the detectors.
type Gate struct {
	cfg   Config
	store Attachments
	log   *slog.Logger
	now   func() time.Time

	deny   atomic.Pointer[map[string]*regexp.Regexp] // by pattern regex
	tools  toolMemo
	judger *judge
	m      *gateMetrics
	// alarm is told each time a judge gives no verdict; nil tells no one.
	alarm func(model string, err error)
	// lastErr is each judge's last failure to give a verdict.
	errMu   sync.Mutex
	lastErr map[string]JudgeError

	mu       sync.Mutex
	attached map[string][]string // "scope/id" → policy names
	loadedAt time.Time
	loading  bool
	merged   map[string]Policy // by the chain of names
}

// NewGate builds the gate. The library must be valid.
func NewGate(cfg Config, store Attachments, log *slog.Logger) (*Gate, error) {
	if err := cfg.Validate(); err != nil {
		return nil, err
	}
	if log == nil {
		log = slog.Default()
	}
	g := &Gate{cfg: cfg, store: store, log: log, now: time.Now,
		merged: map[string]Policy{}, m: newGateMetrics()}
	g.deny.Store(compileDeny(cfg))
	g.reload(context.Background())
	return g, nil
}

func compileDeny(cfg Config) *map[string]*regexp.Regexp {
	deny := map[string]*regexp.Regexp{}
	for _, p := range cfg.Policies {
		for _, r := range []Rules{p.Input, p.Output} {
			for _, d := range r.Deny {
				deny[d.Regex] = regexp.MustCompile(d.Regex)
			}
		}
	}
	return &deny
}

// Replace puts a new library of policies in force, for requests from the
// next one on: the cluster's settings changed. Refused, and the old one
// kept, when it is not valid.
func (g *Gate) Replace(cfg Config) error {
	if err := cfg.Validate(); err != nil {
		return err
	}
	// The patterns first: a policy naming one must find it compiled.
	old := *g.deny.Load()
	next := *compileDeny(cfg)
	for k, v := range old {
		if _, ok := next[k]; !ok {
			next[k] = v
		}
	}
	g.deny.Store(&next)
	g.mu.Lock()
	g.cfg = cfg
	g.merged = map[string]Policy{}
	g.mu.Unlock()
	g.reload(context.Background())
	return nil
}

// Library is the policies in force.
func (g *Gate) Library() Config {
	g.mu.Lock()
	defer g.mu.Unlock()
	return g.cfg
}

// attachmentsTTL is how stale an attachment may be: one made in the console
// applies within it.
const attachmentsTTL = 30 * time.Second

func (g *Gate) reload(ctx context.Context) {
	if g.store == nil {
		return
	}
	list, err := g.store.GuardAttachments(ctx)
	g.mu.Lock()
	defer g.mu.Unlock()
	g.loading = false
	if err != nil {
		// Kept as they were: a store that cannot be read for a moment is no
		// reason to drop every policy.
		g.log.Warn("could not read guard policy attachments; keeping the last ones", "err", err)
		return
	}
	g.attached = map[string][]string{}
	for _, a := range list {
		if _, ok := g.cfg.Policies[a.Policy]; !ok {
			g.log.Warn("a guard policy is attached that the cluster's settings do not define; ignored",
				"policy", a.Policy, "scope", a.Scope, "id", a.ScopeID)
			continue
		}
		k := a.Scope + "/" + a.ScopeID
		g.attached[k] = append(g.attached[k], a.Policy)
	}
	for k := range g.attached {
		sort.Strings(g.attached[k])
	}
	g.loadedAt = g.now()
	g.merged = map[string]Policy{}
}

// Invalidate has attachments read again on the next request: one just made
// applies at once rather than within attachmentsTTL.
func (g *Gate) Invalidate() { g.reload(context.Background()) }

// Policy is the policy in force for a request, and the names it was made
// from.
func (g *Gate) Policy(s Subject) (Policy, string) {
	g.mu.Lock()
	if !g.cfg.On() {
		g.mu.Unlock()
		return Policy{}, ""
	}
	if !g.loading && g.store != nil && g.now().Sub(g.loadedAt) > attachmentsTTL {
		g.loading = true
		go g.reload(context.WithoutCancel(context.Background()))
	}
	names := append([]string{}, g.cfg.Default...)
	for _, sc := range []struct{ scope, id string }{
		{"org", s.OrgID}, {"team", s.TeamID}, {"project", s.ProjectID}, {"user", s.UserID}, {"key", s.KeyID},
	} {
		if sc.id != "" {
			names = append(names, g.attached[sc.scope+"/"+sc.id]...)
		}
	}
	names = dedupe(names)
	key := strings.Join(names, ",")
	p, ok := g.merged[key]
	if !ok {
		ps := make([]Policy, len(names))
		for i, n := range names {
			ps[i] = g.cfg.Policies[n]
		}
		p = Merge(names, ps)
		g.merged[key] = p
	}
	g.mu.Unlock()
	return p, key
}

func dedupe(in []string) []string {
	seen := map[string]bool{}
	var out []string
	for _, s := range in {
		if !seen[s] {
			seen[s] = true
			out = append(out, s)
		}
	}
	return out
}

// finding is one reason to act, by name.
type finding struct {
	action   Action
	category string
	rule     string
	// failOpen is a judge that did not answer, under a policy that serves
	// anyway.
	failOpen bool
}

// judge scans one piece of text under one direction's rules, returning
// every finding and the spans to redact.
func (g *Gate) judge(text string, r Rules) ([]finding, []detect.Span) {
	var fs []finding
	var redact []detect.Span
	if len(r.PII) > 0 {
		for _, sp := range detect.Spans(text) {
			a, ok := actions[r.PII[sp.Kind]]
			if !ok {
				continue
			}
			fs = append(fs, finding{action: a, category: sp.Kind, rule: "pii:" + sp.Kind})
			if a == Redact {
				redact = append(redact, sp)
			}
		}
	}
	if r.Injection != "" {
		for _, k := range detect.Injection(text) {
			fs = append(fs, finding{action: actions[r.Injection], category: "injection", rule: k})
		}
	}
	if len(r.Deny) > 0 {
		norm := detect.Normalize(text)
		for _, d := range r.Deny {
			re := (*g.deny.Load())[d.Regex]
			if re == nil {
				continue
			}
			a := actions[d.Action]
			inText, inNorm := re.MatchString(text), re.MatchString(norm)
			switch {
			case !inText && !inNorm:
				continue
			case a == Redact && inText:
				for _, m := range re.FindAllStringIndex(text, -1) {
					redact = append(redact, detect.Span{Kind: d.Name, Start: m[0], End: m[1]})
				}
			case a == Redact:
				// Only there once disguises are taken off, so there is no
				// place in the text to cut it out of: stopped instead.
				a = Block
			}
			fs = append(fs, finding{action: a, category: "deny", rule: d.Name})
		}
	}
	return fs, redact
}

// strictest picks the finding that decides.
func strictest(fs []finding) (finding, bool) {
	var best finding
	found := false
	for _, f := range fs {
		if !found || rank(f.action) > rank(best.action) {
			best, found = f, true
		}
	}
	return best, found
}

// placeholders rewrites spans as [KIND_n], the same value always the same
// placeholder within one request.
type placeholders struct {
	byValue map[string]string
	count   map[string]int
}

func newPlaceholders() *placeholders {
	return &placeholders{byValue: map[string]string{}, count: map[string]int{}}
}

func (p *placeholders) rewrite(text string, spans []detect.Span) string {
	if len(spans) == 0 {
		return text
	}
	sort.Slice(spans, func(i, j int) bool { return spans[i].Start < spans[j].Start })
	var b strings.Builder
	last := 0
	for _, s := range spans {
		if s.Start < last {
			continue // overlapping: the earlier already covers it
		}
		v := text[s.Start:s.End]
		ph, ok := p.byValue[v]
		if !ok {
			p.count[s.Kind]++
			ph = fmt.Sprintf("[%s_%d]", strings.ToUpper(s.Kind), p.count[s.Kind])
			p.byValue[v] = ph
		}
		b.WriteString(text[last:s.Start])
		b.WriteString(ph)
		last = s.End
	}
	b.WriteString(text[last:])
	return b.String()
}

var _ Guard = (*Gate)(nil)

// Directives are the policy in force's demands of the gateway. Not in
// shadow mode, which changes nothing.
func (g *Gate) Directives(s Subject) Directives {
	p, _ := g.Policy(s)
	return Directives{NoStore: p.NoStore && p.Mode == "enforce"}
}

// Input judges a request body.
func (g *Gate) Input(ctx context.Context, s Subject, body []byte) (Verdict, error) {
	start := g.now()
	p, names := g.Policy(s)
	if !p.active() {
		return Verdict{}, nil
	}
	d, err := parseDoc(body)
	if err != nil {
		// Not this gate's to refuse: the engine says what is wrong with it.
		return Verdict{}, nil
	}
	v := Verdict{Stage: Input, Policy: names, Shadow: p.Mode == "shadow"}
	done := func(f finding) Verdict {
		v.Action, v.Category, v.Rule = f.action, f.category, f.rule
		if f.action == Block {
			v.Message = fmt.Sprintf("this request was refused by the gateway's policy (%s)", refusedFor(f))
			v.Body = nil
		}
		v.Took = g.now().Sub(start)
		return v
	}

	var fs []finding
	changed := false
	// Tools first: a denied tool is stripped before anything else is judged.
	if t := p.Tools; len(t.Allow) > 0 || len(t.Deny) > 0 {
		for _, tl := range d.tools() {
			if toolAllowed(t, tl.name) {
				continue
			}
			if t.OnDenied == "block" {
				fs = append(fs, finding{action: Block, category: "tool", rule: "tool:" + tl.name})
			} else {
				fs = append(fs, finding{action: Flag, category: "tool", rule: "tool:" + tl.name})
			}
		}
		if t.OnDenied != "block" {
			d.keepTools(func(n string) bool { return toolAllowed(t, n) })
			changed = true
		}
	}

	var answer json.RawMessage
	if p.ValidateStructured {
		answer = d.responseSchema()
	}
	var offered []tool
	if p.Tools.ValidateArguments {
		offered = d.tools()
	}
	user := lastUser(segsOf(d))
	var judgeUser string
	if p.Judge != nil && p.Judge.Output != "" {
		judgeUser = user
	}
	g.tools.put(s.RequestID, offered, answer, judgeUser, start)
	if p.RequireCapabilities {
		if len(d.tools()) > 0 && !s.Tools {
			fs = append(fs, finding{action: Block, category: "capability", rule: "tools"})
		}
		if d.hasImages() && !s.Vision {
			fs = append(fs, finding{action: Block, category: "capability", rule: "vision"})
		}
	}

	segs := d.segments()
	total := 0
	ph := newPlaceholders()
	type rewrite struct {
		seg   segment
		text  string
		spans []detect.Span
	}
	var rewrites []rewrite
	for _, sg := range segs {
		total += len(sg.text)
		f, red := g.judge(sg.text, p.Input)
		fs = append(fs, f...)
		if len(red) > 0 {
			rewrites = append(rewrites, rewrite{sg, ph.rewrite(sg.text, red), red})
		}
	}
	if p.Input.MaxChars > 0 && total > p.Input.MaxChars {
		fs = append(fs, finding{action: Block, category: "size", rule: "max_chars"})
	}
	best, found := strictest(fs)
	if found && best.action == Block {
		return done(best), nil
	}
	// The judge last: what a pattern already stopped need not wait for it.
	if jp := p.Judge; jp != nil && jp.Input != "" && user != "" {
		res, err := g.judger.ask(ctx, jp, user, "")
		if err != nil {
			g.log.Warn("the judge gave no verdict on a request", "judge", jp.Name(), "err", err)
			g.noJudger(jp, err)
		}
		if err != nil && p.Fail == "closed" && !v.Shadow {
			v.Took = g.now().Sub(start)
			return v, fmt.Errorf("the judge could not be asked: %w", err)
		}
		jf := judgeFinding(p, jp.Input, res, err)
		if jf.failOpen {
			g.m.failOpen.Inc()
		}
		fs = append(fs, jf)
		if best, found = strictest(fs); found && best.action == Block {
			return done(best), nil
		}
	}
	if !v.Shadow {
		for _, rw := range rewrites {
			g.countRedactions(rw.spans)
			rw.seg.set(rw.text)
			changed = true
		}
		if d.pin(p.SystemPrompt, p.DropClientSystem) {
			changed = true
		}
		if d.capOutput(p.MaxOutputTokens) {
			changed = true
		}
		if changed {
			v.Body = d.encode()
		}
	}
	if found {
		return done(best), nil
	}
	v.Took = g.now().Sub(start)
	return v, nil
}

// judgeFinding is what a judge's answer, or its failing to give one, comes
// to under a policy.
// SetJudgeAlarm says whom to tell when a judge gives no verdict: set before
// the gate serves.
func (g *Gate) SetJudgeAlarm(f func(model string, err error)) { g.alarm = f }

// noJudger records, against each judge a policy asks, that there was no way
// to ask any: a gate built without SetJudge. With one, each judge's failure
// is recorded as it is asked.
func (g *Gate) noJudger(jp *JudgePolicy, err error) {
	if g.judger != nil {
		return
	}
	for _, m := range jp.Judges() {
		g.judgeFailed(m, err)
	}
}

func (g *Gate) judgeFailed(model string, err error) {
	g.errMu.Lock()
	if g.lastErr == nil {
		g.lastErr = map[string]JudgeError{}
	}
	g.lastErr[model] = JudgeError{Err: err.Error(), At: g.now()}
	g.errMu.Unlock()
	if g.alarm != nil {
		g.alarm(model, err)
	}
}

// JudgeError is a judge's last failure to give a verdict.
type JudgeError struct {
	Err string
	At  time.Time
}

// JudgeErrors are each judge's last failure to give a verdict, by model.
func (g *Gate) JudgeErrors() map[string]JudgeError {
	g.errMu.Lock()
	defer g.errMu.Unlock()
	return maps.Clone(g.lastErr)
}

func judgeFinding(p Policy, action string, res judgment, err error) finding {
	if err != nil {
		// Unanswered: closed stops it, open lets it through, said.
		if p.Fail == "closed" {
			return finding{action: Block, category: "judge", rule: "unavailable"}
		}
		return finding{action: Flag, category: "judge", rule: "unavailable", failOpen: true}
	}
	if ok, cat := p.Judge.counts(res); ok {
		return finding{action: actions[action], category: "judge", rule: cat}
	}
	return finding{}
}

func segsOf(d *doc) []segment { return d.segments() }

func (g *Gate) countRedactions(spans []detect.Span) {
	for _, s := range spans {
		g.m.redactions.WithLabelValues(s.Kind).Inc()
	}
}

// lastUser is the last thing the user said: what a judge reads.
func lastUser(segs []segment) string {
	for i := len(segs) - 1; i >= 0; i-- {
		if segs[i].role == "user" {
			return segs[i].text
		}
	}
	return ""
}

func toolAllowed(t Tools, name string) bool {
	for _, d := range t.Deny {
		if d == name {
			return false
		}
	}
	if len(t.Allow) == 0 {
		return true
	}
	for _, a := range t.Allow {
		if a == name {
			return true
		}
	}
	return false
}

// refusedFor is what a refusal names: the category, and for a judge the
// hazard it saw — "judge: S6,Specialized Advice" — so whoever is refused
// can say why, and a judge refusing ordinary questions can be told apart.
// Never the text.
func refusedFor(f finding) string {
	if f.category == "judge" && f.rule != "" && f.rule != "unavailable" {
		return "judge: " + f.rule
	}
	return f.category
}