Cyrene_v2/internal/runtime/guardrail/builtin.go

package guardrail

import (
	"context"
	"errors"
	"regexp"
	"strings"
	"unicode/utf8"

	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/model"
)

// allow returns a passing verdict.
func allow() *Verdict { return &Verdict{Allow: true} }

// deny returns a blocking verdict with the given reason.
func deny(reason string) *Verdict { return &Verdict{Allow: false, Reason: reason} }

// requestText returns the concatenation of every text-bearing content part in a
// request's messages. Image parts are ignored.
func requestText(input *model.Request) string {
	var b strings.Builder
	for _, m := range input.Messages {
		b.WriteString(messageText(m))
	}
	return b.String()
}

// responseText returns the text-bearing output of a response. It prefers the
// convenience Content field and falls back to the normalized Message text parts
// when Content is empty, so a response that duplicates its content in both is
// measured once rather than twice.
func responseText(output *model.Response) string {
	if output.Content != "" {
		return output.Content
	}
	var b strings.Builder
	for _, c := range output.Message.Content {
		if c.Type == model.ContentTypeText {
			b.WriteString(c.Text)
		}
	}
	return b.String()
}

// messageText returns the text-bearing content parts of a single message.
func messageText(msg model.Message) string {
	var b strings.Builder
	for _, c := range msg.Content {
		if c.Type == model.ContentTypeText {
			b.WriteString(c.Text)
		}
	}
	return b.String()
}

// KeywordGuard is a content-policy guardrail that blocks any input or output
// containing one of its configured keywords. Matching is a case-insensitive
// substring search, so each keyword is matched as a plain fragment rather than
// a regexp.
//
// NewKeywordGuard builds a KeywordGuard from a word list. With no keywords the
// guard allows everything.
type KeywordGuard struct {
	keywords []string // lowercased
}

// NewKeywordGuard builds a KeywordGuard that blocks payloads containing any of
// the given keywords.
func NewKeywordGuard(keywords ...string) *KeywordGuard {
	normalized := make([]string, 0, len(keywords))
	for _, k := range keywords {
		if k = strings.ToLower(strings.TrimSpace(k)); k != "" {
			normalized = append(normalized, k)
		}
	}
	return &KeywordGuard{keywords: normalized}
}

// CheckInput blocks the request when any keyword is present in its text.
func (g *KeywordGuard) CheckInput(ctx context.Context, input *model.Request) (*Verdict, error) {
	if err := ctx.Err(); err != nil {
		return nil, err
	}
	if input == nil {
		return nil, errors.New("guardrail: nil request")
	}
	return g.check(requestText(input))
}

// CheckOutput blocks the response when any keyword is present in its text.
func (g *KeywordGuard) CheckOutput(ctx context.Context, output *model.Response) (*Verdict, error) {
	if err := ctx.Err(); err != nil {
		return nil, err
	}
	if output == nil {
		return nil, errors.New("guardrail: nil response")
	}
	return g.check(responseText(output))
}

func (g *KeywordGuard) check(text string) (*Verdict, error) {
	if len(g.keywords) == 0 {
		return allow(), nil
	}
	lower := strings.ToLower(text)
	for _, kw := range g.keywords {
		if strings.Contains(lower, kw) {
			return deny("content_policy_blocked: " + kw), nil
		}
	}
	return allow(), nil
}

// SensitiveMode selects how a SensitiveGuard reacts to a match.
type SensitiveMode int

const (
	// SensitiveBlock denies any payload that matches a pattern.
	SensitiveBlock SensitiveMode = iota
	// SensitiveMark allows the payload but flags it in Verdict.Reason.
	SensitiveMark
)

// SensitivePattern pairs a compiled regular expression with a human label. The
// label is what appears in Verdict.Reason; the pattern's captured secret is
// never echoed anywhere.
type SensitivePattern struct {
	Label  string
	Regexp *regexp.Regexp
}

