piguard

package
v1.1.0 Latest Latest
Warning

This package is not in the latest version of its module.

Go to latest
Published: Jul 21, 2026 License: Apache-2.0 Imports: 23 Imported by: 0

Documentation

Overview

Package piguard ("prompt-injection guard") is the content-screening seam for inbound and outbound agent email. It answers one question — "is this content an attack?" — and is deliberately separate from the identity gates (internal/inboundpolicy), the review queue, and the audit log.

It is a stdlib-friendly leaf: it depends only on the standard library and the low-level emailauth package (for the parsed inbound auth verdict carried on a Request). Detection backends implement Detector; the Engine runs them in parallel and normalizes their heterogeneous outputs into a single Result. v1 ships exactly one detector — the dependency-free heuristics detector — but external providers (Lakera, Bedrock, Model Armor, Prompt Guard, …) plug in behind the same interface without reshaping the contract.

See docs/design/2026-06-20-agent-screening-hitl.md §4.2.

Index

Constants

View Source
const (
	CategoryInjectionDirect   = "prompt_injection_direct"   // OWASP LLM01 / ATLAS T0051.000
	CategoryInjectionIndirect = "prompt_injection_indirect" // ATLAS T0051.001 / NISTAML.015
	CategoryJailbreak         = "jailbreak"                 // ATLAS T0054
	CategoryExfiltration      = "data_exfiltration"
	CategoryObfuscation       = "obfuscation"
	CategorySensitive         = "sensitive_disclosure" // OWASP LLM02
)

Normalized category names, mapped to public taxonomies for portable audit: OWASP LLM Top 10, MITRE ATLAS, NIST AI 100-2e2025.

View Source
const DefaultScanCapBytes = 1 << 20 // 1 MiB of extracted text

DefaultScanCapBytes bounds how much decoded text the extractor will inspect. A multipart bomb or a huge attachment must never OOM the relay; content beyond the cap is dropped and DecodedSignals.Truncated is set so the caller can fail-to-review rather than treat "no finding" as safe. See design §5.

Variables

This section is empty.

Functions

func Extract

func Extract(raw []byte, capBytes int) ([]Segment, DecodedSignals, error)

Extract decodes a raw RFC-2822 message into normalized Segments and the cheap DecodedSignals, inspecting at most capBytes of decoded text (DefaultScanCapBytes when capBytes <= 0). It never returns an error for malformed MIME — it extracts what it can and reports nothing for the rest — because adversarial input is the norm here; a parse failure must not crash or skip screening. The only returned error is for a fundamentally unreadable header block.

func LooksTextual

func LooksTextual(s string) bool

LooksTextual reports whether a byte string is predominantly text (a low ratio of binary control bytes). Used to decide whether a part whose declared Content-Type is non-text actually carries scannable text — the declared type is attacker- controlled and not trusted.

Types

type Action

type Action string

Action is the screening decision a producer policy emits. Severity orders allow < flag < review < block.

const (
	ActionAllow  Action = "allow"
	ActionFlag   Action = "flag"
	ActionReview Action = "review"
	ActionBlock  Action = "block"
)

func ActionForScore

func ActionForScore(score, reviewThreshold, blockThreshold float64) Action

ActionForScore maps an aggregate score to an action band using the per-agent threshold ladder: below review = allow; [review, block) = review; ≥ block = block. Thresholds must satisfy 0 ≤ reviewThreshold ≤ blockThreshold ≤ 1; callers validate that upstream (ValidateScanConfig). This is the scan equivalent of SpamAssassin's score bands.

func MoreSevere

func MoreSevere(a, b Action) Action

MoreSevere returns whichever action is more severe (block > review > flag > allow). Used to combine a gate verdict and a scan verdict into one applied action.

type Aggregate

type Aggregate struct {
	Result
	// Degraded is true when fewer than MinOK detectors returned StatusOK. The caller
	// must map a degraded aggregate to review — NOT allow, NOT block.
	Degraded bool
	// MinAction is a deterministic force-override floor: when a high-confidence
	// deterministic marker is present (e.g. Unicode Tags-block smuggling, or
	// truncated/unscannable content), the applied action is at least this severe
	// regardless of score. ActionAllow when no override fires.
	MinAction Action
	// Truncated is true when DecodedSignals.Truncated (the content exceeded the
	// extraction-level scan cap) OR any StatusOK detector reported its own
	// Result.Truncated (e.g. Gemini's provider-side content-length cap cut off part
	// of the message before it was even sent) — either way, only partially
	// inspected. Folded into MinAction (→ review) but exposed for audit/metrics.
	Truncated bool
	// PerDetector retains every detector's raw Result (including timeouts/errors)
	// for audit and for writing protection_events rows.
	PerDetector []Result
}

