feat(guardrail): 输入/输出护栏(关键词/敏感信息/长度,pipeline fail-fast)
This commit is contained in:
1 parent
e33b005612
commit
02aae2bbb4
3 files changed
+863
No files matched your search
@@ -0,0 +1,247 @@
|
|||||||
|
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
|
||||||
|
}
|
||||||
@@ -0,0 +1,166 @@
|
|||||||
|
// Package guardrail provides input/output validation and safety interception
|
||||||
|
// for the Cyrene agent runtime.
|
||||||
|
//
|
||||||
|
// It implements the Guardrail design in DevDocs/design/04-运行时.md §5: one or
|
||||||
|
// more checks run in parallel with agent execution and fail fast on the first
|
||||||
|
// denial. A guardrail evaluates a model.Request (before it reaches the LLM) or
|
||||||
|
// a model.Response (before it is returned to the user) and returns a Verdict
|
||||||
|
// that either allows the payload or blocks it with a human-readable Reason.
|
||||||
|
//
|
||||||
|
// The package ships three built-in, configurable checks:
|
||||||
|
//
|
||||||
|
// - KeywordGuard: content-policy keyword/category matching.
|
||||||
|
// - SensitiveGuard: regex detection of API keys, tokens, passwords and other
|
||||||
|
// secrets, in either block or mark mode. It never echoes the matched secret
|
||||||
|
// into a Reason, log or return value.
|
||||||
|
// - LengthGuard: input/output character-length limits.
|
||||||
|
//
|
||||||
|
// Checks are composed into a Pipeline (see New). All errors surfaced by a
|
||||||
|
// Pipeline wrap the sentinel ErrGuardrail, so callers can test with errors.Is.
|
||||||
|
//
|
||||||
|
// Guardrail is a leaf package: it depends only on the standard library and
|
||||||
|
// internal/runtime/model types, and it never imports the agent loop or any other
|
||||||
|
// runtime package, so it can sit alongside the loop without an import cycle.
|
||||||
|
package guardrail
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"errors"
|
||||||
|
"fmt"
|
||||||
|
|
||||||
|
"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/model"
|
||||||
|
)
|
||||||
|
|
||||||
|
// Verdict is the outcome of a single guardrail evaluation.
|
||||||
|
//
|
||||||
|
// When Allow is true the payload is accepted; when it is false the payload is
|
||||||
|
// blocked and Reason explains why. Reason is also set on an allow verdict to
|
||||||
|
// flag context (for example a sensitive-information mark) without carrying the
|
||||||
|
// offending secret.
|
||||||
|
type Verdict struct {
|
||||||
|
Allow bool
|
||||||
|
Reason string
|
||||||
|
}
|
||||||
|
|
||||||
|
// Guardrail evaluates a request or a response and returns a Verdict.
|
||||||
|
//
|
||||||
|
// Implementations must be safe for concurrent use and honor ctx cancellation:
|
||||||
|
// a canceled context returns immediately with a non-nil error. A nil Verdict
|
||||||
|
// on a nil error is treated as allow by the Pipeline.
|
||||||
|
type Guardrail interface {
|
||||||
|
// CheckInput evaluates a completion request before it reaches the LLM.
|
||||||
|
CheckInput(ctx context.Context, input *model.Request) (*Verdict, error)
|
||||||
|
// CheckOutput evaluates a completion response before it is returned.
|
||||||
|
CheckOutput(ctx context.Context, output *model.Response) (*Verdict, error)
|
||||||
|
}
|
||||||
|
|
||||||
|
// ErrGuardrail is the sentinel wrapped around any error returned by a
|
||||||
|
// Pipeline's CheckInput or CheckOutput. Use errors.Is to detect it.
|
||||||
|
var ErrGuardrail = errors.New("guardrail: evaluation failed")
|
||||||
|
|
||||||
|
// wrapErr attaches the ErrGuardrail sentinel to a guard's internal error so the
|
||||||
|
// caller can tell a guardrail failure apart from any other runtime error.
|
||||||
|
func wrapErr(err error) error {
|
||||||
|
return fmt.Errorf("%w: %w", ErrGuardrail, err)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Pipeline runs a sequence of Guardrails in the order they were registered and
|
||||||
|
// fails fast: the first deny verdict short-circuits the rest and is returned
|
||||||
|
// as-is. Because the sentinel is non-fatal, allow verdicts let the remaining
|
||||||
|
// guards run, and the reason of the last flagged verdict is preserved.
|
||||||
|
type Pipeline struct {
|
||||||
|
guards []Guardrail
|
||||||
|
}
|
||||||
|
|
||||||
|
// New builds a Pipeline from the given options. Options are applied in order,
|
||||||
|
// so the guards run in the same order they are passed. With no options the
|
||||||
|
// pipeline allows every payload.
|
||||||
|
func New(opts ...Option) *Pipeline {
|
||||||
|
p := &Pipeline{}
|
||||||
|
for _, o := range opts {
|
||||||
|
o(p)
|
||||||
|
}
|
||||||
|
return p
|
||||||
|
}
|
||||||
|
|
||||||
|
// Option configures a Pipeline. Options are applied in the order passed to New.
|
||||||
|
type Option func(*Pipeline)
|
||||||
|
|
||||||
|
// WithGuard appends a custom Guardrail to the pipeline.
|
||||||
|
func WithGuard(g Guardrail) Option {
|
||||||
|
return func(p *Pipeline) { p.guards = append(p.guards, g) }
|
||||||
|
}
|
||||||
|
|
||||||
|
// WithKeywords appends a KeywordGuard that blocks any payload containing one of
|
||||||
|
// the given keywords (case-insensitive substring match).
|
||||||
|
func WithKeywords(keywords ...string) Option {
|
||||||
|
return func(p *Pipeline) { p.guards = append(p.guards, NewKeywordGuard(keywords...)) }
|
||||||
|
}
|
||||||
|
|
||||||
|
// WithSensitiveInfo appends a SensitiveGuard using the built-in default patterns
|
||||||
|
// in the given mode. Use WithSensitivePatterns to supply your own patterns.
|
||||||
|
func WithSensitiveInfo(mode SensitiveMode) Option {
|
||||||
|
return WithSensitivePatterns(mode, DefaultSensitivePatterns()...)
|
||||||
|
}
|
||||||
|
|
||||||
|
// WithSensitivePatterns appends a SensitiveGuard that matches the given patterns
|
||||||
|
// in the given mode.
|
||||||
|
func WithSensitivePatterns(mode SensitiveMode, patterns ...SensitivePattern) Option {
|
||||||
|
return func(p *Pipeline) { p.guards = append(p.guards, NewSensitiveGuard(patterns, mode)) }
|
||||||
|
}
|
||||||
|
|
||||||
|
// WithLengthLimits appends a LengthGuard capping input and output in characters.
|
||||||
|
// A zero limit disables that side.
|
||||||
|
func WithLengthLimits(maxInputChars, maxOutputChars int) Option {
|
||||||
|
return func(p *Pipeline) { p.guards = append(p.guards, NewLengthGuard(maxInputChars, maxOutputChars)) }
|
||||||
|
}
|
||||||
|
|
||||||
|
// CheckInput runs every registered guard against the request in order and
|
||||||
|
// returns the first deny verdict. An internal guard error is wrapped with
|
||||||
|
// ErrGuardrail.
|
||||||
|
func (p *Pipeline) CheckInput(ctx context.Context, input *model.Request) (*Verdict, error) {
|
||||||
|
return p.run(ctx, input)
|
||||||
|
}
|
||||||
|
|
||||||
|
// CheckOutput runs every registered guard against the response in order and
|
||||||
|
// returns the first deny verdict. An internal guard error is wrapped with
|
||||||
|
// ErrGuardrail.
|
||||||
|
func (p *Pipeline) CheckOutput(ctx context.Context, output *model.Response) (*Verdict, error) {
|
||||||
|
return p.run(ctx, output)
|
||||||
|
}
|
||||||
|
|
||||||
|
// run is the shared fail-fast loop. eval dispatches the payload to the guard
|
||||||
|
// and returns a verdict, an error, or both (an error wins). The first deny
|
||||||
|
// verdict stops the loop, while a flagged-allow verdict is remembered as the
|
||||||
|
// final reason.
|
||||||
|
func (p *Pipeline) run(ctx context.Context, payload any) (*Verdict, error) {
|
||||||
|
var reason string
|
||||||
|
for _, g := range p.guards {
|
||||||
|
v, err := p.eval(ctx, g, payload)
|
||||||
|
if err != nil {
|
||||||
|
return nil, wrapErr(err)
|
||||||
|
}
|
||||||
|
if v == nil {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
if !v.Allow {
|
||||||
|
return v, nil // fail-fast: first denial
|
||||||
|
}
|
||||||
|
if v.Reason != "" {
|
||||||
|
reason = v.Reason
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return &Verdict{Allow: true, Reason: reason}, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// eval dispatches a payload to the matching guard method.
|
||||||
|
func (p *Pipeline) eval(ctx context.Context, g Guardrail, payload any) (*Verdict, error) {
|
||||||
|
switch t := payload.(type) {
|
||||||
|
case *model.Request:
|
||||||
|
return g.CheckInput(ctx, t)
|
||||||
|
case *model.Response:
|
||||||
|
return g.CheckOutput(ctx, t)
|
||||||
|
default:
|
||||||
|
return nil, fmt.Errorf("guardrail: unsupported payload %T", payload)
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,450 @@
|
|||||||
|
package guardrail
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"errors"
|
||||||
|
"regexp"
|
||||||
|
"strings"
|
||||||
|
"testing"
|
||||||
|
|
||||||
|
"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/model"
|
||||||
|
)
|
||||||
|
|
||||||
|
// textRequest builds a text-only request from the given message bodies.
|
||||||
|
func textRequest(t *testing.T, texts ...string) *model.Request {
|
||||||
|
t.Helper()
|
||||||
|
msgs := make([]model.Message, 0, len(texts))
|
||||||
|
for _, s := range texts {
|
||||||
|
msgs = append(msgs, model.TextMessage(model.RoleUser, s))
|
||||||
|
}
|
||||||
|
return &model.Request{Messages: msgs}
|
||||||
|
}
|
||||||
|
|
||||||
|
// textResponse builds a response carrying content in both Content and Message,
|
||||||
|
// mirroring how the adapter populates it.
|
||||||
|
func textResponse(t *testing.T, content string) *model.Response {
|
||||||
|
t.Helper()
|
||||||
|
return &model.Response{
|
||||||
|
Content: content,
|
||||||
|
Message: model.TextMessage(model.RoleAssistant, content),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// stub is a configurable Guardrail used to exercise pipeline ordering, fail-fast
|
||||||
|
// and error wrapping without depending on a concrete implementation.
|
||||||
|
type stub struct {
|
||||||
|
inFn func(context.Context, *model.Request) (*Verdict, error)
|
||||||
|
outFn func(context.Context, *model.Response) (*Verdict, error)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *stub) CheckInput(ctx context.Context, r *model.Request) (*Verdict, error) {
|
||||||
|
if s.inFn != nil {
|
||||||
|
return s.inFn(ctx, r)
|
||||||
|
}
|
||||||
|
return allow(), nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *stub) CheckOutput(ctx context.Context, r *model.Response) (*Verdict, error) {
|
||||||
|
if s.outFn != nil {
|
||||||
|
return s.outFn(ctx, r)
|
||||||
|
}
|
||||||
|
return allow(), nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestKeywordGuard_BlocksOnViolence(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
ctx := context.Background()
|
||||||
|
|
||||||
|
g := NewKeywordGuard("violence", "hate speech")
|
||||||
|
v, err := g.CheckInput(ctx, textRequest(t, "please give me instructions on violence"))
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("CheckInput() error: %v", err)
|
||||||
|
}
|
||||||
|
if v.Allow {
|
||||||
|
t.Fatalf("CheckInput() = %+v, want deny", v)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestKeywordGuard_MatchesCaseInsensitively(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
ctx := context.Background()
|
||||||
|
|
||||||
|
g := NewKeywordGuard("violence")
|
||||||
|
v, err := g.CheckInput(ctx, textRequest(t, "VIOLENCE is never acceptable"))
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("CheckInput() error: %v", err)
|
||||||
|
}
|
||||||
|
if v.Allow {
|
||||||
|
t.Fatalf("CheckInput() = %+v, want deny for uppercase keyword", v)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestKeywordGuard_BlocksOutput(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
ctx := context.Background()
|
||||||
|
|
||||||
|
g := NewKeywordGuard("hate speech")
|
||||||
|
v, err := g.CheckOutput(ctx, textResponse(t, "the answer involves hate speech"))
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("CheckOutput() error: %v", err)
|
||||||
|
}
|
||||||
|
if v.Allow {
|
||||||
|
t.Fatalf("CheckOutput() = %+v, want deny", v)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestKeywordGuard_AllowsWhenNoMatch(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
ctx := context.Background()
|
||||||
|
|
||||||
|
g := NewKeywordGuard("violence", "hate speech")
|
||||||
|
v, err := g.CheckInput(ctx, textRequest(t, "a friendly greeting"))
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("CheckInput() error: %v", err)
|
||||||
|
}
|
||||||
|
if !v.Allow {
|
||||||
|
t.Fatalf("CheckInput() = %+v, want allow", v)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestKeywordGuard_AllowsEmptyKeywordList(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
ctx := context.Background()
|
||||||
|
|
||||||
|
g := NewKeywordGuard()
|
||||||
|
v, err := g.CheckInput(ctx, textRequest(t, "anything at all"))
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("CheckInput() error: %v", err)
|
||||||
|
}
|
||||||
|
if !v.Allow {
|
||||||
|
t.Fatalf("CheckInput() = %+v, want allow for empty keyword list", v)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestSensitiveGuard_BlocksOnMatch(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
ctx := context.Background()
|
||||||
|
|
||||||
|
patterns := []SensitivePattern{{
|
||||||
|
Label: "openai_key",
|
||||||
|
Regexp: regexp.MustCompile(`\bsk-[A-Za-z0-9_-]{20,}\b`),
|
||||||
|
}}
|
||||||
|
g := NewSensitiveGuard(patterns, SensitiveBlock)
|
||||||
|
|
||||||
|
secret := "Bearer token is sk-1234567890abcdefghijkl"
|
||||||
|
v, err := g.CheckInput(ctx, textRequest(t, secret))
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("CheckInput() error: %v", err)
|
||||||
|
}
|
||||||
|
if v.Allow {
|
||||||
|
t.Fatalf("CheckInput() = %+v, want deny", v)
|
||||||
|
}
|
||||||
|
if !strings.Contains(v.Reason, "openai_key") {
|
||||||
|
t.Fatalf("Reason = %q, want it to name the matched pattern", v.Reason)
|
||||||
|
}
|
||||||
|
if strings.Contains(v.Reason, "sk-1234567890abcdefghijkl") {
|
||||||
|
t.Fatalf("Reason = %q, must not echo the secret", v.Reason)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestSensitiveGuard_MarksOnMatch(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
ctx := context.Background()
|
||||||
|
|
||||||
|
patterns := []SensitivePattern{{
|
||||||
|
Label: "mysql_password",
|
||||||
|
Regexp: regexp.MustCompile(`(?i)password\s*[:=]\s*\S+`),
|
||||||
|
}}
|
||||||
|
g := NewSensitiveGuard(patterns, SensitiveMark)
|
||||||
|
|
||||||
|
v, err := g.CheckInput(ctx, textRequest(t, "connect using password=hunter2"))
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("CheckInput() error: %v", err)
|
||||||
|
}
|
||||||
|
if !v.Allow {
|
||||||
|
t.Fatalf("CheckInput() = %+v, want allow in mark mode", v)
|
||||||
|
}
|
||||||
|
if !strings.Contains(v.Reason, "mysql_password") {
|
||||||
|
t.Fatalf("Reason = %q, want the flagged label", v.Reason)
|
||||||
|
}
|
||||||
|
if strings.Contains(v.Reason, "hunter2") {
|
||||||
|
t.Fatalf("Reason = %q, must not echo the secret", v.Reason)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestSensitiveGuard_AllowsWhenNoMatch(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
ctx := context.Background()
|
||||||
|
|
||||||
|
patterns := []SensitivePattern{{
|
||||||
|
Label: "openai_key",
|
||||||
|
Regexp: regexp.MustCompile(`\bsk-[A-Za-z0-9_-]{20,}\b`),
|
||||||
|
}}
|
||||||
|
g := NewSensitiveGuard(patterns, SensitiveBlock)
|
||||||
|
|
||||||
|
v, err := g.CheckInput(ctx, textRequest(t, "a harmless sentence"))
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("CheckInput() error: %v", err)
|
||||||
|
}
|
||||||
|
if !v.Allow {
|
||||||
|
t.Fatalf("CheckInput() = %+v, want allow", v)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestSensitiveGuard_SkipsNilPattern(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
ctx := context.Background()
|
||||||
|
|
||||||
|
g := NewSensitiveGuard([]SensitivePattern{{Label: "broken", Regexp: nil}}, SensitiveBlock)
|
||||||
|
v, err := g.CheckInput(ctx, textRequest(t, "any text"))
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("CheckInput() error: %v", err)
|
||||||
|
}
|
||||||
|
if !v.Allow {
|
||||||
|
t.Fatalf("CheckInput() = %+v, want allow when pattern regexp is nil", v)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestSensitiveGuard_DefaultPatternsCatchBearerToken(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
ctx := context.Background()
|
||||||
|
|
||||||
|
g := NewSensitiveGuard(DefaultSensitivePatterns(), SensitiveBlock)
|
||||||
|
v, err := g.CheckInput(ctx, textRequest(t, "Authorization: Bearer abc.def-ghi_jkl"))
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("CheckInput() error: %v", err)
|
||||||
|
}
|
||||||
|
if v.Allow {
|
||||||
|
t.Fatalf("CheckInput() = %+v, want deny for a bearer token", v)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestSensitiveGuard_ChecksOutput(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
ctx := context.Background()
|
||||||
|
|
||||||
|
patterns := []SensitivePattern{{
|
||||||
|
Label: "api_key",
|
||||||
|
Regexp: regexp.MustCompile(`(?i)api[_-]?key\s*[:=]\s*\S+`),
|
||||||
|
}}
|
||||||
|
g := NewSensitiveGuard(patterns, SensitiveBlock)
|
||||||
|
v, err := g.CheckOutput(ctx, textResponse(t, "API_KEY=sk-secret-value"))
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("CheckOutput() error: %v", err)
|
||||||
|
}
|
||||||
|
if v.Allow {
|
||||||
|
t.Fatalf("CheckOutput() = %+v, want deny", v)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestLengthGuard_BlocksOverLimit(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
ctx := context.Background()
|
||||||
|
|
||||||
|
tests := []struct {
|
||||||
|
name string
|
||||||
|
limit int
|
||||||
|
text string
|
||||||
|
want string
|
||||||
|
}{
|
||||||
|
{name: "input_over", limit: 5, text: "hello world", want: "input_length_exceeded"},
|
||||||
|
{name: "output_over", limit: 5, text: "a very long answer", want: "output_length_exceeded"},
|
||||||
|
{name: "unicode_runes_not_bytes", limit: 2, text: "你好啊", want: "input_length_exceeded"},
|
||||||
|
}
|
||||||
|
for _, tt := range tests {
|
||||||
|
tt := tt
|
||||||
|
t.Run(tt.name, func(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
g := NewLengthGuard(tt.limit, tt.limit)
|
||||||
|
var v *Verdict
|
||||||
|
var err error
|
||||||
|
if tt.name == "output_over" {
|
||||||
|
v, err = g.CheckOutput(ctx, textResponse(t, tt.text))
|
||||||
|
} else {
|
||||||
|
v, err = g.CheckInput(ctx, textRequest(t, tt.text))
|
||||||
|
}
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("check error: %v", err)
|
||||||
|
}
|
||||||
|
if v.Allow {
|
||||||
|
t.Fatalf("got %+v, want deny", v)
|
||||||
|
}
|
||||||
|
if !strings.Contains(v.Reason, tt.want) {
|
||||||
|
t.Fatalf("Reason = %q, want %q", v.Reason, tt.want)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestLengthGuard_AllowsWithinLimit(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
ctx := context.Background()
|
||||||
|
|
||||||
|
g := NewLengthGuard(20, 20)
|
||||||
|
v, err := g.CheckInput(ctx, textRequest(t, "ok"))
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("CheckInput() error: %v", err)
|
||||||
|
}
|
||||||
|
if !v.Allow {
|
||||||
|
t.Fatalf("CheckInput() = %+v, want allow under limit", v)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestLengthGuard_ZeroLimitIsUnlimited(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
ctx := context.Background()
|
||||||
|
|
||||||
|
g := NewLengthGuard(0, 0)
|
||||||
|
v, err := g.CheckInput(ctx, textRequest(t, strings.Repeat("x", 100)))
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("CheckInput() error: %v", err)
|
||||||
|
}
|
||||||
|
if !v.Allow {
|
||||||
|
t.Fatalf("CheckInput() = %+v, want allow for unlimited", v)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestPipeline_AllowsCleanPayload(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
ctx := context.Background()
|
||||||
|
|
||||||
|
p := New(
|
||||||
|
WithKeywords("violence", "hate"),
|
||||||
|
WithSensitiveInfo(SensitiveBlock),
|
||||||
|
WithLengthLimits(100, 100),
|
||||||
|
)
|
||||||
|
res, err := p.CheckInput(ctx, textRequest(t, "a friendly greeting"))
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("CheckInput() error: %v", err)
|
||||||
|
}
|
||||||
|
if !res.Allow {
|
||||||
|
t.Fatalf("CheckInput() = %+v, want allow", res)
|
||||||
|
}
|
||||||
|
|
||||||
|
res, err = p.CheckOutput(ctx, textResponse(t, "a helpful reply"))
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("CheckOutput() error: %v", err)
|
||||||
|
}
|
||||||
|
if !res.Allow {
|
||||||
|
t.Fatalf("CheckOutput() = %+v, want allow", res)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestPipeline_EmptyPipelineAllows(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
ctx := context.Background()
|
||||||
|
|
||||||
|
p := New()
|
||||||
|
v, err := p.CheckInput(ctx, textRequest(t, "anything"))
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("CheckInput() error: %v", err)
|
||||||
|
}
|
||||||
|
if !v.Allow {
|
||||||
|
t.Fatalf("CheckInput() = %+v, want allow with no guards", v)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestPipeline_FailFastOnInput(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
ctx := context.Background()
|
||||||
|
|
||||||
|
denyFirst := &stub{inFn: func(context.Context, *model.Request) (*Verdict, error) {
|
||||||
|
return deny("first"), nil
|
||||||
|
}}
|
||||||
|
keywordCalled := false
|
||||||
|
second := &stub{inFn: func(context.Context, *model.Request) (*Verdict, error) {
|
||||||
|
keywordCalled = true
|
||||||
|
return deny("second"), nil
|
||||||
|
}}
|
||||||
|
|
||||||
|
p := New(WithGuard(denyFirst), WithGuard(second))
|
||||||
|
v, err := p.CheckInput(ctx, textRequest(t, "text"))
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("CheckInput() error: %v", err)
|
||||||
|
}
|
||||||
|
if v.Allow || v.Reason != "first" {
|
||||||
|
t.Fatalf("verdict = %+v, want deny reason %q", v, "first")
|
||||||
|
}
|
||||||
|
if keywordCalled {
|
||||||
|
t.Fatal("second guard ran after a denial; want fail-fast")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestPipeline_FailFastOnOutput(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
ctx := context.Background()
|
||||||
|
|
||||||
|
denyFirst := &stub{outFn: func(context.Context, *model.Response) (*Verdict, error) {
|
||||||
|
return deny("blocked-output"), nil
|
||||||
|
}}
|
||||||
|
second := &stub{outFn: func(context.Context, *model.Response) (*Verdict, error) {
|
||||||
|
t.Fatal("second guard ran; want fail-fast")
|
||||||
|
return allow(), nil
|
||||||
|
}}
|
||||||
|
p := New(WithGuard(denyFirst), WithGuard(second))
|
||||||
|
v, err := p.CheckOutput(ctx, textResponse(t, "text"))
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("CheckOutput() error: %v", err)
|
||||||
|
}
|
||||||
|
if v.Allow || v.Reason != "blocked-output" {
|
||||||
|
t.Fatalf("verdict = %+v, want deny reason %q", v, "blocked-output")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestPipeline_ErrorWrapsErrGuardrail(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
ctx := context.Background()
|
||||||
|
|
||||||
|
boom := errors.New("boom")
|
||||||
|
g := &stub{inFn: func(context.Context, *model.Request) (*Verdict, error) {
|
||||||
|
return nil, boom
|
||||||
|
}}
|
||||||
|
p := New(WithGuard(g))
|
||||||
|
_, err := p.CheckInput(ctx, textRequest(t, "text"))
|
||||||
|
if err == nil {
|
||||||
|
t.Fatal("CheckInput() returned nil error, want ErrGuardrail")
|
||||||
|
}
|
||||||
|
if !errors.Is(err, ErrGuardrail) {
|
||||||
|
t.Fatalf("err = %v, want it to wrap ErrGuardrail", err)
|
||||||
|
}
|
||||||
|
if !errors.Is(err, boom) {
|
||||||
|
t.Fatalf("err = %v, want it to preserve the cause", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestPipeline_CanceledContextWraps(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
|
||||||
|
ctx, cancel := context.WithCancel(context.Background())
|
||||||
|
cancel()
|
||||||
|
g := NewKeywordGuard("x")
|
||||||
|
p := New(WithGuard(g))
|
||||||
|
_, err := p.CheckInput(ctx, textRequest(t, "text"))
|
||||||
|
if err == nil {
|
||||||
|
t.Fatal("CheckInput() returned nil error, want wrapped ctx error")
|
||||||
|
}
|
||||||
|
if !errors.Is(err, ErrGuardrail) {
|
||||||
|
t.Fatalf("err = %v, want it to wrap ErrGuardrail", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestPipeline_PreservesMarkReason(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
ctx := context.Background()
|
||||||
|
|
||||||
|
patterns := []SensitivePattern{{
|
||||||
|
Label: "api_key",
|
||||||
|
Regexp: regexp.MustCompile(`(?i)api[_-]?key\s*[:=]\s*\S+`),
|
||||||
|
}}
|
||||||
|
p := New(WithSensitivePatterns(SensitiveMark, patterns...))
|
||||||
|
v, err := p.CheckInput(ctx, textRequest(t, "API_KEY=secret"))
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("CheckInput() error: %v", err)
|
||||||
|
}
|
||||||
|
if !v.Allow {
|
||||||
|
t.Fatalf("CheckInput() = %+v, want allow in mark mode", v)
|
||||||
|
}
|
||||||
|
if !strings.Contains(v.Reason, "api_key") {
|
||||||
|
t.Fatalf("Reason = %q, want the flagged label preserved", v.Reason)
|
||||||
|
}
|
||||||
|
}
|
||||||
Reference in new issue
Block a user