// DefaultSensitivePatterns returns the built-in set of patterns used to detect
// common API keys, bearer tokens, cloud credentials, passwords and private-key
// material. Callers may extend or replace these with their own patterns.
func DefaultSensitivePatterns() []SensitivePattern {
	return []SensitivePattern{
		{Label: "api_key", Regexp: regexp.MustCompile(`(?i)\bapi[_-]?key\s*[:=]\s*\S+`)},
		{Label: "openai_key", Regexp: regexp.MustCompile(`\bsk-[A-Za-z0-9_-]{20,}\b`)},
		{Label: "aws_access_key", Regexp: regexp.MustCompile(`\bAKIA[0-9A-Z]{16}\b`)},
		{Label: "bearer_token", Regexp: regexp.MustCompile(`(?i)\bbearer\s+[A-Za-z0-9._~-]+`)},
		{Label: "password", Regexp: regexp.MustCompile(`(?i)\b(?:bind_?)?password\s*[:=]\s*\S+`)},
		{Label: "private_key", Regexp: regexp.MustCompile(`-----BEGIN (?:RSA |EC |DSA |OPENSSH |PGP )?PRIVATE KEY-----`)},
	}
}

// SensitiveGuard is a guardrail that detects sensitive information such as API
// keys, tokens and passwords using regular expressions. Matching is
// case-insensitive where the pattern says so; it never writes the matched secret
// into a Reason, log or return value.
//
// NewSensitiveGuard builds a SensitiveGuard from patterns and a mode.
type SensitiveGuard struct {
	patterns []SensitivePattern
	mode     SensitiveMode
}

// NewSensitiveGuard builds a SensitiveGuard. In SensitiveBlock mode a match
// denies the payload; in SensitiveMark mode a match allows it but records the
// flagged label in Verdict.Reason.
func NewSensitiveGuard(patterns []SensitivePattern, mode SensitiveMode) *SensitiveGuard {
	return &SensitiveGuard{patterns: patterns, mode: mode}
}

// CheckInput evaluates the request text against every pattern.
func (g *SensitiveGuard) CheckInput(ctx context.Context, input *model.Request) (*Verdict, error) {
	if err := ctx.Err(); err != nil {
		return nil, err
	}
	if input == nil {
		return nil, errors.New("guardrail: nil request")
	}
	return g.check(requestText(input))
}

// CheckOutput evaluates the response text against every pattern.
func (g *SensitiveGuard) CheckOutput(ctx context.Context, output *model.Response) (*Verdict, error) {
	if err := ctx.Err(); err != nil {
		return nil, err
	}
	if output == nil {
		return nil, errors.New("guardrail: nil response")
	}
	return g.check(responseText(output))
}

func (g *SensitiveGuard) check(text string) (*Verdict, error) {
	for _, p := range g.patterns {
		if p.Regexp == nil || !p.Regexp.MatchString(text) {
			continue
		}
		if g.mode == SensitiveBlock {
			return deny("sensitive_information_blocked: " + p.Label), nil
		}
		return &Verdict{Allow: true, Reason: "sensitive_information_detected: " + p.Label}, nil
	}
	return allow(), nil
}

// LengthGuard is a guardrail that blocks input or output exceeding a configured
// character limit. A limit of zero disables that side.
//
// NewLengthGuard builds a LengthGuard with per-side limits.
type LengthGuard struct {
	maxInputChars  int
	maxOutputChars int
}

// NewLengthGuard builds a LengthGuard. Limits are measured in Unicode code
// points; a limit of zero means unlimited.
func NewLengthGuard(maxInputChars, maxOutputChars int) *LengthGuard {
	return &LengthGuard{maxInputChars: maxInputChars, maxOutputChars: maxOutputChars}
}

// CheckInput blocks the request when its text exceeds the input character limit.
func (g *LengthGuard) CheckInput(ctx context.Context, input *model.Request) (*Verdict, error) {
	if err := ctx.Err(); err != nil {
		return nil, err
	}
	if input == nil {
		return nil, errors.New("guardrail: nil request")
	}
	if g.maxInputChars <= 0 {
		return allow(), nil
	}
	if n := utf8.RuneCountInString(requestText(input)); n > g.maxInputChars {
		return deny("input_length_exceeded"), nil
	}
	return allow(), nil
}

// CheckOutput blocks the response when its text exceeds the output limit.
func (g *LengthGuard) CheckOutput(ctx context.Context, output *model.Response) (*Verdict, error) {
	if err := ctx.Err(); err != nil {
		return nil, err
	}
	if output == nil {
		return nil, errors.New("guardrail: nil response")
	}
	if g.maxOutputChars <= 0 {
		return allow(), nil
	}
	if n := utf8.RuneCountInString(responseText(output)); n > g.maxOutputChars {
		return deny("output_length_exceeded"), nil
	}
	return allow(), nil
}