Aggregate is the combined verdict plus the control signals the wiring layer needs to choose an action without knowing per-detector internals.

func (Aggregate) Action

func (a Aggregate) Action(reviewThreshold, blockThreshold float64) Action

Action folds the aggregate into a single applied action using the per-agent threshold ladder, guaranteeing the fail-safe defaults so call sites can't accidentally silent-allow: a degraded aggregate (too few usable verdicts) routes to review; the deterministic force-override floor (MinAction — truncated→review, Unicode-tags→flag) always applies; and a NaN score can never read as allow. Prefer this over calling ActionForScore(agg.Score, …) directly.

func (Aggregate) DetectorLabel

func (a Aggregate) DetectorLabel() string

DetectorLabel returns a stable, sorted, comma-joined list of the names of detectors that actually contributed a StatusOK verdict to this aggregate. Callers writing an audit row (e.g. protection_events.detector) should use this instead of hardcoding a single detector name — with more than one detector wired into the Engine, a hardcoded name silently mislabels which detector(s) actually drove the recorded score/action. Empty when no detector returned StatusOK (a Degraded aggregate); PerDetector still has every detector's Result, including timeouts/errors, for a caller that wants the full breakdown (e.g. marshaled into an audit row's raw/JSONB column).

type Category

type Category struct {
	Name       string  `json:"name"`
	NativeCode string  `json:"native_code,omitempty"`
	Score      float64 `json:"score,omitempty"`
}

Category is a normalized detection label. NativeCode preserves the provider's own code (e.g. Bedrock "PROMPT_ATTACK", Llama Guard "S1") so policy can route on a stable name while audit keeps the source.

type DecodedSignals

type DecodedSignals struct {
	// UnicodeTags reports U+E0000–U+E007F "ASCII smuggling" characters: ASCII
	// mirrored into an invisible Unicode block that LLM tokenizers still read.
	UnicodeTags bool
	// ZeroWidth reports zero-width/invisible separators (U+200B–U+200D, U+2060,
	// U+FEFF) used to hide or fragment instructions.
	ZeroWidth bool
	// HiddenCSSText reports text hidden from humans via CSS (display:none,
	// font-size:0, white-on-white, visibility:hidden, mso-hide:all, off-screen).
	HiddenCSSText bool
	// HomoglyphRatio is the fraction of letters that are non-ASCII confusables
	// (Cyrillic/Greek lookalikes) — a spoofing/obfuscation signal. 0..1.
	HomoglyphRatio float64
	// Unscannable reports a part whose bytes could not be scanned (a genuinely
	// binary attachment, no OCR in v1). Routed to review — no finding != benign.
	Unscannable bool
	// FragmentedURL reports reassembly-style obfuscation ("join 'h','ttp',…").
	FragmentedURL bool
	// PlainHTMLDiverge reports that the text/plain and visible text/html parts carry
	// materially different content (the human reads one, the agent may read the
	// other).
	PlainHTMLDiverge bool
	// Truncated reports that extraction stopped at the scan size cap — content beyond
	// the cap was not inspected, so "no finding" is not a safety guarantee.
	Truncated bool
}

DecodedSignals are the cheap, deterministic, near-zero-false-positive markers the extractor computes once. They drive both heuristic scoring and the Engine's force-overrides (e.g. Unicode Tags-block present ⇒ floor the action at flag).

type Detector

type Detector interface {
	// Inspect screens req and returns a normalized verdict. A returned error and a
	// non-OK Result.Status both signal "no usable verdict"; the Engine treats them
	// identically (excluded from the aggregate).
	//
	// Implementations MUST return promptly when ctx is cancelled or times out. The
	// Engine bounds each call with a timeout, but Go cannot force-cancel a goroutine:
	// a detector that ignores ctx and blocks leaks one goroutine per call. Network
	// backends must pass ctx into their request.
	Inspect(ctx context.Context, req Request) (*Result, error)
	// Name is the stable adapter id, e.g. "heuristics", "lakera".
	Name() string
}

Detector is the pluggable screening backend. Implementations must be safe for concurrent use (the Engine fans out across them) and must not panic on adversarial input — return a StatusError Result instead.

type Direction

type Direction int

Direction is the screening direction. Inbound screens a received message aimed at an agent; Outbound screens what an agent is about to send (exfiltration/leakage).

const (
	DirectionInput Direction = iota
	DirectionOutput
)

