feat(guardrail): 输入/输出护栏（关键词/敏感信息/长度，pipeline fail-fast）
This commit is contained in:
Aska committed
1 parent
e33b005612
commit
02aae2bbb4
3 files changed
+863
No files matched your search
@@ -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