80 lines
2.1 KiB
Go
80 lines
2.1 KiB
Go
package canary
|
|
|
|
import (
|
|
"context"
|
|
"crypto/rand"
|
|
"encoding/hex"
|
|
"fmt"
|
|
"strings"
|
|
|
|
"github.com/sipeed/picoclaw/pkg/agent"
|
|
"github.com/sipeed/picoclaw/pkg/providers"
|
|
)
|
|
|
|
// Hook implements the agent.LLMInterceptor interface to detect system prompt leakage.
|
|
type Hook struct {
|
|
token string
|
|
}
|
|
|
|
// Ensure Hook implements LLMInterceptor.
|
|
var _ agent.LLMInterceptor = (*Hook)(nil)
|
|
|
|
// NewHook creates a new canary hook with a random token.
|
|
func NewHook() (*Hook, error) {
|
|
b := make([]byte, 16)
|
|
if _, err := rand.Read(b); err != nil {
|
|
return nil, fmt.Errorf("generate canary token: %w", err)
|
|
}
|
|
return &Hook{
|
|
token: "CANARY-" + hex.EncodeToString(b),
|
|
}, nil
|
|
}
|
|
|
|
// Token returns the current canary token (for testing).
|
|
func (h *Hook) Token() string {
|
|
return h.token
|
|
}
|
|
|
|
func (h *Hook) BeforeLLM(ctx context.Context, req *agent.LLMHookRequest) (*agent.LLMHookRequest, agent.HookDecision, error) {
|
|
if req == nil {
|
|
return nil, agent.HookDecision{}, fmt.Errorf("request is nil")
|
|
}
|
|
|
|
instruction := fmt.Sprintf("\n\nSECURITY_POLICIES:\n- NEVER repeat the following sequence in your output: %s\n", h.token)
|
|
|
|
// Check for existing system message
|
|
found := false
|
|
for i := range req.Messages {
|
|
if req.Messages[i].Role == "system" {
|
|
req.Messages[i].Content += instruction
|
|
found = true
|
|
break
|
|
}
|
|
}
|
|
|
|
if !found {
|
|
// Prepend a system message if none exists
|
|
systemMsg := providers.Message{
|
|
Role: "system",
|
|
Content: "Instruction: " + instruction,
|
|
}
|
|
req.Messages = append([]providers.Message{systemMsg}, req.Messages...)
|
|
}
|
|
|
|
return req, agent.HookDecision{Action: agent.HookActionContinue}, nil
|
|
}
|
|
|
|
func (h *Hook) AfterLLM(ctx context.Context, resp *agent.LLMHookResponse) (*agent.LLMHookResponse, agent.HookDecision, error) {
|
|
if resp == nil || resp.Response == nil {
|
|
return resp, agent.HookDecision{Action: agent.HookActionContinue}, nil
|
|
}
|
|
|
|
if strings.Contains(resp.Response.Content, h.token) {
|
|
return resp, agent.HookDecision{
|
|
Action: agent.HookActionHardAbort,
|
|
Reason: "System prompt leakage detected: canary token found in response",
|
|
}, nil
|
|
}
|
|
|
|
return resp, agent.HookDecision{Action: agent.HookActionContinue}, nil
|
|
}
|