func (Direction) String

func (d Direction) String() string

type Engine

type Engine struct {
	// contains filtered or unexported fields
}

Engine runs the registered detectors in parallel and normalizes their verdicts into one Aggregate. It is safe for concurrent use.

func NewEngine

func NewEngine(cfg EngineConfig, detectors ...Detector) *Engine

NewEngine builds an Engine. At least one detector is expected; an Engine with no detectors always reports Degraded (fail-to-review), which is the safe posture.

func (*Engine) Evaluate

func (e *Engine) Evaluate(ctx context.Context, req Request) Aggregate

Evaluate fans out to all detectors with a per-detector timeout, then aggregates. It never panics: a detector that panics is recorded as StatusError and excluded.

func (*Engine) Timeout

func (e *Engine) Timeout() time.Duration

Timeout returns the effective per-detector timeout, e.g. for wiring-layer tests that assert a caller configured a non-default timeout (see buildScreenEngine in internal/relay/server.go, which widens this when the Gemini detector is present).

type EngineConfig

type EngineConfig struct {
	// Weights is the per-detector weight in the aggregate (by Detector.Name).
	// Missing → 1.0.
	Weights map[string]float64
	// Timeout bounds each detector's Inspect. A detector that exceeds it is recorded
	// as StatusTimeout and EXCLUDED from the aggregate (never counted as benign).
	Timeout time.Duration
	// MinOK is the minimum number of detectors that must return StatusOK for the
	// aggregate to be trusted. Below it, the Engine reports Degraded → the caller
	// maps to review (fail-to-review). Default 1.
	MinOK int
}

EngineConfig tunes aggregation. These are operator (deployment-level) knobs; the zero value is usable via sensible defaults applied in NewEngine.

type GeminiConfig

type GeminiConfig struct {
	// Model is the Gemini model name. When empty, NewGeminiDetector falls back to
	// the GEMINI_EVAL_MODEL environment variable, then to geminiDefaultModel.
	Model string
	// APIKey is the Google AI Studio key. When empty, NewGeminiDetector falls back
	// to the GEMINI_API_KEY and GOOGLE_API_KEY environment variables. Never log or
	// include this value in error messages.
	APIKey string
	// MaxRetries is the number of retries on transient API errors (429, 5xx).
	// Default geminiDefaultMaxRetries.
	MaxRetries int
	// HTTPClient allows injecting a custom *http.Client (e.g. for tests). When nil
	// a default client with a 30 s timeout is used.
	HTTPClient *http.Client
}

GeminiConfig configures the Gemini detector.

type GeminiDetector

type GeminiDetector struct {
	// contains filtered or unexported fields
}

GeminiDetector is a piguard.Detector backed by the Google Gemini API. It asks the model to classify inbound email for two threat types — prompt injection and phishing — and reports the more severe of the two as its primary Score/Flagged (so a high-confidence phishing verdict routes through review/block exactly like injection does, not just into the audit trail); both scores also remain visible individually as Categories for audit and per-threat-type reasoning. Safe for concurrent use.

The API key is sent only in the x-goog-api-key request header and is never written to logs or included in error messages.

func NewGeminiDetector

func NewGeminiDetector(cfg GeminiConfig) (*GeminiDetector, error)

NewGeminiDetector constructs a GeminiDetector. Returns a non-nil error when no API key is available (cfg.APIKey empty and neither GEMINI_API_KEY nor GOOGLE_API_KEY is set in the environment).

func (*GeminiDetector) Inspect

func (g *GeminiDetector) Inspect(ctx context.Context, req Request) (*Result, error)

Inspect implements Detector. It concatenates the email's extracted segments, sends them to Gemini, and maps the more severe of injection_confidence and phishing_confidence to the primary piguard Score/Flagged — so a purely-phishing email (no injection component) still crosses the Engine's review/block thresholds on its own, the same way a purely-injection email already does. Before this, phishing_confidence only ever reached a Category (audit-only, invisible to Aggregate.Action), so a 100%-confidence phishing verdict with zero injection signal produced an aggregate Score of 0 and was silently delivered. Both scores remain individually visible as Categories either way.

func (*GeminiDetector) Model

func (g *GeminiDetector) Model() string

Model returns the configured Gemini model name, e.g. for startup logging.

func (*GeminiDetector) Name

func (g *GeminiDetector) Name() string

Name implements Detector.

type HeuristicsDetector

type HeuristicsDetector struct{}

HeuristicsDetector is the built-in, dependency-free detector shipped in v1. It is deterministic, requires no network or model, and is tuned for near-zero false positives: it leans on the high-confidence DecodedSignals (which the extractor already computed) plus a small set of well-known injection / exfiltration content patterns. It is the reference Detector implementation and the baseline that external providers augment.

func NewHeuristicsDetector

func NewHeuristicsDetector() *HeuristicsDetector

NewHeuristicsDetector returns the built-in detector.

func (*HeuristicsDetector) Inspect

func (h *HeuristicsDetector) Inspect(_ context.Context, req Request) (*Result, error)

Inspect screens req using DecodedSignals and content patterns. It never returns a non-nil error (it cannot fail on adversarial input) and always reports StatusOK.

func (*HeuristicsDetector) Name

func (h *HeuristicsDetector) Name() string

type ProviderMeta

type ProviderMeta struct {
	Name          string         `json:"name"`
	ModelVersion  string         `json:"model_version,omitempty"`
	NativeVerdict string         `json:"native_verdict,omitempty"`
	NativeScore   string         `json:"native_score,omitempty"`
	LatencyMS     int            `json:"latency_ms,omitempty"`
	Raw           map[string]any `json:"raw,omitempty"`
}

ProviderMeta carries adapter identity and the raw provider response for forensics.

type Request

type Request struct {
	Direction Direction
	Segments  []Segment
	Signals   DecodedSignals
	// Sender is the claimed RFC 5322 From (inbound) or the agent identity
	// (outbound). Consult Auth.DMARC before treating the inbound domain as verified.
	Sender string
	// Auth is the parsed inbound auth verdict; nil on outbound.
	Auth *emailauth.Authentication
	// SizeBytes is the total extracted text size (post-cap).
	SizeBytes int
}

Request is the normalized input to every Detector. The caller extracts MIME once (via Extract) and passes Segments + Signals so each detector — raw-text classifier or structure-aware API — works off the same view without re-parsing.

type Result

type Result struct {
	Flagged    bool         `json:"flagged"`
	Score      float64      `json:"score"`
	Categories []Category   `json:"categories,omitempty"`
	Spans      []Span       `json:"spans,omitempty"`
	Status     Status       `json:"status"`
	Provider   ProviderMeta `json:"provider"`
	// Truncated reports that this detector only inspected a prefix of the content
	// (e.g. a provider-side context-window cap) — a "no finding" from THIS detector
	// is not a safety guarantee. Distinct from DecodedSignals.Truncated, which
	// reports extraction-level truncation (the 1 MiB scan cap) shared by all
	// detectors; this field is per-detector and additive to the Engine's
	// force-override floor (see Engine.forceFloor).
	Truncated bool `json:"truncated,omitempty"`
}

Result is the normalized verdict from a single detector. Flagged and Score are required (every adapter derives them, mapping enum/boolean providers onto a 0..1 score); Categories, Spans, and the richer ProviderMeta fields are optional.

type Segment

type Segment struct {
	Type    SegmentType
	Content string
	Ref     string
}

Segment is one extracted unit of content. Ref is a stable locator (e.g. "html#hidden[2]") used for span/offending-segment reporting.

type SegmentType

type SegmentType string

SegmentType identifies the provenance of a piece of extracted content. Splitting visible from hidden HTML, and per-part, lets a verdict point at *where* a payload hid and lets text-only detectors downcast while structure-aware providers see the full breakdown.

const (
	SegmentSubject        SegmentType = "subject"
	SegmentTextPlain      SegmentType = "text_plain"
	SegmentHTMLVisible    SegmentType = "html_visible"
	SegmentHTMLHidden     SegmentType = "html_hidden"
	SegmentAttachmentText SegmentType = "attachment_text"
	// SegmentImageOCR is reserved by the contract but never produced in v1 (no OCR).
	SegmentImageOCR SegmentType = "image_ocr"
)

type Span

type Span struct {
	Start int    `json:"start"`
	End   int    `json:"end"`
	Text  string `json:"text,omitempty"`
	Label string `json:"label,omitempty"`
	Ref   string `json:"ref,omitempty"`
}

Span is an optional flagged character range. Few providers populate it (essentially only Lakera), so it is always optional.

type Status

type Status int

Status reports whether a detector actually produced a verdict. A non-OK status is load-bearing: the Engine excludes it from the aggregate rather than counting it as benign, so an outage never silently allows.

const (
	StatusOK Status = iota
	StatusTimeout
	StatusError
	StatusUnsupported
)

func (Status) String

func (s Status) String() string

Jump to

Keyboard shortcuts

? : This menu
/ : Search site
f or F : Jump to
y or Y : Canonical URL