diff --git a/deepscan/event.go b/deepscan/event.go new file mode 100644 index 0000000..375133a --- /dev/null +++ b/deepscan/event.go @@ -0,0 +1,62 @@ +// Package deepscan is the shared wire + queue layer for AiKey's asynchronous +// deep scan (深扫 / 补扫). +// +// It has two wire generations living side by side on purpose: +// +// v1 — [1B version=1][4B LE len][{prompt, spans}]. Fire-and-forget: the sender +// never reads a response and the RECEIVER uploads its own findings. The +// sender is the detector (it owns the fast layer's findings, so its frames +// can carry `spans` with category/entity_type). Kept byte-for-byte so a +// new node still serves an old proxy's on-machine daemon. +// +// v2 — [1B version=2][4B LE len][FrameV2] out, [1B version=2][4B LE len][ResultFrame] +// back on the SAME connection, matched by job_id. The sender is the PROXY +// and the receiver returns a result instead of uploading, because the proxy +// is the only party that knows who the content belongs to. +// +// The package deliberately depends on nothing but the standard library: it is +// imported by aikey-proxy (Go), mirrored by ai-compliance-workers (Python), and +// wrapped by ai-compliance-detector. A dependency here would have to be vendored +// into all three. +package deepscan + +import "encoding/json" + +// ProtocolVersion is the v1 deep-scan wire version. The framing mirrors the +// apphook stdio protocol ([1B version][4B LE length][payload]). +const ProtocolVersion byte = 1 + +// FrameVersionV2 is the v2 task/result frame version (design §4b.1). +const FrameVersionV2 byte = 2 + +// Span is the minimal description of one span the sync (快层) layer already +// matched. +// +// 🔴 v2 frames do NOT carry Spans — see FrameV2. This type survives because the +// v1 frame still does, and v1 exists so a node can serve an old proxy's +// on-machine daemon. Do not reintroduce it into FrameV2 without reopening the +// decision recorded there. +// +// A Span carries OFFSETS + structural type only — never the matched substring. +type Span struct { + Start int `json:"start"` + End int `json:"end"` + Category string `json:"category"` + EntityType string `json:"entity_type"` + Detector string `json:"detector"` + Confidence int `json:"confidence"` +} + +// payloadV1 is the JSON body of one v1 deep-scan task: the raw prompt to +// re-scan plus the fast layer's hit list. Span offsets index Prompt. +type payloadV1 struct { + Prompt string `json:"prompt"` + Spans []Span `json:"spans"` +} + +// EncodePayloadV1 marshals one v1 task body (raw prompt + fast-layer spans). +// Exported so the detector's thin wrapper and the golden-frame fixture both +// produce bytes through this one function rather than two hand-written copies. +func EncodePayloadV1(prompt string, spans []Span) ([]byte, error) { + return json.Marshal(payloadV1{Prompt: prompt, Spans: spans}) +} diff --git a/deepscan/forwarder.go b/deepscan/forwarder.go new file mode 100644 index 0000000..adcc3ae --- /dev/null +++ b/deepscan/forwarder.go @@ -0,0 +1,186 @@ +package deepscan + +import ( + "context" + "fmt" + "os" + "sync" + "sync/atomic" + "time" +) + +// Config holds Forwarder settings. Zero values get sane defaults in +// NewForwarder. +type Config struct { + // QueueCapacity bounds the in-memory task queue. When full, Enqueue drops + // the task and counts it (never blocks). Default 256 — deliberately smaller + // than intake's queue because each task holds a full prompt, so worst-case + // memory ≈ QueueCapacity × average prompt size. + QueueCapacity int + + // SendTimeout bounds a single Sink.Send so a wedged endpoint can't stall + // the worker forever. Default 2s. (Even an unbounded stall would only fill + // the queue → drop, never block the hot path — this just bounds recovery.) + SendTimeout time.Duration +} + +// Forwarder is the background deep-scan forwarding goroutine. +// +// Lifecycle: NewForwarder → Start → Enqueue (many) → Close. +type Forwarder struct { + cfg Config + sink Sink + + queue chan task + stopCh chan struct{} + doneCh chan struct{} + + // Metrics — atomics, safe to read from any goroutine (e.g. /status). + enqueued atomic.Uint64 + dropped atomic.Uint64 + forwarded atomic.Uint64 + failed atomic.Uint64 + + closeOnce sync.Once +} + +// task is one queued deep-scan unit: the raw prompt + a snapshot of the fast +// layer's spans (already projected off engine structs by the caller, so we +// retain no Finding references and no Evidence). +type task struct { + prompt string + spans []Span +} + +// NewForwarder returns a configured (but not started) Forwarder that writes to +// sink. Call Start to begin the background loop. +func NewForwarder(cfg Config, sink Sink) *Forwarder { + if cfg.QueueCapacity <= 0 { + cfg.QueueCapacity = 256 + } + if cfg.SendTimeout <= 0 { + cfg.SendTimeout = 2 * time.Second + } + return &Forwarder{ + cfg: cfg, + sink: sink, + queue: make(chan task, cfg.QueueCapacity), + stopCh: make(chan struct{}), + doneCh: make(chan struct{}), + } +} + +// Start launches the background forwarding goroutine. Call exactly once. +func (f *Forwarder) Start() { + go f.loop() +} + +// Enqueue queues one deep-scan task. It NEVER blocks: if the queue is full +// (receiver slow / down), the task is dropped and counted. Safe to call on a nil +// *Forwarder (deep-scan disabled) — it's a no-op, so callers don't need to branch. +// +// 🔴 The signature takes []Span, not the detector's []types.Finding, because this +// package is imported by aikey-proxy and mirrored by ai-compliance-workers and +// must depend on nothing but the standard library. The detector keeps its +// findings→Span projection in its own thin wrapper (internal/deepscan), which is +// the only place that has ever needed it. +// +// spans MUST already be in the RAW prompt frame (the caller remaps normalized +// offsets first) so their offsets index prompt. +func (f *Forwarder) Enqueue(prompt string, spans []Span) { + if f == nil { + return + } + t := task{prompt: prompt, spans: spans} + select { + case f.queue <- t: + f.enqueued.Add(1) + default: + f.dropped.Add(1) + } +} + +// Close signals shutdown, drains/forwards what's queued (best-effort), and +// waits for the worker to exit — bounded by ctx so process shutdown isn't held +// up by a wedged endpoint. Safe to call on a nil *Forwarder. +func (f *Forwarder) Close(ctx context.Context) error { + if f == nil { + return nil + } + f.closeOnce.Do(func() { close(f.stopCh) }) + select { + case <-f.doneCh: + return nil + case <-ctx.Done(): + return ctx.Err() + } +} + +// Stats snapshots the forwarder counters for /status / debug / tests. +type Stats struct { + Enqueued uint64 + Dropped uint64 + Forwarded uint64 + Failed uint64 + QueueLen int +} + +func (f *Forwarder) Stats() Stats { + if f == nil { + return Stats{} + } + return Stats{ + Enqueued: f.enqueued.Load(), + Dropped: f.dropped.Load(), + Forwarded: f.forwarded.Load(), + Failed: f.failed.Load(), + QueueLen: len(f.queue), + } +} + +// loop is the background worker: one task at a time (deep-scan is best-effort, +// and single-flight keeps the Sink connection simple). It exits when stopCh +// closes, draining whatever is still queued first, then closes the sink. +func (f *Forwarder) loop() { + defer close(f.doneCh) + defer func() { _ = f.sink.Close() }() + + for { + select { + case <-f.stopCh: + drain: + for { + select { + case t := <-f.queue: + f.forward(t) + default: + break drain + } + } + return + case t := <-f.queue: + f.forward(t) + } + } +} + +// forward encodes one task and writes it to the sink, bounded by SendTimeout. +// Failures are counted + logged to stderr (matching the detector's logging +// convention) and swallowed — a dead deep-scan must never crash or stall the +// detector. +func (f *Forwarder) forward(t task) { + body, err := EncodePayloadV1(t.prompt, t.spans) + if err != nil { + f.failed.Add(1) + fmt.Fprintf(os.Stderr, "warn: deepscan encode failed: %v\n", err) + return + } + ctx, cancel := context.WithTimeout(context.Background(), f.cfg.SendTimeout) + defer cancel() + if err := f.sink.Send(ctx, encodeFrame(body)); err != nil { + f.failed.Add(1) + fmt.Fprintf(os.Stderr, "warn: deepscan forward failed: %v\n", err) + return + } + f.forwarded.Add(1) +} diff --git a/deepscan/forwarder_test.go b/deepscan/forwarder_test.go new file mode 100644 index 0000000..023b58e --- /dev/null +++ b/deepscan/forwarder_test.go @@ -0,0 +1,148 @@ +package deepscan + +import ( + "context" + "errors" + "sync" + "sync/atomic" + "testing" + "time" +) + +// blockingSink blocks each Send until ctx fires — simulates a slow/wedged +// deep-scan endpoint. Used to prove the hot path (Enqueue) never blocks on it. +type blockingSink struct{ sends atomic.Int64 } + +func (s *blockingSink) Send(ctx context.Context, _ []byte) error { + s.sends.Add(1) + <-ctx.Done() + return ctx.Err() +} +func (s *blockingSink) Close() error { return nil } + +// errorSink fails every Send immediately — simulates a deep-scan that's up but +// rejecting. Proves the worker keeps draining and never deadlocks. +type errorSink struct{} + +func (errorSink) Send(context.Context, []byte) error { return errors.New("sink down") } +func (errorSink) Close() error { return nil } + +// countingSink accepts every Send. Proves drain-on-close forwards everything. +type countingSink struct{ count atomic.Int64 } + +func (s *countingSink) Send(context.Context, []byte) error { s.count.Add(1); return nil } +func (s *countingSink) Close() error { return nil } + +// Enqueue must drop (not block, not grow) once the bounded queue is full. With +// no worker started, the queue fills to capacity and every extra Enqueue drops. +func TestForwarder_EnqueueDropsWhenFull(t *testing.T) { + const cap, extra = 4, 6 + f := NewForwarder(Config{QueueCapacity: cap}, &countingSink{}) + for i := 0; i < cap+extra; i++ { + f.Enqueue("prompt", nil) + } + st := f.Stats() + if st.Enqueued != cap { + t.Errorf("enqueued = %d, want %d (capacity)", st.Enqueued, cap) + } + if st.Dropped != extra { + t.Errorf("dropped = %d, want %d", st.Dropped, extra) + } +} + +// The I3 core assertion: even with the worker stuck on a blocking Send, every +// Enqueue returns ~instantly. Enqueue is a pure channel op — it must be fully +// decoupled from sink latency. +func TestForwarder_EnqueueNeverBlocksWhileSinkStuck(t *testing.T) { + sink := &blockingSink{} + f := NewForwarder(Config{QueueCapacity: 8, SendTimeout: 500 * time.Millisecond}, sink) + f.Start() + t.Cleanup(func() { + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + _ = f.Close(ctx) + }) + + for i := 0; i < 100; i++ { + start := time.Now() + f.Enqueue("prompt", nil) + if d := time.Since(start); d > 20*time.Millisecond { + t.Fatalf("Enqueue blocked %v on iter %d (must be ~instant regardless of sink)", d, i) + } + } + if st := f.Stats(); st.Dropped == 0 { + t.Errorf("expected drops with a stuck sink, got %+v", st) + } +} + +// A sink that errors on every Send must not deadlock the worker; Close drains +// and returns promptly, and every task is counted as failed. +func TestForwarder_SinkErrorNoDeadlock(t *testing.T) { + const n = 10 + f := NewForwarder(Config{QueueCapacity: 16}, errorSink{}) + f.Start() + for i := 0; i < n; i++ { + f.Enqueue("prompt", nil) + } + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + if err := f.Close(ctx); err != nil { + t.Fatalf("Close = %v (worker deadlocked on erroring sink?)", err) + } + if st := f.Stats(); st.Failed != n { + t.Errorf("failed = %d, want %d", st.Failed, n) + } +} + +// Close drains the queue: every enqueued task reaches the sink, none lost. +func TestForwarder_CloseDrainsAll(t *testing.T) { + const n = 20 + sink := &countingSink{} + f := NewForwarder(Config{QueueCapacity: 64}, sink) + f.Start() + for i := 0; i < n; i++ { + f.Enqueue("prompt", nil) + } + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + if err := f.Close(ctx); err != nil { + t.Fatalf("Close = %v", err) + } + if st := f.Stats(); st.Forwarded != n { + t.Errorf("forwarded = %d, want %d (drain lost tasks?)", st.Forwarded, n) + } + if got := sink.count.Load(); got != int64(n) { + t.Errorf("sink received %d, want %d", got, n) + } +} + +// A nil *Forwarder (deep-scan disabled) must be safe for every method so +// callers never branch. +func TestForwarder_NilSafe(t *testing.T) { + var f *Forwarder + f.Enqueue("prompt", nil) // must not panic + if err := f.Close(context.Background()); err != nil { + t.Errorf("nil Close = %v, want nil", err) + } + if st := f.Stats(); st != (Stats{}) { + t.Errorf("nil Stats = %+v, want zero", st) + } +} + +// Close is idempotent / ctx-bounded: a second Close (or one racing the worker) +// must not panic on a double channel close. +func TestForwarder_CloseIdempotent(t *testing.T) { + f := NewForwarder(Config{}, &countingSink{}) + f.Start() + var wg sync.WaitGroup + for i := 0; i < 3; i++ { + wg.Add(1) + go func() { + defer wg.Done() + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + _ = f.Close(ctx) + }() + } + wg.Wait() +} diff --git a/deepscan/frame_v2.go b/deepscan/frame_v2.go new file mode 100644 index 0000000..f34e0cf --- /dev/null +++ b/deepscan/frame_v2.go @@ -0,0 +1,102 @@ +package deepscan + +import ( + "encoding/json" + "fmt" +) + +// Source values for FrameV2.Source (design §3.1). +const ( + SourceRequest = "request" + SourceResponse = "response" +) + +// Engine names carried in FrameV2.Engines and Finding.Engine. +const ( + EngineBGE = "bge" + EngineRules = "rules" +) + +// Range is a half-open [Start, End) byte interval into the piece. +type Range struct { + Start int `json:"start"` + End int `json:"end"` +} + +// FrameV2 is one asynchronous scan task: proxy → scan node over TLS, or proxy → +// on-machine daemon over a unix socket (same shape, no token on the socket). +// +// 🔴 WHAT IS DELIBERATELY ABSENT, AND WHY (design §3.1, R-scan-node-deepscan-3.S1): +// +// - seat_id / virtual_key_id / session_id / trace_id / user_id. A node is a +// SHARED box that several employees' content passes through, and it uploads +// nothing — the proxy stamps identity onto the event after the result comes +// back. Sending identity would tell a compromised node WHO said WHAT and buy +// nothing. TenantID is the one exception: without it the node cannot refuse +// content belonging to another org. +// +// - spans (the fast layer's hits). v1 frames DO carry them, because in v1 the +// sender is the detector, which owns the findings. In v1.1 the sender is the +// proxy, and apphook.Response / pkg/pipewire.Response hand the proxy OFFSETS +// ONLY — by design, not by omission: "the proxy never learns what the token +// stands for business-wise" (aikey-proxy internal/apphook/apphook.go +// invariant #16), re-verified in baseline-forensics §F1 on 2026-09-11. +// User decision 2026-09-11: the node re-scans [0, HeadBytes) itself to get +// the equivalent hit list for bge position-dedup, instead of the proxy +// learning categories it is forbidden to know. +// Fence: TestFrameV2_NoSeatOrSessionFields. +type FrameV2 struct { + Version int `json:"version"` + JobID string `json:"job_id"` + // Token is the `sct1` org-scoped scan token. Empty on the unix-socket form + // (same machine, same trust domain). + Token string `json:"token,omitempty"` + TenantID string `json:"tenant_id"` + AuditUnitID string `json:"audit_unit_id"` + ContentSHA256 string `json:"content_sha256"` + Source string `json:"source"` + // HeadBytes is how many bytes the synchronous fast layer already inspected. + // Rule findings that end at or before it are the fast layer's, not new + // coverage; 0 when the fast layer degraded and inspected nothing. + HeadBytes int `json:"head_bytes"` + // Engines is the subset of {bge, rules} the receiver should run. + Engines []string `json:"engines"` + // RuleChunks are computed ONCE by the sender (pkg/scanchunk) and carried on + // the wire so Go and Python never hold two chunking implementations that can + // drift apart. + RuleChunks []Range `json:"rule_chunks,omitempty"` + Prompt string `json:"prompt"` +} + +// EncodeFrameV2 marshals a v2 task and wraps it in the shared frame header. +func EncodeFrameV2(f FrameV2) ([]byte, error) { + if f.Version == 0 { + f.Version = int(FrameVersionV2) + } + body, err := json.Marshal(f) + if err != nil { + return nil, fmt.Errorf("deepscan: encode frame v2: %w", err) + } + if len(body)+frameHeaderLen > MaxFrameBytes { + return nil, fmt.Errorf("deepscan: frame v2 is %d bytes, over the %d-byte cap", len(body)+frameHeaderLen, MaxFrameBytes) + } + return EncodeFrame(FrameVersionV2, body), nil +} + +// DecodeFrameV2 parses a framed v2 task. It refuses any other version rather +// than guessing: a v1 body decoded as v2 would silently produce a frame with an +// empty tenant, which the node would then refuse for the wrong reason. +func DecodeFrameV2(frame []byte) (FrameV2, error) { + version, body, err := SplitFrame(frame) + if err != nil { + return FrameV2{}, err + } + if version != FrameVersionV2 { + return FrameV2{}, fmt.Errorf("deepscan: frame version %d is not v2", version) + } + var f FrameV2 + if err := json.Unmarshal(body, &f); err != nil { + return FrameV2{}, fmt.Errorf("deepscan: decode frame v2: %w", err) + } + return f, nil +} diff --git a/deepscan/frame_v2_test.go b/deepscan/frame_v2_test.go new file mode 100644 index 0000000..789d479 --- /dev/null +++ b/deepscan/frame_v2_test.go @@ -0,0 +1,145 @@ +package deepscan + +import ( + "encoding/json" + "strings" + "testing" +) + +// TestFrameV2_NoSeatOrSessionFields is the wire-privacy fence for the v2 task +// frame. +// +// WHY this fence exists (design §3.1, R-scan-node-deepscan-3.S1): a scan node is +// a SHARED box that several employees' content flows through, and it does not +// upload anything itself — the proxy stamps identity onto the event afterwards. +// So the node needs exactly one identity field, `tenant_id`, to refuse content +// from another org. Every other identity field (seat, virtual key, session, +// trace, user) would tell a compromised node WHO said WHAT while buying nothing. +// +// It also asserts the frame carries NO `spans`. v1 frames did carry the fast +// layer's hits WITH their category, because in v1 the sender was the DETECTOR, +// which owns the findings. In v1.1 the sender is the PROXY, and the proxy is +// forbidden from learning what the detector matched (apphook.go invariant #16: +// "the proxy never learns what the token stands for business-wise"). Measured in +// baseline-forensics §F1: apphook.Response and pkg/pipewire.Response hand the +// proxy offsets only. User decision 2026-09-11: the node re-scans [0, head_bytes) +// itself instead. A `spans` field reappearing here means someone reintroduced the +// proxy→business-semantics coupling that decision removed. +func TestFrameV2_NoSeatOrSessionFields(t *testing.T) { + f := FrameV2{ + Version: int(FrameVersionV2), + JobID: "job-1", + Token: "sct1.org_a.1757620000.k1.abc", + TenantID: "org_a", + AuditUnitID: "au_deadbeef", + ContentSHA256: "b1946ac92492d2347c6235b4d2611184", + Source: SourceRequest, + HeadBytes: 16384, + Engines: []string{EngineBGE, EngineRules}, + RuleChunks: []Range{{Start: 14336, End: 30720}}, + Prompt: "hello", + } + b, err := EncodeFrameV2(f) + if err != nil { + t.Fatalf("EncodeFrameV2: %v", err) + } + if len(b) < frameHeaderLen || b[0] != FrameVersionV2 { + t.Fatalf("frame header: want version byte %d, got % x", FrameVersionV2, b[:min(len(b), 8)]) + } + body := string(b[frameHeaderLen:]) + + for _, banned := range []string{"seat_id", "virtual_key_id", "session_id", "trace_id", "user_id", "spans"} { + if strings.Contains(body, banned) { + t.Errorf("v2 frame carries %q — it must not; body=%s", banned, body) + } + } + + var got map[string]any + if err := json.Unmarshal(b[frameHeaderLen:], &got); err != nil { + t.Fatalf("frame body is not JSON: %v", err) + } + if got["tenant_id"] != "org_a" { + t.Errorf("tenant_id: want org_a, got %v — the node cannot refuse cross-org content without it", got["tenant_id"]) + } + if got["audit_unit_id"] != "au_deadbeef" { + t.Errorf("audit_unit_id: want au_deadbeef, got %v", got["audit_unit_id"]) + } +} + +// TestResultFrame_RoundTrip pins the result frame both ways. The node answers on +// the SAME connection and results may come back out of order, so job_id has to +// survive the trip; findings carry ABSOLUTE offsets into the piece (design §3.6) +// because the proxy merges them against head_bytes without re-deriving anything. +func TestResultFrame_RoundTrip(t *testing.T) { + in := ResultFrame{ + JobID: "job-1", + Status: StatusPartial, + Reason: ReasonWindowCap, + ScannedBytes: 240000, + TotalBytes: 262144, + Engines: EngineResults{ + BGE: EngineResult{Status: StatusPartial, Windows: 4096, Reason: ReasonWindowCap}, + Rules: EngineResult{Status: StatusComplete, Chunks: 16, DetectorVersion: "95eebb1", ContentVersion: "cv-7"}, + }, + Findings: []Finding{{ + Engine: EngineRules, RuleID: "secret.connection-uri-password", + Category: "secret", EntityType: "CREDENTIAL_DSN", Severity: "high", + Confidence: 95, Start: 250000, End: 250072, Detector: "regex", + }}, + RuleVerdicts: []RangeVerdict{{Start: 245760, End: 262144, Action: "block"}}, + } + b, err := EncodeResult(in) + if err != nil { + t.Fatalf("EncodeResult: %v", err) + } + out, err := DecodeResult(b) + if err != nil { + t.Fatalf("DecodeResult: %v", err) + } + if out.JobID != in.JobID || out.Status != in.Status || out.Reason != in.Reason { + t.Errorf("header round-trip: got %+v want %+v", out, in) + } + if out.ScannedBytes != in.ScannedBytes || out.TotalBytes != in.TotalBytes { + t.Errorf("coverage round-trip: got %d/%d want %d/%d", + out.ScannedBytes, out.TotalBytes, in.ScannedBytes, in.TotalBytes) + } + if len(out.Findings) != 1 || out.Findings[0].Start != 250000 || out.Findings[0].End != 250072 { + t.Fatalf("finding offsets did not survive: %+v", out.Findings) + } + if out.Findings[0].EntityType != "CREDENTIAL_DSN" || out.Findings[0].RuleID != "secret.connection-uri-password" { + t.Errorf("finding identity did not survive: %+v", out.Findings[0]) + } + if len(out.RuleVerdicts) != 1 || out.RuleVerdicts[0].Action != "block" { + t.Fatalf("rule verdicts did not survive: %+v", out.RuleVerdicts) + } + if out.Engines.Rules.DetectorVersion != "95eebb1" || out.Engines.BGE.Windows != 4096 { + t.Errorf("engine block did not survive: %+v", out.Engines) + } +} + +// TestResultFrame_RejectIsSelfContained: a refused frame answers with job_id + +// reject ONLY. A node that refused a frame has not scanned it, so any coverage +// number it reported would be a lie the proxy would then store as `complete`. +func TestResultFrame_RejectIsSelfContained(t *testing.T) { + b, err := EncodeResult(ResultFrame{JobID: "job-9", Reject: RejectTenantMismatch}) + if err != nil { + t.Fatalf("EncodeResult: %v", err) + } + out, err := DecodeResult(b) + if err != nil { + t.Fatalf("DecodeResult: %v", err) + } + if out.Reject != RejectTenantMismatch { + t.Fatalf("reject code lost: %+v", out) + } + if out.Status != "" || out.ScannedBytes != 0 || len(out.Findings) != 0 { + t.Errorf("a rejected frame must carry no scan result, got %+v", out) + } +} + +func min(a, b int) int { + if a < b { + return a + } + return b +} diff --git a/deepscan/gen_golden_test.go b/deepscan/gen_golden_test.go new file mode 100644 index 0000000..c89f4e2 --- /dev/null +++ b/deepscan/gen_golden_test.go @@ -0,0 +1,148 @@ +package deepscan + +import ( + "encoding/json" + "os" + "path/filepath" + "testing" +) + +// TestGoldenFrames_MatchFixture pins the exact BYTES of the v1 and v2 frames in +// testdata/golden_frames.json. +// +// WHY a byte fixture and not just a Go round-trip: ai-compliance-workers decodes +// these frames in PYTHON. A Go round-trip proves Go agrees with itself, which is +// precisely the failure this fixture exists to catch — the two languages drifting +// apart on field names, the 5-byte header, or little-endian length. The Python +// side reads the SAME file (tests/testdata/golden_frames.json) and must produce +// byte-identical output. Regenerate deliberately with -update, never casually: +// changing a byte here is a wire-compat decision. +func TestGoldenFrames_MatchFixture(t *testing.T) { + path := filepath.Join("testdata", "golden_frames.json") + got := buildGoldenFrames(t) + + if os.Getenv("UPDATE_GOLDEN") == "1" { + b, err := json.MarshalIndent(got, "", " ") + if err != nil { + t.Fatalf("marshal golden: %v", err) + } + if err := os.MkdirAll("testdata", 0o755); err != nil { + t.Fatalf("mkdir testdata: %v", err) + } + if err := os.WriteFile(path, append(b, '\n'), 0o644); err != nil { + t.Fatalf("write golden: %v", err) + } + t.Logf("golden fixture regenerated at %s", path) + return + } + + want, err := os.ReadFile(path) + if err != nil { + t.Fatalf("read golden fixture (regenerate with UPDATE_GOLDEN=1): %v", err) + } + var wantMap map[string]any + if err := json.Unmarshal(want, &wantMap); err != nil { + t.Fatalf("golden fixture is not JSON: %v", err) + } + gotJSON, _ := json.Marshal(got) + var gotMap map[string]any + _ = json.Unmarshal(gotJSON, &gotMap) + + for _, k := range []string{"v1_task", "v2_task", "v2_result", "v2_result_reject"} { + if gotMap[k] == nil { + t.Fatalf("generator produced no %q case", k) + } + if wantMap[k] == nil { + t.Fatalf("fixture has no %q case — regenerate with UPDATE_GOLDEN=1", k) + } + g, _ := json.Marshal(gotMap[k]) + w, _ := json.Marshal(wantMap[k]) + if string(g) != string(w) { + t.Errorf("%s drifted from the fixture the Python mirror reads.\n got: %s\nwant: %s", k, g, w) + } + } +} + +type goldenCase struct { + Description string `json:"description"` + FrameHex string `json:"frame_hex"` + BodyJSON string `json:"body_json"` + Version int `json:"version"` +} + +func buildGoldenFrames(t *testing.T) map[string]goldenCase { + t.Helper() + out := map[string]goldenCase{} + + v1Body, err := EncodePayloadV1("客户手机号 13800138000 请核对", []Span{ + {Start: 6, End: 17, Category: "pii", EntityType: "CN_PHONE", Detector: "regex", Confidence: 95}, + }) + if err != nil { + t.Fatalf("v1 body: %v", err) + } + out["v1_task"] = mkCase(t, "v1 task frame: fire-and-forget, carries the fast layer's spans because the DETECTOR sends it", + EncodeFrame(ProtocolVersion, v1Body), int(ProtocolVersion)) + + v2, err := EncodeFrameV2(FrameV2{ + JobID: "job-golden-1", + Token: "sct1.org_a.1757620000.k1.Zm9vYmFy", + TenantID: "org_a", + AuditUnitID: "au_0123456789abcdef", + ContentSHA256: "b1946ac92492d2347c6235b4d2611184", + Source: SourceRequest, + HeadBytes: 16384, + Engines: []string{EngineBGE, EngineRules}, + RuleChunks: []Range{{Start: 14336, End: 30720}, {Start: 28672, End: 45056}}, + Prompt: "客户手机号 13800138000 请核对", + }) + if err != nil { + t.Fatalf("v2 frame: %v", err) + } + out["v2_task"] = mkCase(t, "v2 task frame: NO spans and NO seat/session identity — the node re-scans the head itself", v2, int(FrameVersionV2)) + + res, err := EncodeResult(ResultFrame{ + JobID: "job-golden-1", Status: StatusPartial, Reason: ReasonWindowCap, + ScannedBytes: 240000, TotalBytes: 262144, + Engines: EngineResults{ + BGE: EngineResult{Status: StatusPartial, Windows: 4096, Reason: ReasonWindowCap}, + Rules: EngineResult{Status: StatusComplete, Chunks: 2, DetectorVersion: "95eebb1", ContentVersion: "cv-7"}, + }, + Findings: []Finding{{ + Engine: EngineRules, RuleID: "secret.connection-uri-password", Category: "secret", + EntityType: "CREDENTIAL_DSN", Severity: "high", Confidence: 95, + Start: 250000, End: 250072, Detector: "regex", + }}, + RuleVerdicts: []RangeVerdict{{Start: 245760, End: 262144, Action: "block"}}, + }) + if err != nil { + t.Fatalf("v2 result: %v", err) + } + out["v2_result"] = mkCase(t, "v2 result frame: absolute offsets, no raw matched text", res, int(FrameVersionV2)) + + rej, err := EncodeResult(ResultFrame{JobID: "job-golden-1", Reject: RejectTenantMismatch}) + if err != nil { + t.Fatalf("v2 reject: %v", err) + } + out["v2_result_reject"] = mkCase(t, "v2 reject: job_id + reject and nothing else", rej, int(FrameVersionV2)) + return out +} + +func mkCase(t *testing.T, desc string, frame []byte, version int) goldenCase { + t.Helper() + _, body, err := SplitFrame(frame) + if err != nil { + t.Fatalf("%s: SplitFrame: %v", desc, err) + } + return goldenCase{Description: desc, FrameHex: hexOf(frame), BodyJSON: string(body), Version: version} +} + +const hexDigits = "0123456789abcdef" + +func hexOf(b []byte) string { + out := make([]byte, len(b)*2) + for i, c := range b { + out[i*2] = hexDigits[c>>4] + out[i*2+1] = hexDigits[c&0x0f] + } + return string(out) +} diff --git a/deepscan/go.mod b/deepscan/go.mod new file mode 100644 index 0000000..31291c1 --- /dev/null +++ b/deepscan/go.mod @@ -0,0 +1,3 @@ +module github.com/AiKeyLabs/pkg/deepscan + +go 1.26.1 diff --git a/deepscan/result.go b/deepscan/result.go new file mode 100644 index 0000000..77eb438 --- /dev/null +++ b/deepscan/result.go @@ -0,0 +1,162 @@ +package deepscan + +import ( + "encoding/binary" + "encoding/json" + "fmt" + "io" +) + +// Coverage status values (design §4b.12). `partial` is not a soft `complete`: +// it is the difference between "we looked at all of it and found nothing" and +// "we ran out of budget", which is the whole reason scan_coverage exists. +const ( + StatusComplete = "complete" + StatusPartial = "partial" + StatusFailed = "failed" +) + +// Reasons a result is partial. +const ( + ReasonPieceCap = "piece_cap" // the piece was truncated before it was sent + ReasonWindowCap = "window_cap" // bge hit AIKEY_DEEPSCAN_MAX_WINDOWS + ReasonTransient = "transient" // finalised after repeated transient failures +) + +// Reject codes (design §4b.1). The node answers within 50ms with {job_id, +// reject} and nothing else. +// +// 🔴 These constants are mirrored in ai-compliance-workers wire.py and pinned by +// a cross-language fixture test. Same name, same value, or a node silently +// refuses frames for a reason the proxy cannot act on. +const ( + RejectBusy = "busy" // admission full — try the next node + RejectUnauthorized = "unauthorized" // token invalid/expired — try the next node + RejectTenantMismatch = "tenant_mismatch" // wrong org — do NOT try another node + RejectBadFrame = "bad_frame" // malformed — do NOT try another node + RejectVersionUnsupported = "version_unsupported" // wire skew — do NOT try another node +) + +// EngineResult is one engine's outcome inside a result frame. +type EngineResult struct { + Status string `json:"status"` + Reason string `json:"reason,omitempty"` + // Windows is how many bge windows were encoded (bge only). + Windows int `json:"windows,omitempty"` + // Chunks is how many rule chunks were scanned (rules only). + Chunks int `json:"chunks,omitempty"` + // DetectorVersion / ContentVersion let the proxy count version skew between + // what the node ran and what it runs itself, instead of silently merging + // results produced by a different ruleset. + DetectorVersion string `json:"detector_version,omitempty"` + ContentVersion string `json:"content_version,omitempty"` +} + +// EngineResults groups the per-engine outcomes. +type EngineResults struct { + BGE EngineResult `json:"bge,omitempty"` + Rules EngineResult `json:"rules,omitempty"` +} + +// Finding is one hit the asynchronous scan produced. +// +// 🔴 Offsets are ABSOLUTE into the piece the proxy sent, never relative to a +// chunk or window. The receiver adds the chunk/window origin before it answers, +// because the proxy merges findings against HeadBytes and cannot re-derive an +// origin it never saw. Fence: TestResultFrame_RoundTrip + the workers-side +// test_tail_chunks_absolute_offsets. +// +// 🚫 No raw matched text. The result travels back over the network and lands in +// an audit event; the evidence stays on the box that produced it. +type Finding struct { + Engine string `json:"engine"` + RuleID string `json:"rule_id,omitempty"` + Category string `json:"category"` + EntityType string `json:"entity_type"` + Severity string `json:"severity"` + Confidence int `json:"confidence"` + Start int `json:"start"` + End int `json:"end"` + Detector string `json:"detector,omitempty"` +} + +// RangeVerdict is one rule chunk's raw action, kept separate from Findings so +// the proxy can decide "this WOULD have been blocked" against its own ceilings +// rather than trusting a verdict computed on a box with a different policy. +type RangeVerdict struct { + Start int `json:"start"` + End int `json:"end"` + Action string `json:"action"` +} + +// ResultFrame is what a node (or the on-machine daemon, for v2 frames) answers +// on the same connection. Results may come back out of order; JobID matches +// them up. +type ResultFrame struct { + JobID string `json:"job_id"` + // Reject, when set, is the ONLY other field present: a refused frame was + // never scanned, so any coverage number alongside it would be a lie the + // proxy would go on to store as a real scan result. + Reject string `json:"reject,omitempty"` + + Status string `json:"status,omitempty"` + Reason string `json:"reason,omitempty"` + ScannedBytes int `json:"scanned_bytes,omitempty"` + TotalBytes int `json:"total_bytes,omitempty"` + Engines EngineResults `json:"engines,omitempty"` + Findings []Finding `json:"findings,omitempty"` + RuleVerdicts []RangeVerdict `json:"rule_verdicts,omitempty"` +} + +// EncodeResult marshals a result and wraps it in the shared frame header. +func EncodeResult(r ResultFrame) ([]byte, error) { + body, err := json.Marshal(r) + if err != nil { + return nil, fmt.Errorf("deepscan: encode result: %w", err) + } + return EncodeFrame(FrameVersionV2, body), nil +} + +// DecodeResult parses a framed result. +func DecodeResult(frame []byte) (ResultFrame, error) { + version, body, err := SplitFrame(frame) + if err != nil { + return ResultFrame{}, err + } + if version != FrameVersionV2 { + return ResultFrame{}, fmt.Errorf("deepscan: result frame version %d is not v2", version) + } + var r ResultFrame + if err := json.Unmarshal(body, &r); err != nil { + return ResultFrame{}, fmt.Errorf("deepscan: decode result: %w", err) + } + return r, nil +} + +// ReadResult reads exactly one framed result — [1B version][4B LE length][JSON] +// — from r. The declared length is checked against MaxFrameBytes BEFORE any +// allocation, so a hostile or corrupt peer cannot make the reader allocate +// whatever it claims. Every sink that reads a node's answer uses this, so the +// framing cannot drift between the TLS and the local-socket paths. +func ReadResult(r io.Reader) (ResultFrame, error) { + var head [5]byte + if _, err := io.ReadFull(r, head[:]); err != nil { + return ResultFrame{}, fmt.Errorf("deepscan: read result header: %w", err) + } + n := binary.LittleEndian.Uint32(head[1:5]) + if n > MaxFrameBytes { + return ResultFrame{}, fmt.Errorf("deepscan: result declares %d bytes, over the %d cap", n, MaxFrameBytes) + } + body := make([]byte, n) + if _, err := io.ReadFull(r, body); err != nil { + return ResultFrame{}, fmt.Errorf("deepscan: read result body: %w", err) + } + return DecodeResult(append(head[:], body...)) +} + +// RejectError is a node's refusal of one frame, carried as an error so a +// delivery loop can tell "this node said no, and why" from "this node is +// unreachable". Code is one of the Reject* constants. +type RejectError struct{ Code string } + +func (e *RejectError) Error() string { return "deepscan: node rejected the frame: " + e.Code } diff --git a/deepscan/sink.go b/deepscan/sink.go new file mode 100644 index 0000000..1c572e3 --- /dev/null +++ b/deepscan/sink.go @@ -0,0 +1,126 @@ +package deepscan + +import ( + "context" + "encoding/binary" + "fmt" + "net" + "sync" + "time" +) + +// Sink is the transport the forwarder writes encoded frames to. The seam keeps +// the forwarder's bounded-queue logic independent of the wire: tests inject a +// slow/dead/erroring sink, and the scan-node TLS sink (pkg/scannode) plugs in +// here — both without touching the forwarder (invariant I8: the endpoint +// changes, the protocol and forwarder code do not). +type Sink interface { + // Send writes one already-framed message. It must honor ctx for + // cancellation/deadline. Returning an error makes the forwarder count a + // failure and move on; either way the detect hot path is never affected + // (the bounded queue + non-blocking Enqueue absorb a slow/dead Sink). + Send(ctx context.Context, frame []byte) error + // Close releases any underlying connection. + Close() error +} + +// frameHeaderLen is the fixed [1B version][4B LE length] header, shared by v1 +// and v2 so one parser shape reads both. +const frameHeaderLen = 5 + +// MaxFrameBytes bounds a single frame (design §4b.1: "整帧 ≤1MiB"). A reader +// must refuse anything larger rather than allocate it: the length prefix is +// attacker-influenced on the TLS path. +const MaxFrameBytes = 1 << 20 + +// EncodeFrame wraps a body as [1B version][4B LE length][body]. +func EncodeFrame(version byte, body []byte) []byte { + frame := make([]byte, frameHeaderLen+len(body)) + frame[0] = version + binary.LittleEndian.PutUint32(frame[1:frameHeaderLen], uint32(len(body))) + copy(frame[frameHeaderLen:], body) + return frame +} + +// encodeFrame keeps the v1 call sites byte-identical after the migration. +func encodeFrame(body []byte) []byte { return EncodeFrame(ProtocolVersion, body) } + +// SplitFrame validates a framed message and returns its version and body. +// Length-prefix first, allocation never: a declared length over MaxFrameBytes is +// rejected without reading it. +func SplitFrame(frame []byte) (version byte, body []byte, err error) { + if len(frame) < frameHeaderLen { + return 0, nil, fmt.Errorf("deepscan frame: %d bytes is shorter than the %d-byte header", len(frame), frameHeaderLen) + } + n := binary.LittleEndian.Uint32(frame[1:frameHeaderLen]) + if n > MaxFrameBytes { + return 0, nil, fmt.Errorf("deepscan frame: declared body %d bytes exceeds the %d-byte cap", n, MaxFrameBytes) + } + if int(n) != len(frame)-frameHeaderLen { + return 0, nil, fmt.Errorf("deepscan frame: header says %d body bytes, got %d", n, len(frame)-frameHeaderLen) + } + return frame[0], frame[frameHeaderLen:], nil +} + +// unixSink writes frames to a unix domain socket. It dials lazily and redials +// on the next Send after a write/dial error — a dead deep-scan endpoint +// degrades to dropped 补漏 tasks, never a stuck hot path (the forwarder's +// bounded queue takes up the slack and drops on overflow). +type unixSink struct { + path string + dialTimeout time.Duration + writeTimeout time.Duration + + mu sync.Mutex + conn net.Conn +} + +// NewUnixSink returns a Sink connecting to the unix socket at path. Dialing is +// lazy (first Send) so the sender starts even if the deep-scan process isn't up +// yet — it just drops/redials until the socket appears. +func NewUnixSink(path string, dialTimeout, writeTimeout time.Duration) Sink { + if dialTimeout <= 0 { + dialTimeout = 1 * time.Second + } + if writeTimeout <= 0 { + writeTimeout = 1 * time.Second + } + return &unixSink{path: path, dialTimeout: dialTimeout, writeTimeout: writeTimeout} +} + +func (s *unixSink) Send(ctx context.Context, frame []byte) error { + s.mu.Lock() + defer s.mu.Unlock() + + if s.conn == nil { + d := net.Dialer{Timeout: s.dialTimeout} + conn, err := d.DialContext(ctx, "unix", s.path) + if err != nil { + return fmt.Errorf("deepscan dial %s: %w", s.path, err) + } + s.conn = conn + } + + if s.writeTimeout > 0 { + _ = s.conn.SetWriteDeadline(time.Now().Add(s.writeTimeout)) + } + if _, err := s.conn.Write(frame); err != nil { + // Broken pipe / timeout → drop the connection so the next Send redials + // instead of reusing a wedged one. + _ = s.conn.Close() + s.conn = nil + return fmt.Errorf("deepscan write %s: %w", s.path, err) + } + return nil +} + +func (s *unixSink) Close() error { + s.mu.Lock() + defer s.mu.Unlock() + if s.conn == nil { + return nil + } + err := s.conn.Close() + s.conn = nil + return err +} diff --git a/deepscan/sink_test.go b/deepscan/sink_test.go new file mode 100644 index 0000000..fe35f4a --- /dev/null +++ b/deepscan/sink_test.go @@ -0,0 +1,111 @@ +package deepscan + +import ( + "context" + "encoding/binary" + "encoding/json" + "io" + "net" + "path/filepath" + "testing" + "time" +) + +// The unix sink must write a well-formed [version][LE len][JSON] frame that a +// listener can parse back into {prompt, spans} — the contract the deep-scan +// reader (D2, Python) will implement. +func TestUnixSink_WritesFrame(t *testing.T) { + // Short filename — macOS caps the full unix socket path at ~104 bytes. + sockPath := filepath.Join(t.TempDir(), "s") + ln, err := net.Listen("unix", sockPath) + if err != nil { + t.Fatalf("listen: %v", err) + } + defer ln.Close() + + type received struct { + frame []byte + err error + } + recvCh := make(chan received, 1) + go func() { + conn, err := ln.Accept() + if err != nil { + recvCh <- received{err: err} + return + } + defer conn.Close() + header := make([]byte, frameHeaderLen) + if _, err := io.ReadFull(conn, header); err != nil { + recvCh <- received{err: err} + return + } + n := binary.LittleEndian.Uint32(header[1:frameHeaderLen]) + body := make([]byte, n) + if _, err := io.ReadFull(conn, body); err != nil { + recvCh <- received{err: err} + return + } + recvCh <- received{frame: append(header, body...)} + }() + + sink := NewUnixSink(sockPath, time.Second, time.Second) + defer sink.Close() + + const prompt = "re-scan me 0123" + spans := []Span{{Start: 0, End: 2, Category: "pii", EntityType: "PHONE", Detector: "regex", Confidence: 90}} + body, err := EncodePayloadV1(prompt, spans) + if err != nil { + t.Fatalf("EncodePayloadV1: %v", err) + } + if err := sink.Send(context.Background(), encodeFrame(body)); err != nil { + t.Fatalf("Send: %v", err) + } + + select { + case got := <-recvCh: + if got.err != nil { + t.Fatalf("listener: %v", got.err) + } + if got.frame[0] != ProtocolVersion { + t.Errorf("version = %d, want %d", got.frame[0], ProtocolVersion) + } + var p payloadV1 + if err := json.Unmarshal(got.frame[frameHeaderLen:], &p); err != nil { + t.Fatalf("unmarshal payload: %v", err) + } + if p.Prompt != prompt { + t.Errorf("prompt = %q, want %q", p.Prompt, prompt) + } + if len(p.Spans) != 1 || p.Spans[0].EntityType != "PHONE" || p.Spans[0].Detector != "regex" { + t.Errorf("spans = %+v, want one PHONE/regex span", p.Spans) + } + case <-time.After(2 * time.Second): + t.Fatal("listener did not receive a frame") + } +} + +// Dialing an absent socket must surface an error (so the forwarder counts a +// failure) and must not panic — the deep-scan process may not be up yet. +func TestUnixSink_DialErrorNoPanic(t *testing.T) { + sink := NewUnixSink(filepath.Join(t.TempDir(), "absent"), 200*time.Millisecond, 200*time.Millisecond) + defer sink.Close() + if err := sink.Send(context.Background(), encodeFrame([]byte("x"))); err == nil { + t.Error("Send to absent socket should return an error") + } +} + +// encodeFrame header round-trips: version byte + little-endian length. +func TestEncodeFrame_Header(t *testing.T) { + body := []byte(`{"prompt":"x"}`) + frame := encodeFrame(body) + if frame[0] != ProtocolVersion { + t.Errorf("version = %d, want %d", frame[0], ProtocolVersion) + } + if n := binary.LittleEndian.Uint32(frame[1:frameHeaderLen]); int(n) != len(body) { + t.Errorf("length = %d, want %d", n, len(body)) + } + if string(frame[frameHeaderLen:]) != string(body) { + t.Errorf("body mismatch") + } +} diff --git a/deepscan/testdata/golden_frames.json b/deepscan/testdata/golden_frames.json new file mode 100644 index 0000000..8e5980e --- /dev/null +++ b/deepscan/testdata/golden_frames.json @@ -0,0 +1,26 @@ +{ + "v1_task": { + "description": "v1 task frame: fire-and-forget, carries the fast layer's spans because the DETECTOR sends it", + "frame_hex": "019e0000007b2270726f6d7074223a22e5aea2e688b7e6898be69cbae58fb720313338303031333830303020e8afb7e6a0b8e5afb9222c227370616e73223a5b7b227374617274223a362c22656e64223a31372c2263617465676f7279223a22706969222c22656e746974795f74797065223a22434e5f50484f4e45222c226465746563746f72223a227265676578222c22636f6e666964656e6365223a39357d5d7d", + "body_json": "{\"prompt\":\"客户手机号 13800138000 请核对\",\"spans\":[{\"start\":6,\"end\":17,\"category\":\"pii\",\"entity_type\":\"CN_PHONE\",\"detector\":\"regex\",\"confidence\":95}]}", + "version": 1 + }, + "v2_result": { + "description": "v2 result frame: absolute offsets, no raw matched text", + "frame_hex": "02290200007b226a6f625f6964223a226a6f622d676f6c64656e2d31222c22737461747573223a227061727469616c222c22726561736f6e223a2277696e646f775f636170222c227363616e6e65645f6279746573223a3234303030302c22746f74616c5f6279746573223a3236323134342c22656e67696e6573223a7b22626765223a7b22737461747573223a227061727469616c222c22726561736f6e223a2277696e646f775f636170222c2277696e646f7773223a343039367d2c2272756c6573223a7b22737461747573223a22636f6d706c657465222c226368756e6b73223a322c226465746563746f725f76657273696f6e223a2239356565626231222c22636f6e74656e745f76657273696f6e223a2263762d37227d7d2c2266696e64696e6773223a5b7b22656e67696e65223a2272756c6573222c2272756c655f6964223a227365637265742e636f6e6e656374696f6e2d7572692d70617373776f7264222c2263617465676f7279223a22736563726574222c22656e746974795f74797065223a2243524544454e5449414c5f44534e222c227365766572697479223a2268696768222c22636f6e666964656e6365223a39352c227374617274223a3235303030302c22656e64223a3235303037322c226465746563746f72223a227265676578227d5d2c2272756c655f7665726469637473223a5b7b227374617274223a3234353736302c22656e64223a3236323134342c22616374696f6e223a22626c6f636b227d5d7d", + "body_json": "{\"job_id\":\"job-golden-1\",\"status\":\"partial\",\"reason\":\"window_cap\",\"scanned_bytes\":240000,\"total_bytes\":262144,\"engines\":{\"bge\":{\"status\":\"partial\",\"reason\":\"window_cap\",\"windows\":4096},\"rules\":{\"status\":\"complete\",\"chunks\":2,\"detector_version\":\"95eebb1\",\"content_version\":\"cv-7\"}},\"findings\":[{\"engine\":\"rules\",\"rule_id\":\"secret.connection-uri-password\",\"category\":\"secret\",\"entity_type\":\"CREDENTIAL_DSN\",\"severity\":\"high\",\"confidence\":95,\"start\":250000,\"end\":250072,\"detector\":\"regex\"}],\"rule_verdicts\":[{\"start\":245760,\"end\":262144,\"action\":\"block\"}]}", + "version": 2 + }, + "v2_result_reject": { + "description": "v2 reject: job_id + reject and nothing else", + "frame_hex": "026a0000007b226a6f625f6964223a226a6f622d676f6c64656e2d31222c2272656a656374223a2274656e616e745f6d69736d61746368222c22656e67696e6573223a7b22626765223a7b22737461747573223a22227d2c2272756c6573223a7b22737461747573223a22227d7d7d", + "body_json": "{\"job_id\":\"job-golden-1\",\"reject\":\"tenant_mismatch\",\"engines\":{\"bge\":{\"status\":\"\"},\"rules\":{\"status\":\"\"}}}", + "version": 2 + }, + "v2_task": { + "description": "v2 task frame: NO spans and NO seat/session identity — the node re-scans the head itself", + "frame_hex": "02780100007b2276657273696f6e223a322c226a6f625f6964223a226a6f622d676f6c64656e2d31222c22746f6b656e223a22736374312e6f72675f612e313735373632303030302e6b312e5a6d3976596d4679222c2274656e616e745f6964223a226f72675f61222c2261756469745f756e69745f6964223a2261755f30313233343536373839616263646566222c22636f6e74656e745f736861323536223a226231393436616339323439326432333437633632333562346432363131313834222c22736f75726365223a2272657175657374222c22686561645f6279746573223a31363338342c22656e67696e6573223a5b22626765222c2272756c6573225d2c2272756c655f6368756e6b73223a5b7b227374617274223a31343333362c22656e64223a33303732307d2c7b227374617274223a32383637322c22656e64223a34353035367d5d2c2270726f6d7074223a22e5aea2e688b7e6898be69cbae58fb720313338303031333830303020e8afb7e6a0b8e5afb9227d", + "body_json": "{\"version\":2,\"job_id\":\"job-golden-1\",\"token\":\"sct1.org_a.1757620000.k1.Zm9vYmFy\",\"tenant_id\":\"org_a\",\"audit_unit_id\":\"au_0123456789abcdef\",\"content_sha256\":\"b1946ac92492d2347c6235b4d2611184\",\"source\":\"request\",\"head_bytes\":16384,\"engines\":[\"bge\",\"rules\"],\"rule_chunks\":[{\"start\":14336,\"end\":30720},{\"start\":28672,\"end\":45056}],\"prompt\":\"客户手机号 13800138000 请核对\"}", + "version": 2 + } +} diff --git a/scanchunk/chunk.go b/scanchunk/chunk.go new file mode 100644 index 0000000..9bad7a9 --- /dev/null +++ b/scanchunk/chunk.go @@ -0,0 +1,110 @@ +// Package scanchunk splits a content piece into the fixed-size, overlapping +// ranges the asynchronous rule lane scans, and merges what comes back. +// +// ONE implementation, on purpose (design §3.6): the proxy computes the ranges +// and ships them inside the task frame, and the receiver — Go on the local +// executor, PYTHON on a scan node — only slices what it was told to slice. A +// second chunker on the Python side would be a second thing that can drift, and +// the symptom of that drift is a missed credential, not a crash. +package scanchunk + +import ( + "unicode/utf8" + + "github.com/AiKeyLabs/pkg/deepscan" +) + +// SizeBytes is the maximum bytes in one chunk. It matches the detector's +// small-lane input cap (ai-compliance-detector engine.go) — a chunk larger than +// this is silently truncated by the detector, which is the blind spot this whole +// lane exists to close. +const SizeBytes = 16 * 1024 + +// OverlapDefault is how much consecutive chunks overlap, so an entity sitting on +// a boundary still appears WHOLE in at least one chunk. +// +// 🔴 2048, NOT 512. MEASURED 2026-09-11 with ai-compliance-detector's +// tools/rulespan, which computes each built-in rule's maximum possible match +// length from its parsed syntax tree: 9 rules can match more than 512 bytes, the +// longest being secret.connection-uri-password (CREDENTIAL_DSN) at 1939 bytes. +// With a 512-byte overlap such a credential can span the boundary so that +// NEITHER chunk contains it whole — both miss it, and nothing anywhere reports a +// gap. User decision 2026-09-11 after the measurement was presented. +// +// 22 further rules have an UNBOUNDED maximum match (secret.private-key — a PEM +// key — secret.jwt, secret.generic-api-key, …). No fixed overlap can cover those; +// they are a registered known limitation (design §10 R15), not something this +// constant should be inflated to chase. +// +// Before changing this, re-run: go run ./tools/rulespan internal/baselines/built-in +// Fence: TestChunker_OverlapCoversLongestBoundedRule. +const OverlapDefault = 2048 + +// Range is a half-open [Start, End) byte interval into the piece. +type Range = deepscan.Range + +// Chunks splits text[start:] into ranges of at most size bytes, overlapping by +// overlap bytes, with every boundary pulled back to a UTF-8 rune boundary. +// +// Returns nil when there is nothing left to scan (start at or past the end, or +// empty text) — a piece the fast layer already covered in full must produce no +// rule work at all. +func Chunks(text string, start, size, overlap int) []Range { + if size <= 0 { + size = SizeBytes + } + if overlap < 0 { + overlap = 0 + } + if overlap >= size { + // An overlap at or above the chunk size would never advance. + overlap = size / 2 + } + if start < 0 { + start = 0 + } + if start >= len(text) { + return nil + } + // Pull the very first boundary back to a rune start too: `start` is derived + // from head_bytes minus the overlap and has no reason to be rune-aligned. + start = runeStartAtOrBefore(text, start) + + var out []Range + pos := start + for pos < len(text) { + end := pos + size + if end >= len(text) { + out = append(out, Range{Start: pos, End: len(text)}) + break + } + end = runeStartAtOrBefore(text, end) + if end <= pos { + // A single rune longer than the whole chunk cannot happen in UTF-8 + // (max 4 bytes), but refusing to emit a non-advancing chunk keeps this + // loop structurally unable to spin. + end = pos + size + } + out = append(out, Range{Start: pos, End: end}) + next := end - overlap + if next <= pos { + next = end + } + pos = runeStartAtOrBefore(text, next) + } + return out +} + +// runeStartAtOrBefore returns the largest index <= i that begins a UTF-8 rune. +func runeStartAtOrBefore(s string, i int) int { + if i <= 0 { + return 0 + } + if i >= len(s) { + return len(s) + } + for i > 0 && !utf8.RuneStart(s[i]) { + i-- + } + return i +} diff --git a/scanchunk/chunk_test.go b/scanchunk/chunk_test.go new file mode 100644 index 0000000..622c294 --- /dev/null +++ b/scanchunk/chunk_test.go @@ -0,0 +1,161 @@ +package scanchunk + +import ( + "strings" + "testing" + "unicode/utf8" + + "github.com/AiKeyLabs/pkg/deepscan" +) + +// TestChunker_ChunksWithinSmallLaneAndRuneBoundary: every chunk must fit the +// detector's small lane (16 KiB input cap) and must not cut a UTF-8 rune in +// half, and together the chunks must cover every byte from start to the end of +// the piece. +// +// The rune-boundary half is not cosmetic: a chunk ending mid-rune hands the +// detector invalid UTF-8, and the regex engine's behaviour on invalid UTF-8 is +// not something a compliance guarantee should rest on. +func TestChunker_ChunksWithinSmallLaneAndRuneBoundary(t *testing.T) { + // CJK: 3 bytes per rune, so byte-cutting will land mid-rune constantly + // unless the implementation backs off deliberately. + text := strings.Repeat("客户资料需要保密不要外发。", 6000) // ~234 KiB + for _, start := range []int{0, 1, 12345, len(text) - 10} { + chunks := Chunks(text, start, SizeBytes, OverlapDefault) + if len(chunks) == 0 { + t.Fatalf("start=%d produced no chunks for a %d-byte piece", start, len(text)) + } + covered := make([]bool, len(text)) + for i, c := range chunks { + if c.End-c.Start > SizeBytes { + t.Errorf("start=%d chunk %d is %d bytes, over the %d-byte small lane", start, i, c.End-c.Start, SizeBytes) + } + if c.Start < 0 || c.End > len(text) || c.Start >= c.End { + t.Fatalf("start=%d chunk %d has impossible bounds %+v (text %d bytes)", start, i, c, len(text)) + } + if !utf8.ValidString(text[c.Start:c.End]) { + t.Errorf("start=%d chunk %d [%d,%d) is not valid UTF-8 — a rune was cut in half", start, i, c.Start, c.End) + } + for b := c.Start; b < c.End; b++ { + covered[b] = true + } + } + for b := chunks[0].Start; b < len(text); b++ { + if !covered[b] { + t.Fatalf("start=%d byte %d is not covered by any chunk", start, b) + break + } + } + if chunks[len(chunks)-1].End != len(text) { + t.Errorf("start=%d last chunk ends at %d, not the end of the piece (%d)", start, chunks[len(chunks)-1].End, len(text)) + } + } +} + +// TestChunker_OverlapCatchesBoundaryEntity: an entity straddling a chunk +// boundary must appear WHOLE inside at least one chunk, or both chunks see half +// of it and neither reports it — a silent miss with no signal anywhere. +func TestChunker_OverlapCatchesBoundaryEntity(t *testing.T) { + const secret = "postgres://svc_user:S3cr3tP4ssw0rd@10.2.3.4:5432/analytics" + filler := strings.Repeat("a", SizeBytes-len(secret)/2) + text := filler + secret + strings.Repeat("b", SizeBytes) + + chunks := Chunks(text, 0, SizeBytes, OverlapDefault) + idx := strings.Index(text, secret) + whole := false + for _, c := range chunks { + if c.Start <= idx && idx+len(secret) <= c.End { + whole = true + break + } + } + if !whole { + t.Fatalf("a %d-byte credential straddling the boundary at %d appears whole in NO chunk: %+v", len(secret), SizeBytes, chunks) + } +} + +// TestChunker_OverlapCoversLongestBoundedRule is the fence for the P0 §F8 +// measurement that changed this package's constant. +// +// MEASURED 2026-09-11 (ai-compliance-detector/tools/rulespan): 9 built-in rules +// can match MORE than the 512-byte overlap the design originally specified, the +// longest being secret.connection-uri-password at 1939 bytes. With overlap=512 a +// 1939-byte DSN sitting across a chunk boundary is split so that neither chunk +// contains it whole — both miss it, and nothing reports a gap. User decision +// 2026-09-11: raise the overlap to 2048. +// +// If this test ever fails because someone lowered OverlapDefault, the question +// to ask is NOT "how do I make the test pass" — it is "has the rule set's longest +// bounded match changed?", and the way to answer it is to re-run +// `go run ./tools/rulespan internal/baselines/built-in` in ai-compliance-detector. +func TestChunker_OverlapCoversLongestBoundedRule(t *testing.T) { + const longestBoundedRuleBytes = 1939 // secret.connection-uri-password, measured + + if OverlapDefault < longestBoundedRuleBytes { + t.Fatalf("OverlapDefault is %d, below the measured longest bounded rule match (%d bytes). "+ + "A credential that long can straddle a chunk boundary and be missed by BOTH chunks.", + OverlapDefault, longestBoundedRuleBytes) + } + + // And prove it end to end, not just by comparing constants. + entity := strings.Repeat("X", longestBoundedRuleBytes) + // Place it so it starts just before a boundary and runs past it. + head := strings.Repeat("a", SizeBytes-20) + text := head + entity + strings.Repeat("b", SizeBytes) + idx := len(head) + + for _, c := range Chunks(text, 0, SizeBytes, OverlapDefault) { + if c.Start <= idx && idx+len(entity) <= c.End { + return // found whole in one chunk + } + } + t.Fatalf("a %d-byte entity starting at %d (boundary %d) appears whole in no chunk", len(entity), idx, SizeBytes) +} + +// TestChunker_SinglePieceNoChunks: a piece the fast layer already covered in +// full produces no rule work at all. Without this the async lane re-scans every +// short prompt for nothing (design §3.6 / R-scan-node-deepscan-16.S3). +func TestChunker_SinglePieceNoChunks(t *testing.T) { + text := strings.Repeat("a", 8000) + if got := Chunks(text, len(text), SizeBytes, OverlapDefault); len(got) != 0 { + t.Errorf("start at end of piece must yield no chunks, got %+v", got) + } + if got := Chunks("", 0, SizeBytes, OverlapDefault); len(got) != 0 { + t.Errorf("empty text must yield no chunks, got %+v", got) + } +} + +// TestChunker_MergeFindingsDeduplicatesOverlapRegion: the overlap means a +// credential sitting in it is found TWICE, once per chunk. Merging on +// (entity_type, absolute start, absolute end) collapses those, while two genuinely +// different entity types at the same offsets stay separate. +func TestChunker_MergeFindingsDeduplicatesOverlapRegion(t *testing.T) { + dup := deepscan.Finding{Engine: "rules", EntityType: "CREDENTIAL_DSN", Category: "secret", Start: 16300, End: 16360, Severity: "high", Confidence: 90} + sameSpanOtherType := dup + sameSpanOtherType.EntityType = "CREDENTIAL_PASSWORD" + other := dup + other.Start, other.End = 40000, 40011 + + got := MergeFindings([]ChunkFindings{ + {Range: Range{Start: 0, End: 16384}, Findings: []deepscan.Finding{dup}}, + {Range: Range{Start: 14336, End: 30720}, Findings: []deepscan.Finding{dup, sameSpanOtherType}}, + {Range: Range{Start: 30720, End: 47104}, Findings: []deepscan.Finding{other}}, + }) + if len(got) != 3 { + t.Fatalf("want 3 merged findings (dup collapsed, other type kept, third kept), got %d: %+v", len(got), got) + } + seen := map[string]int{} + for _, f := range got { + seen[f.EntityType]++ + } + if seen["CREDENTIAL_DSN"] != 2 || seen["CREDENTIAL_PASSWORD"] != 1 { + t.Errorf("merge collapsed the wrong things: %v", seen) + } + // Output must be ordered by offset so the audit event reads top to bottom. + for i := 1; i < len(got); i++ { + if got[i-1].Start > got[i].Start { + t.Errorf("findings are not ordered by offset: %+v", got) + break + } + } +} diff --git a/scanchunk/go.mod b/scanchunk/go.mod new file mode 100644 index 0000000..eb54189 --- /dev/null +++ b/scanchunk/go.mod @@ -0,0 +1,7 @@ +module github.com/AiKeyLabs/pkg/scanchunk + +go 1.26.1 + +require github.com/AiKeyLabs/pkg/deepscan v0.0.0 + +replace github.com/AiKeyLabs/pkg/deepscan => ../deepscan diff --git a/scanchunk/merge.go b/scanchunk/merge.go new file mode 100644 index 0000000..f66c93a --- /dev/null +++ b/scanchunk/merge.go @@ -0,0 +1,67 @@ +package scanchunk + +import ( + "sort" + + "github.com/AiKeyLabs/pkg/deepscan" +) + +// ChunkFindings is one chunk's range plus what was found inside it. Offsets in +// Findings are ABSOLUTE into the piece — the producer adds the chunk origin +// before handing them over, because only the producer knows it. +type ChunkFindings struct { + Range Range + Findings []deepscan.Finding +} + +// MergeFindings collapses the duplicates the overlap creates and returns one +// list ordered by position. +// +// The identity is (entity_type, absolute start, absolute end). Why those three +// and not the whole finding: a credential inside the overlap region is found +// once per chunk, and the two copies can legitimately differ in confidence or +// rule_id (a different rule may win in a different context window) — keeping +// both would double-count one credential in the audit record. Two DIFFERENT +// entity types at the same offsets are kept apart on purpose: a string that is +// both a password and part of a DSN is two findings a reviewer needs to see. +// +// Engine is NOT part of the identity because callers pass one engine's findings +// at a time; bge and rules results are never merged (design §3.6: 跨引擎不去重). +func MergeFindings(chunks []ChunkFindings) []deepscan.Finding { + type key struct { + entityType string + start, end int + } + seen := make(map[key]deepscan.Finding) + order := make([]key, 0) + for _, c := range chunks { + for _, f := range c.Findings { + k := key{f.EntityType, f.Start, f.End} + prev, ok := seen[k] + if !ok { + seen[k] = f + order = append(order, k) + continue + } + // Same span, same type, seen twice: keep the higher-confidence copy so + // a merge never lowers a finding's severity signal. + if f.Confidence > prev.Confidence { + seen[k] = f + } + } + } + out := make([]deepscan.Finding, 0, len(order)) + for _, k := range order { + out = append(out, seen[k]) + } + sort.SliceStable(out, func(i, j int) bool { + if out[i].Start != out[j].Start { + return out[i].Start < out[j].Start + } + if out[i].End != out[j].End { + return out[i].End < out[j].End + } + return out[i].EntityType < out[j].EntityType + }) + return out +} diff --git a/scannode/client_test.go b/scannode/client_test.go new file mode 100644 index 0000000..84401e7 --- /dev/null +++ b/scannode/client_test.go @@ -0,0 +1,267 @@ +package scannode + +import ( + "context" + "crypto/ecdsa" + "crypto/elliptic" + "crypto/rand" + "crypto/sha256" + "crypto/tls" + "crypto/x509" + "crypto/x509/pkix" + "encoding/hex" + "errors" + "math/big" + "net" + "sync/atomic" + "testing" + "time" + + "github.com/AiKeyLabs/pkg/deepscan" +) + +// --------------------------------------------------------------------------- +// These four fences all guard ONE property, from four directions: +// +// a piece of an employee's raw prompt must never reach a box we did not +// independently decide to trust. +// +// A scan node receives RAW USER CONTENT. That is the whole point of it and also +// the whole risk: every weakening here is a direct content-disclosure path, so +// each fence asserts BYTES RECEIVED == 0, not merely "an error was returned". +// An implementation that writes the frame and then notices the problem passes a +// naive error-check test and still leaked the prompt. +// --------------------------------------------------------------------------- + +// tlsProbe is a real TLS listener with a self-signed cert. It counts how many +// APPLICATION bytes it managed to read, which is the only number these tests +// actually care about. +type tlsProbe struct { + addr string + fingerprint string + bytesRead atomic.Int64 + conns atomic.Int64 + // reply, when set, is written back after the first bytes of each read — + // a node answering on the same connection. + reply atomic.Pointer[[]byte] + close func() +} + +func startTLSProbe(t *testing.T) *tlsProbe { + t.Helper() + key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader) + if err != nil { + t.Fatalf("gen key: %v", err) + } + tmpl := &x509.Certificate{ + SerialNumber: big.NewInt(1), + Subject: pkix.Name{CommonName: "scan-node-probe"}, + NotBefore: time.Now().Add(-time.Hour), + NotAfter: time.Now().Add(time.Hour), + IPAddresses: []net.IP{net.ParseIP("127.0.0.1")}, + KeyUsage: x509.KeyUsageDigitalSignature | x509.KeyUsageCertSign, + ExtKeyUsage: []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth}, + IsCA: true, + } + der, err := x509.CreateCertificate(rand.Reader, tmpl, tmpl, &key.PublicKey, key) + if err != nil { + t.Fatalf("create cert: %v", err) + } + sum := sha256.Sum256(der) + p := &tlsProbe{fingerprint: hex.EncodeToString(sum[:])} + + ln, err := tls.Listen("tcp", "127.0.0.1:0", &tls.Config{ + Certificates: []tls.Certificate{{Certificate: [][]byte{der}, PrivateKey: key}}, + MinVersion: tls.VersionTLS12, + }) + if err != nil { + t.Fatalf("tls listen: %v", err) + } + p.addr = "https://" + ln.Addr().String() + done := make(chan struct{}) + go func() { + for { + c, err := ln.Accept() + if err != nil { + close(done) + return + } + p.conns.Add(1) + go func(c net.Conn) { + defer c.Close() + buf := make([]byte, 4096) + for { + _ = c.SetReadDeadline(time.Now().Add(2 * time.Second)) + n, err := c.Read(buf) + p.bytesRead.Add(int64(n)) + if rp := p.reply.Load(); rp != nil && n > 0 { + _, _ = c.Write(*rp) + } + if err != nil { + return + } + } + }(c) + } + }() + p.close = func() { _ = ln.Close(); <-done } + t.Cleanup(p.close) + return p +} + +func frame(t *testing.T) []byte { + t.Helper() + b, err := deepscan.EncodeFrameV2(deepscan.FrameV2{ + JobID: "job-1", TenantID: "org_a", ContentSHA256: "abc", + Source: deepscan.SourceRequest, Prompt: "客户手机号 13800138000", + }) + if err != nil { + t.Fatalf("encode: %v", err) + } + return b +} + +func trustAll(string) bool { return true } +func trustNone(string) bool { return false } + +// TestScanNodeClient_PlaintextNodeRefused: a node URL that is not https is +// refused before a socket is opened. There is no "plaintext mode" — user +// decision D21, 2026-09-11: 「强制 TLS…永无明文模式」. +func TestScanNodeClient_PlaintextNodeRefused(t *testing.T) { + p := startTLSProbe(t) + host := p.addr[len("https://"):] + + sink := NewTLSSink(Node{ID: "n1", Addr: "http://" + host, Fingerprint: p.fingerprint}, + trustAll, time.Second, time.Second) + err := sink.Send(context.Background(), frame(t)) + if err == nil { + t.Fatal("a plaintext node URL was accepted") + } + if p.conns.Load() != 0 || p.bytesRead.Load() != 0 { + t.Errorf("the refusal happened AFTER touching the node: conns=%d bytes=%d", + p.conns.Load(), p.bytesRead.Load()) + } +} + +// TestScanNodeClient_FingerprintMismatchAbortsBeforeBody: the node presents a +// valid, working TLS cert — just not the one we pinned. The handshake must fail +// and ZERO application bytes may be written. This is the fence for "someone +// stood up a look-alike node inside the network". +func TestScanNodeClient_FingerprintMismatchAbortsBeforeBody(t *testing.T) { + p := startTLSProbe(t) + wrong := "00" + p.fingerprint[2:] + + sink := NewTLSSink(Node{ID: "n1", Addr: p.addr, Fingerprint: wrong}, + trustAll, 2*time.Second, 2*time.Second) + if err := sink.Send(context.Background(), frame(t)); err == nil { + t.Fatal("a node presenting an unpinned certificate was accepted") + } + if got := p.bytesRead.Load(); got != 0 { + t.Errorf("%d application bytes reached a node whose fingerprint did not match — the prompt leaked", got) + } +} + +// TestScanNodeClient_UnlistedFingerprintNeverReceivesBytes: the pin matches what +// the node presents, but the fingerprint is not in the trust set the control +// plane handed us. Same outcome, different gate — this is the one that catches +// "the node self-reported its own fingerprint and we believed it". +func TestScanNodeClient_UnlistedFingerprintNeverReceivesBytes(t *testing.T) { + p := startTLSProbe(t) + + sink := NewTLSSink(Node{ID: "n1", Addr: p.addr, Fingerprint: p.fingerprint}, + trustNone, 2*time.Second, 2*time.Second) + if err := sink.Send(context.Background(), frame(t)); err == nil { + t.Fatal("a node outside the trust set was accepted") + } + if got := p.bytesRead.Load(); got != 0 { + t.Errorf("%d application bytes reached an untrusted node — the prompt leaked", got) + } + if got := p.conns.Load(); got != 0 { + t.Errorf("the untrusted node was dialed at all (conns=%d); the allowlist must be checked before the socket", got) + } +} + +// TestScanNodeClient_PublicAddressRefused: R-scan-node-deepscan-10 — a node +// address that resolves outside RFC1918 / RFC4193 / loopback is refused. A +// misconfigured or tampered node list must not become an exfiltration channel +// for raw prompts. +func TestScanNodeClient_PublicAddressRefused(t *testing.T) { + for _, addr := range []string{ + "https://8.8.8.8:27411", + "https://1.1.1.1:27411", + "https://[2001:4860:4860::8888]:27411", + } { + sink := NewTLSSink(Node{ID: "n1", Addr: addr, Fingerprint: "aa"}, trustAll, time.Second, time.Second) + if err := sink.Send(context.Background(), frame(t)); err == nil { + t.Errorf("public address %s was accepted", addr) + } + } + // ...and the private ones it exists to allow are not collateral damage. + for _, addr := range []string{ + "https://10.2.3.4:27411", + "https://192.168.1.9:27411", + "https://172.16.0.1:27411", + "https://127.0.0.1:27411", + "https://[fd00::1]:27411", + } { + if err := checkPrivateAddress(addr); err != nil { + t.Errorf("private address %s was refused: %v", addr, err) + } + } +} + +func replyWith(t *testing.T, p *tlsProbe, r deepscan.ResultFrame) { + t.Helper() + b, err := deepscan.EncodeResult(r) + if err != nil { + t.Fatalf("encode result: %v", err) + } + p.reply.Store(&b) +} + +// TestScanNodeClient_ReadsTheResultBack — the sink returns the node's answer. +// +// 🔴 It used to return right after writing the frame, so no finding a node +// produced ever reached the proxy. +// bugfix: workflow/CI/bugfix/20260913-async-scan-lane-never-returned-findings.md +func TestScanNodeClient_ReadsTheResultBack(t *testing.T) { + p := startTLSProbe(t) + replyWith(t, p, deepscan.ResultFrame{JobID: "job-1", Status: deepscan.StatusComplete, + Findings: []deepscan.Finding{{Engine: deepscan.EngineRules, EntityType: "CN_PHONE", Start: 10, End: 21}}}) + sink := NewTLSSink(Node{ID: "n1", Addr: p.addr, Fingerprint: p.fingerprint, Weight: 1}, trustAll, time.Second, 2*time.Second) + defer sink.Close() + + if err := sink.Send(context.Background(), frame(t)); err != nil { + t.Fatalf("send: %v", err) + } + rs, ok := sink.(interface { + Results() <-chan deepscan.ResultFrame + }) + if !ok { + t.Fatal("the TLS sink does not expose Results()") + } + select { + case r := <-rs.Results(): + if r.JobID != "job-1" || len(r.Findings) != 1 { + t.Fatalf("result altered in transit: %+v", r) + } + default: + t.Fatal("Send returned without the node's result being readable — the answer was not read back") + } +} + +// TestScanNodeClient_RejectIsATypedError — a refusal comes back as a +// RejectError carrying the code, never as a result and never as a generic error +// (the forwarder decides by code whether another node may be tried). +func TestScanNodeClient_RejectIsATypedError(t *testing.T) { + p := startTLSProbe(t) + replyWith(t, p, deepscan.ResultFrame{JobID: "job-1", Reject: deepscan.RejectTenantMismatch}) + sink := NewTLSSink(Node{ID: "n1", Addr: p.addr, Fingerprint: p.fingerprint, Weight: 1}, trustAll, time.Second, 2*time.Second) + defer sink.Close() + + err := sink.Send(context.Background(), frame(t)) + var rej *deepscan.RejectError + if !errors.As(err, &rej) || rej.Code != deepscan.RejectTenantMismatch { + t.Fatalf("send error = %v, want RejectError{%s}", err, deepscan.RejectTenantMismatch) + } +} diff --git a/scannode/go.mod b/scannode/go.mod new file mode 100644 index 0000000..f84a22c --- /dev/null +++ b/scannode/go.mod @@ -0,0 +1,12 @@ +module github.com/AiKeyLabs/pkg/scannode + +go 1.26.1 + +require ( + github.com/AiKeyLabs/pkg/deepscan v0.0.0 + github.com/AiKeyLabs/pkg/seatassign v0.0.0 +) + +replace github.com/AiKeyLabs/pkg/deepscan => ../deepscan + +replace github.com/AiKeyLabs/pkg/seatassign => ../seatassign diff --git a/scannode/rank.go b/scannode/rank.go new file mode 100644 index 0000000..a16b215 --- /dev/null +++ b/scannode/rank.go @@ -0,0 +1,137 @@ +package scannode + +import ( + "sync" + "time" + + "github.com/AiKeyLabs/pkg/seatassign" +) + +// Rank orders nodes by preference for one content key, best first. +// +// The key is `tenant_id‖content_sha256` (design §4b.2). Same content → same node, +// every time, from every proxy — which is what makes the node's own local record +// a de-facto shared "already scanned" layer WITHOUT a shared table (D18/D19). +// +// 🔴 It reuses pkg/seatassign rather than hashing here. That package is the +// weighted-rendezvous implementation master and proxy already agree on +// byte-for-byte, and a second implementation of the same idea is exactly how two +// sides of a distributed system quietly stop choosing the same box. The +// vocabulary mismatch (Account/seatID vs Node/contentKey) is cosmetic: both are +// "pick one of N by a stable key, weighted". +func Rank(nodes []Node, key string) []Node { + if len(nodes) == 0 { + return nil + } + accounts := make([]seatassign.Account, 0, len(nodes)) + byID := make(map[string]Node, len(nodes)) + for _, n := range nodes { + w := float64(n.Weight) + if w <= 0 { + w = 1 + } + accounts = append(accounts, seatassign.Account{AccountID: n.ID, Weight: w}) + byID[n.ID] = n + } + ranked := seatassign.Rank(key, accounts) + out := make([]Node, 0, len(ranked)) + for _, a := range ranked { + out = append(out, byID[a.AccountID]) + } + return out +} + +// Breaker decides whether a node may be tried right now. +type Breaker interface { + Allow(nodeID string) bool + Report(nodeID string, ok bool) +} + +// NewBreaker returns a per-node consecutive-failure breaker with the default +// half-open cooldown. +func NewBreaker(threshold int) Breaker { + return NewBreakerWithCooldown(threshold, DefaultBreakerCooldown) +} + +// DefaultBreakerCooldown is how long an open breaker refuses a node before it +// lets ONE attempt through again. +const DefaultBreakerCooldown = 30 * time.Second + +// NewBreakerWithCooldown returns a per-node breaker that opens after threshold +// consecutive failures and, once open, refuses the node until cooldown has +// elapsed — then allows attempts again so a recovered node can prove itself. +// +// 🔴 THE COOLDOWN IS NOT OPTIONAL, and leaving it out is not a smaller design — +// it is a broken one. A breaker that opens on failures and only closes on a +// success, while also refusing every attempt, can never close: the one request +// that would have proven the node healthy is the request it blocks. That is not +// hypothetical; it is what this package did until +// TestDeepScanForward_DegradesAfterThreeFailuresAndRecovers caught it on +// 2026-09-12. The node came back, and the lane stayed dark until a proxy restart. +// +// Half-open is deliberately generous: after the cooldown EVERY attempt is +// allowed again (not a single probe), because the cost of being wrong is one +// dropped background task, never a failed user request — and being slow to +// notice a recovery costs real compliance coverage. +func NewBreakerWithCooldown(threshold int, cooldown time.Duration) Breaker { + if threshold <= 0 { + threshold = 3 + } + if cooldown <= 0 { + cooldown = DefaultBreakerCooldown + } + return &breaker{threshold: threshold, cooldown: cooldown, state: map[string]*breakerState{}, now: time.Now} +} + +type breakerState struct { + fails int + openedAt time.Time +} + +type breaker struct { + threshold int + cooldown time.Duration + now func() time.Time + + mu sync.Mutex + state map[string]*breakerState +} + +func (b *breaker) Allow(nodeID string) bool { + b.mu.Lock() + defer b.mu.Unlock() + st, ok := b.state[nodeID] + if !ok || st.fails < b.threshold { + return true + } + if b.now().Sub(st.openedAt) >= b.cooldown { + // Half-open: let attempts through again. The counter is NOT reset here — + // only an actual success clears it (Report), so a node that is still down + // re-opens on its next failure instead of flapping every cooldown. + return true + } + return false +} + +func (b *breaker) Report(nodeID string, ok bool) { + b.mu.Lock() + defer b.mu.Unlock() + if ok { + delete(b.state, nodeID) + return + } + st := b.state[nodeID] + if st == nil { + st = &breakerState{} + b.state[nodeID] = st + } + st.fails++ + if st.fails >= b.threshold { + // >= not ==: a failure while the breaker is HALF-OPEN must restart the + // cooldown. With `==` the window would only ever be set once, so after the + // first cooldown elapsed the breaker would allow every attempt forever — + // open in name only, which is worse than no breaker because the health + // section would still report it as protecting something. + st.openedAt = b.now() + } +} diff --git a/scannode/rank_test.go b/scannode/rank_test.go new file mode 100644 index 0000000..1e05a25 --- /dev/null +++ b/scannode/rank_test.go @@ -0,0 +1,139 @@ +package scannode + +import ( + "crypto/sha256" + "encoding/hex" + "fmt" + "testing" + "time" +) + +// TestRank_SameContentSameNode: the SAME content must always rank the same node +// first, and that node must stay first when OTHER nodes come and go. +// +// WHY this matters beyond load balancing (design D18/D19, 2026-09-11): the +// product deliberately does NOT build a shared "already scanned" table. The +// affinity IS the shared layer — the same piece of content always lands on the +// same node, so that node's own local record answers "has anyone in this org +// sent this before?" without a database. Break the affinity and the de-dup +// silently becomes a full re-scan on every node, which looks like nothing at all +// except a rising CPU bill. +// +// The key is `tenant_id‖content_sha256` (design §4b.2). Note what is NOT in it: +// no seat, no session. Two employees pasting the same text hit the same node on +// purpose — that is the whole de-dup. +func TestRank_SameContentSameNode(t *testing.T) { + all := []Node{ + {ID: "scan-1", Addr: "https://10.0.0.1:27411", Weight: 1}, + {ID: "scan-2", Addr: "https://10.0.0.2:27411", Weight: 1}, + {ID: "scan-3", Addr: "https://10.0.0.3:27411", Weight: 1}, + {ID: "scan-4", Addr: "https://10.0.0.4:27411", Weight: 1}, + } + key := contentKey("org_a", "the quarterly figures, verbatim") + + first := Rank(all, key) + if len(first) != len(all) { + t.Fatalf("Rank dropped nodes: %d in, %d out", len(all), len(first)) + } + for i := 0; i < 50; i++ { + got := Rank(all, key) + if got[0].ID != first[0].ID { + t.Fatalf("same key chose %s then %s — affinity is not deterministic", first[0].ID, got[0].ID) + } + } + + // Input order must not matter: two proxies holding the same node set in + // different order have to agree, or the de-dup is per-proxy. + shuffled := []Node{all[2], all[0], all[3], all[1]} + if got := Rank(shuffled, key); got[0].ID != first[0].ID { + t.Errorf("node input order changed the winner: %s vs %s", got[0].ID, first[0].ID) + } + + // Removing an UNRELATED node must not move the winner (that is what + // rendezvous hashing buys over modulo). + var without []Node + for _, n := range all { + if n.ID != first[len(first)-1].ID { + without = append(without, n) + } + } + if got := Rank(without, key); got[0].ID != first[0].ID { + t.Errorf("dropping the LAST-ranked node moved the winner from %s to %s — that is modulo behaviour, not rendezvous", + first[0].ID, got[0].ID) + } + + // Different content must spread; identical content must not. + seen := map[string]int{} + for i := 0; i < 2000; i++ { + seen[Rank(all, contentKey("org_a", fmt.Sprintf("piece-%d", i)))[0].ID]++ + } + if len(seen) != len(all) { + t.Errorf("2000 distinct pieces only ever reached %d of %d nodes: %v", len(seen), len(all), seen) + } +} + +// TestBreaker_OpensAfterThreeFailuresAndRecovers pins the degrade rule the +// health section reports (design §3.4: "连续 3 次投递失败后为 degraded，一次成功归零"). +func TestBreaker_OpensAfterThreeFailuresAndRecovers(t *testing.T) { + b := NewBreaker(3) + if !b.Allow("scan-1") { + t.Fatal("a fresh breaker must allow") + } + b.Report("scan-1", false) + b.Report("scan-1", false) + if !b.Allow("scan-1") { + t.Error("two failures must not open the breaker — a node is allowed to blip") + } + b.Report("scan-1", false) + if b.Allow("scan-1") { + t.Error("three consecutive failures must open the breaker") + } + // A different node is unaffected: the breaker is per-node or a single bad + // box takes the whole lane down. + if !b.Allow("scan-2") { + t.Error("one node's failures opened another node's breaker") + } + b.Report("scan-1", true) + if !b.Allow("scan-1") { + t.Error("one success must close the breaker again") + } +} + +func contentKey(tenant, content string) string { + sum := sha256.Sum256([]byte(content)) + return tenant + "‖" + hex.EncodeToString(sum[:]) +} + +// TestBreaker_HalfOpenLetsARecoveredNodeBackIn is the regression fence for a +// defect this package shipped with until 2026-09-12: the breaker opened after +// three failures and then refused every attempt, including the one that would +// have proven the node healthy again. A node that recovered stayed excluded +// until the proxy restarted, and nothing anywhere said so — the lane just +// quietly produced no coverage. +// +// Caught by TestDeepScanForward_DegradesAfterThreeFailuresAndRecovers in +// aikey-proxy, fenced here at the source. +func TestBreaker_HalfOpenLetsARecoveredNodeBackIn(t *testing.T) { + b := NewBreakerWithCooldown(3, 40*time.Millisecond) + for i := 0; i < 3; i++ { + b.Report("scan-1", false) + } + if b.Allow("scan-1") { + t.Fatal("breaker must be open immediately after the threshold") + } + time.Sleep(60 * time.Millisecond) + if !b.Allow("scan-1") { + t.Fatal("after the cooldown the breaker must allow an attempt, or a recovered node can never come back") + } + // Still down: the next failure re-opens it rather than flapping open/closed. + b.Report("scan-1", false) + if b.Allow("scan-1") { + t.Error("a node that failed again during half-open must be refused again") + } + // Recovered: one success clears everything. + time.Sleep(60 * time.Millisecond) + b.Report("scan-1", true) + if !b.Allow("scan-1") { + t.Error("a success must fully close the breaker") + } +} diff --git a/scannode/sink_tls.go b/scannode/sink_tls.go new file mode 100644 index 0000000..81c209a --- /dev/null +++ b/scannode/sink_tls.go @@ -0,0 +1,192 @@ +package scannode + +import ( + "context" + "crypto/sha256" + "crypto/tls" + "encoding/hex" + "fmt" + "net" + "strings" + "sync" + "time" + + "github.com/AiKeyLabs/pkg/deepscan" +) + +// Ack codes the proxy acts on. AckAccepted is the only non-reject outcome; the +// rest alias pkg/deepscan's reject constants so there is ONE definition of each +// string shared by the encoder, this client and the Python node. +const ( + AckAccepted = "accepted" + AckBusy = deepscan.RejectBusy + AckUnauthorized = deepscan.RejectUnauthorized + AckTenantMismatch = deepscan.RejectTenantMismatch + AckBadFrame = deepscan.RejectBadFrame + AckVersionUnsupported = deepscan.RejectVersionUnsupported +) + +// tlsSink writes frames to one scan node over TLS and reads the result frame +// back on the same connection. +type tlsSink struct { + node Node + trusted func(string) bool + dialTimeout time.Duration + ackTimeout time.Duration + + mu sync.Mutex + conn net.Conn + results chan deepscan.ResultFrame +} + +// NewTLSSink returns a deepscan.Sink that delivers frames to n. +// +// GATE ORDER IS PART OF THE CONTRACT. Cheapest and most decisive first, so the +// expensive/irreversible step happens last: +// +// 1. is it in the trust set? — no socket opened +// 2. is the scheme https? — no socket opened +// 3. does it resolve private? — no socket opened +// 4. does the presented cert match the pin? — handshake fails, no app bytes +// 5. only now is the frame written +// +// Fences: TestScanNodeClient_UnlistedFingerprintNeverReceivesBytes, +// _PlaintextNodeRefused, _PublicAddressRefused, _FingerprintMismatchAbortsBeforeBody. +func NewTLSSink(n Node, trusted func(string) bool, dialTimeout, ackTimeout time.Duration) deepscan.Sink { + if dialTimeout <= 0 { + dialTimeout = 2 * time.Second + } + if ackTimeout <= 0 { + ackTimeout = 30 * time.Second + } + if trusted == nil { + // A nil predicate must mean "trust nothing". Defaulting to trust-all here + // would turn a wiring mistake into a silent content-disclosure path. + trusted = func(string) bool { return false } + } + return &tlsSink{node: n, trusted: trusted, dialTimeout: dialTimeout, ackTimeout: ackTimeout, + results: make(chan deepscan.ResultFrame, 8)} +} + +// Results exposes result frames read back from this node. +func (s *tlsSink) Results() <-chan deepscan.ResultFrame { return s.results } + +func (s *tlsSink) Send(ctx context.Context, frame []byte) error { + s.mu.Lock() + defer s.mu.Unlock() + + if s.conn == nil { + c, err := s.dial(ctx) + if err != nil { + return err + } + s.conn = c + } + _ = s.conn.SetWriteDeadline(time.Now().Add(s.dialTimeout)) + if _, err := s.conn.Write(frame); err != nil { + _ = s.conn.Close() + s.conn = nil + return fmt.Errorf("scannode %s: write: %w", s.node.ID, err) + } + // 🔴 The node answers on the SAME connection, and that answer is the whole + // product of sending the frame. Until 2026-09-13 this returned right after + // the write: the result frame (findings, coverage, and any refusal) was never + // read, so no finding a node produced ever reached the proxy and nothing was + // filed — while delivery counters and node health both looked healthy. + // bugfix: workflow/CI/bugfix/20260913-async-scan-lane-never-returned-findings.md + _ = s.conn.SetReadDeadline(time.Now().Add(s.ackTimeout)) + res, err := deepscan.ReadResult(s.conn) + if err != nil { + _ = s.conn.Close() + s.conn = nil + return fmt.Errorf("scannode %s: read result: %w", s.node.ID, err) + } + if res.Reject != "" { + // A refusal is not a result. Reset the connection (the node may close + // its end) and let the caller decide, by code, whether another node may + // be tried. + _ = s.conn.Close() + s.conn = nil + return &deepscan.RejectError{Code: res.Reject} + } + select { + case s.results <- res: + default: + return fmt.Errorf("scannode %s: result buffer full, result for job %s dropped", s.node.ID, res.JobID) + } + return nil +} + +func (s *tlsSink) dial(ctx context.Context) (net.Conn, error) { + // (1) allowlist — before anything touches the network. + if !s.trusted(strings.ToLower(s.node.Fingerprint)) { + return nil, fmt.Errorf("scannode %s: fingerprint %s is not in the trusted set (a node self-reporting its own fingerprint is not an answer)", + s.node.ID, shortFP(s.node.Fingerprint)) + } + // (2) https and (3) private address. + if err := checkPrivateAddress(s.node.Addr); err != nil { + return nil, err + } + u, err := checkScheme(s.node.Addr) + if err != nil { + return nil, err + } + + want := strings.ToLower(strings.ReplaceAll(s.node.Fingerprint, ":", "")) + d := &net.Dialer{Timeout: s.dialTimeout} + // (4) certificate pin, enforced inside the handshake so a mismatch aborts + // before any application byte can be written. + // + // 🔴 InsecureSkipVerify is TRUE ON PURPOSE and is NOT a weakening here. + // This product has no CA, no issuance and no revocation anywhere in the tree + // (verified 2026-09-11), so chain verification could only ever fail; nodes + // carry installer-generated self-signed certs. The trust decision is moved + // wholesale into VerifyPeerCertificate below, which is STRICTER than a chain + // check: it accepts exactly one certificate, by its SHA-256, and nothing else. + // Removing InsecureSkipVerify without replacing this pin would not make the + // client safer — it would make it non-functional. + cfg := &tls.Config{ + InsecureSkipVerify: true, //nolint:gosec // replaced by the exact-fingerprint pin below; see the comment above + MinVersion: tls.VersionTLS12, + VerifyPeerCertificate: func(rawCerts [][]byte, _ [][]*x509Certificate) error { + if len(rawCerts) == 0 { + return fmt.Errorf("scannode %s: node presented no certificate", s.node.ID) + } + sum := sha256.Sum256(rawCerts[0]) + got := hex.EncodeToString(sum[:]) + if got != want { + return fmt.Errorf("scannode %s: certificate fingerprint %s does not match the pinned %s", + s.node.ID, shortFP(got), shortFP(want)) + } + return nil + }, + } + conn, err := tls.DialWithDialer(d, "tcp", u.Host, cfg) + if err != nil { + return nil, fmt.Errorf("scannode %s: tls dial %s: %w", s.node.ID, u.Host, err) + } + _ = ctx // dial deadline is carried by the dialer; ctx reserved for cancellation wiring + return conn, nil +} + +func (s *tlsSink) Close() error { + s.mu.Lock() + defer s.mu.Unlock() + if s.conn == nil { + return nil + } + err := s.conn.Close() + s.conn = nil + return err +} + +// shortFP renders a fingerprint for logs. Never log the full pin: it is not a +// secret, but a truncated form keeps log lines readable and makes an accidental +// copy-paste into a config obviously wrong. +func shortFP(fp string) string { + fp = strings.ToLower(fp) + if len(fp) <= 16 { + return fp + } + return fp[:16] + "…" +} diff --git a/scannode/trust.go b/scannode/trust.go new file mode 100644 index 0000000..5c77be0 --- /dev/null +++ b/scannode/trust.go @@ -0,0 +1,112 @@ +// Package scannode is the proxy's client side of the scan-node protocol: which +// nodes exist, which of them may be handed raw prompt bytes, and how those bytes +// get there. +// +// 🔴 EVERYTHING IN THIS PACKAGE GUARDS ONE SENTENCE: a piece of an employee's +// raw prompt must never reach a box we did not independently decide to trust. +// The four gates below are ANDed and every one of them is checked BEFORE the +// frame is written, because "returned an error after sending" is indistinguishable +// from success as far as disclosure goes. +package scannode + +import ( + "fmt" + "net" + "net/url" + "strings" +) + +// Node is one scan node the proxy may send work to. +type Node struct { + ID string + Addr string // https://: + // Fingerprint is the SHA-256 of the node's DER certificate, lowercase hex. + // It is the ONLY thing that authenticates the node: there is no CA in this + // product (verified across the whole tree, 2026-09-11), so the chain is not + // checked and this pin replaces it entirely. + Fingerprint string + Weight int +} + +// NodeSet is a resolved node list plus the decision of which fingerprints the +// control plane vouched for. +// +// Trusted is a FUNCTION and not a []string because the two editions answer the +// question from different sources — Production from an HTTPS member-auth +// response, Cluster from a list the installer rendered onto the box — and a node +// self-reporting its own fingerprint may never be an answer. Reason carries why +// the set is empty (no_nodes / insecure_control_plane / disabled) so the health +// surface can say something an operator can act on instead of "0 nodes". +type NodeSet struct { + Nodes []Node + Trusted func(fingerprint string) bool + Reason string +} + +// Reasons a NodeSet is empty or unusable. These strings surface verbatim in the +// proxy's health section (design §4b.3), so they are part of an external contract. +const ( + ReasonOK = "ok" + ReasonNoNodes = "no_nodes" + ReasonInsecureControlPlane = "insecure_control_plane" + ReasonLocalDaemonAbsent = "local_daemon_absent" + ReasonDisabled = "disabled" +) + +// checkScheme refuses anything but https. There is no plaintext mode and no +// opt-out flag: user decision D21 (2026-09-11) 「强制 TLS…永无明文模式」. A flag +// here would become the thing that is set "just for debugging" on the one +// machine that then stays that way. +func checkScheme(addr string) (*url.URL, error) { + u, err := url.Parse(addr) + if err != nil { + return nil, fmt.Errorf("scannode: unparsable node address %q: %w", addr, err) + } + if !strings.EqualFold(u.Scheme, "https") { + return nil, fmt.Errorf("scannode: node address %q is %q, not https — raw prompt content is never sent in the clear", addr, u.Scheme) + } + if u.Host == "" { + return nil, fmt.Errorf("scannode: node address %q has no host", addr) + } + return u, nil +} + +// checkPrivateAddress resolves the node's host and requires every resolved IP to +// be RFC1918, RFC4193 (fc00::/7), link-local or loopback (R-scan-node-deepscan-10). +// +// WHY resolve instead of pattern-matching the string: the guard exists to stop a +// tampered or fat-fingered node list from becoming an exfiltration channel, and +// a hostname is exactly how such a list would point outside without looking like +// it. EVERY resolved address must pass — a name that returns one private and one +// public address is refused, because we do not control which one the dialer picks. +func checkPrivateAddress(addr string) error { + u, err := checkScheme(addr) + if err != nil { + return err + } + host := u.Hostname() + ips := []net.IP{} + if ip := net.ParseIP(host); ip != nil { + ips = append(ips, ip) + } else { + resolved, rerr := net.LookupIP(host) + if rerr != nil { + return fmt.Errorf("scannode: cannot resolve node host %q: %w", host, rerr) + } + ips = resolved + } + if len(ips) == 0 { + return fmt.Errorf("scannode: node host %q resolved to no addresses", host) + } + for _, ip := range ips { + if !isPrivate(ip) { + return fmt.Errorf("scannode: node %q resolves to %s, which is not a private address — refusing to send raw content off the private network", addr, ip) + } + } + return nil +} + +// isPrivate covers RFC1918, RFC4193 unique-local, loopback and link-local. +func isPrivate(ip net.IP) bool { + return ip.IsLoopback() || ip.IsLinkLocalUnicast() || ip.IsPrivate() +} diff --git a/scannode/x509alias.go b/scannode/x509alias.go new file mode 100644 index 0000000..023da2c --- /dev/null +++ b/scannode/x509alias.go @@ -0,0 +1,7 @@ +package scannode + +import "crypto/x509" + +// x509Certificate keeps the VerifyPeerCertificate signature in sink_tls.go +// readable without importing crypto/x509 there for a type we never construct. +type x509Certificate = x509.Certificate diff --git a/scantoken/go.mod b/scantoken/go.mod new file mode 100644 index 0000000..e438ac5 --- /dev/null +++ b/scantoken/go.mod @@ -0,0 +1,3 @@ +module github.com/AiKeyLabs/pkg/scantoken + +go 1.26.1 diff --git a/scantoken/keyset_test.go b/scantoken/keyset_test.go new file mode 100644 index 0000000..ba62b05 --- /dev/null +++ b/scantoken/keyset_test.go @@ -0,0 +1,142 @@ +package scantoken + +import ( + "encoding/json" + "os" + "testing" + "time" +) + +type keysetFixture struct { + Cases []struct { + Name string `json:"name"` + Raw string `json:"raw"` + CurrentKeyID string `json:"current_key_id"` + CurrentSecret string `json:"current_secret"` + PreviousKeyID string `json:"previous_key_id"` + PreviousSecret string `json:"previous_secret"` + PreviousValidUntil int64 `json:"previous_valid_until"` + } `json:"cases"` + RejectCases []struct { + Name string `json:"name"` + Raw string `json:"raw"` + } `json:"reject_cases"` +} + +func loadKeysetFixture(t *testing.T) keysetFixture { + t.Helper() + raw, err := os.ReadFile("testdata/keyset_vectors.json") + if err != nil { + t.Fatalf("read fixture: %v", err) + } + var f keysetFixture + if err := json.Unmarshal(raw, &f); err != nil { + t.Fatalf("parse fixture: %v", err) + } + if len(f.Cases) == 0 || len(f.RejectCases) == 0 { + t.Fatal("fixture is empty — this test would assert nothing") + } + return f +} + +// TestParseKeySet_Fixture pins the file format against the SAME vectors the +// Python node parses (workers/scan_token.py). Two languages, one file on disk: +// if they ever disagree, the master mints with a key the node will not accept +// and every scan fails `unauthorized` with nothing to point at. +func TestParseKeySet_Fixture(t *testing.T) { + f := loadKeysetFixture(t) + for _, c := range f.Cases { + t.Run(c.Name, func(t *testing.T) { + ks, err := ParseKeySet([]byte(c.Raw)) + if err != nil { + t.Fatalf("ParseKeySet: %v", err) + } + if ks.Current.ID != c.CurrentKeyID { + t.Errorf("current id = %q, want %q", ks.Current.ID, c.CurrentKeyID) + } + if string(ks.Current.Secret) != c.CurrentSecret { + t.Errorf("current secret = %q, want %q", ks.Current.Secret, c.CurrentSecret) + } + if c.PreviousSecret == "" { + if ks.Previous != nil { + t.Errorf("previous = %+v, want none", ks.Previous) + } + return + } + if ks.Previous == nil { + t.Fatal("previous is nil, want a key") + } + if ks.Previous.ID != c.PreviousKeyID || string(ks.Previous.Secret) != c.PreviousSecret { + t.Errorf("previous = %s:%s, want %s:%s", + ks.Previous.ID, ks.Previous.Secret, c.PreviousKeyID, c.PreviousSecret) + } + if got := ks.PreviousValidUntil.Unix(); got != c.PreviousValidUntil { + t.Errorf("previous_valid_until = %d, want %d", got, c.PreviousValidUntil) + } + }) + } + for _, c := range f.RejectCases { + t.Run("reject/"+c.Name, func(t *testing.T) { + if _, err := ParseKeySet([]byte(c.Raw)); err == nil { + t.Fatal("ParseKeySet accepted a file it must refuse — " + + "downstream, 'no usable key' looks exactly like 'this deployment has no scan nodes'") + } + }) + } +} + +// TestFormatKeySet_RoundTrips is what makes the installers safe to write: they +// call FormatKeySet, every reader calls ParseKeySet, so the shape exists in one +// place. A round-trip that loses a field would strand tokens on rotation. +func TestFormatKeySet_RoundTrips(t *testing.T) { + until := time.Unix(1789200634, 0) + in := KeySet{ + Current: Key{ID: "k7", Secret: []byte("new-secret-value")}, + Previous: &Key{ID: "k6", Secret: []byte("old-secret-value")}, + PreviousValidUntil: until, + } + out, err := ParseKeySet([]byte(FormatKeySet(in))) + if err != nil { + t.Fatalf("ParseKeySet(FormatKeySet(...)): %v", err) + } + if out.Current.ID != in.Current.ID || string(out.Current.Secret) != string(in.Current.Secret) { + t.Errorf("current round-trip: %s:%s != %s:%s", + out.Current.ID, out.Current.Secret, in.Current.ID, in.Current.Secret) + } + if out.Previous == nil || + out.Previous.ID != in.Previous.ID || + string(out.Previous.Secret) != string(in.Previous.Secret) { + t.Errorf("previous round-trip: %+v != %+v", out.Previous, in.Previous) + } + if !out.PreviousValidUntil.Equal(until) { + t.Errorf("previous_valid_until round-trip: %v != %v", out.PreviousValidUntil, until) + } +} + +// TestParseKeySet_RotatedKeySetActuallyVerifiesBothKeys is the assertion that +// matters operationally: a token minted with the key that was just retired must +// still verify during the grace window, and must stop afterwards. Everything +// else here is about parsing; this is about the outage rotation exists to avoid. +func TestParseKeySet_RotatedKeySetActuallyVerifiesBothKeys(t *testing.T) { + now := time.Unix(1789200034, 0) + old := Key{ID: "k1", Secret: []byte("OLDSECRET-bbbbbbbbbbbbbbbbbbbbbb")} + inFlight := Mint(old, "org_a", now.Add(-30*time.Second)) + + ks, err := ParseKeySet([]byte( + "current: k2:NEWSECRET-aaaaaaaaaaaaaaaaaaaaaa\n" + + "previous: k1:OLDSECRET-bbbbbbbbbbbbbbbbbbbbbb\n" + + "previous_valid_until: " + "1789200634" + "\n")) + if err != nil { + t.Fatalf("ParseKeySet: %v", err) + } + + if org, err := Verify(inFlight, ks, now); err != nil || org != "org_a" { + t.Fatalf("a token minted seconds before the rotation was refused (org=%q err=%v) — "+ + "this is the ten-minute fleet-wide outage `previous` exists to prevent", org, err) + } + // One second past the grace instant the retired key is dead, even for a + // token whose own expiry has not been reached. + if _, err := Verify(inFlight, ks, time.Unix(1789200635, 0)); err == nil { + t.Fatal("the retired key still verified after previous_valid_until — rotation never completes") + } +} diff --git a/scantoken/scantoken.go b/scantoken/scantoken.go new file mode 100644 index 0000000..5d01010 --- /dev/null +++ b/scantoken/scantoken.go @@ -0,0 +1,301 @@ +// Package scantoken mints and verifies the short-lived, org-scoped bearer token +// a proxy presents to a scan node with every task frame. +// +// FORMAT (design §4b.7 — this string is a wire contract, mirrored in Python by +// ai-compliance-workers/workers/scan_token.py and pinned by testdata/vectors.json): +// +// sct1.... +// +// WHY STATELESS, AND WHAT THAT COSTS +// --------------------------------------------------------------------------- +// A node must be able to refuse a frame WITHOUT calling anyone: it sits on the +// data path, and a lookup per frame would put the control plane in the way of +// every scan. The price is that there is no revocation list — a leaked token +// stays usable until it expires. That is why the lifetime is 600s and why +// Verify refuses to be generous about the clock: the TTL *is* the revocation +// mechanism. Rotation (KeySet.Previous + PreviousValidUntil) exists so changing +// the signing key does not blackhole tokens already in flight. +// +// 🚫 The org id travels in PLAINTEXT inside the token. That is deliberate — the +// node needs it before it can verify anything — and it is exactly why every +// field is inside the MAC: editing the org id, the expiry or the key id must +// invalidate the token. Fence: TestScanToken_TamperedMacRejected. +package scantoken + +import ( + "crypto/hmac" + "crypto/sha256" + "encoding/base64" + "errors" + "fmt" + "strconv" + "strings" + "time" +) + +// Prefix is the format marker. A different prefix is a different format and is +// refused rather than best-effort parsed. +const Prefix = "sct1" + +// TTL is how long a minted token stays valid. Short on purpose: with no +// revocation list this is the only bound on a leaked token's usefulness. +const TTL = 600 * time.Second + +// SkewAllowance is how far past expiry a node still accepts a token, to absorb +// clock drift between an employee laptop and a scan node. Kept small for the +// same reason TTL is: every second here is a second a leaked token still works. +const SkewAllowance = 30 * time.Second + +// Errors callers may branch on. A node maps all of them to the single reject +// code `unauthorized`: telling a caller WHICH check failed is free help for +// someone probing the token format. +var ( + ErrMalformed = errors.New("scantoken: malformed token") + ErrUnknownKey = errors.New("scantoken: unknown key id") + ErrExpired = errors.New("scantoken: token expired") + ErrBadMAC = errors.New("scantoken: signature mismatch") + ErrGraceEnded = errors.New("scantoken: previous key is past its grace window") +) + +// Key is one signing key: a stable id (which travels in the token so a verifier +// knows which secret to use) and the secret itself. +type Key struct { + ID string + Secret []byte +} + +// KeySet is what a verifier holds: the key in use now, plus optionally the key +// it just rotated away from and the instant that key stops being accepted. +type KeySet struct { + Current Key + Previous *Key + PreviousValidUntil time.Time +} + +// Mint returns a token for orgID valid for TTL from now. +func Mint(key Key, orgID string, now time.Time) string { + exp := now.Add(TTL).Unix() + return assemble(key, orgID, exp) +} + +// MintUntil returns a token that expires at the given instant. Used by callers +// that must align a token's life with something else (a response's +// token_expires_at, a test vector); Mint is the normal entry point. +func MintUntil(key Key, orgID string, exp time.Time) string { + return assemble(key, orgID, exp.Unix()) +} + +func assemble(key Key, orgID string, exp int64) string { + expStr := strconv.FormatInt(exp, 10) + mac := sign(key.Secret, orgID, expStr, key.ID) + return strings.Join([]string{Prefix, orgID, expStr, key.ID, mac}, ".") +} + +// signingInput is the exact byte string covered by the MAC. Pipe-separated with +// no length prefixes is safe here ONLY because none of the four fields may +// contain a pipe or a dot — enforced by validateField below. Without that check +// ("a|b", "c") and ("a", "b|c") would sign identically. +func signingInput(orgID, expStr, keyID string) string { + return Prefix + "|" + orgID + "|" + expStr + "|" + keyID +} + +func sign(secret []byte, orgID, expStr, keyID string) string { + m := hmac.New(sha256.New, secret) + m.Write([]byte(signingInput(orgID, expStr, keyID))) + return base64.RawURLEncoding.EncodeToString(m.Sum(nil)) +} + +// Verify checks a token against keys at time now and returns the org it names. +// +// Order matters and is deliberate: shape → key lookup → MAC → expiry. The MAC +// is checked BEFORE the expiry so an attacker cannot learn anything by feeding +// tokens with edited expiry fields; and the grace window is checked against the +// KEY, not the token, so a rotated-away key stops working on schedule even for +// tokens minted seconds before the rotation. +func Verify(tok string, keys KeySet, now time.Time) (string, error) { + parts := strings.Split(tok, ".") + if len(parts) != 5 || parts[0] != Prefix { + return "", ErrMalformed + } + orgID, expStr, keyID, mac := parts[1], parts[2], parts[3], parts[4] + if orgID == "" || keyID == "" || mac == "" { + return "", ErrMalformed + } + exp, err := strconv.ParseInt(expStr, 10, 64) + if err != nil { + return "", ErrMalformed + } + + var secret []byte + switch { + case keyID == keys.Current.ID: + secret = keys.Current.Secret + case keys.Previous != nil && keyID == keys.Previous.ID: + if !keys.PreviousValidUntil.IsZero() && now.After(keys.PreviousValidUntil) { + return "", ErrGraceEnded + } + secret = keys.Previous.Secret + default: + return "", ErrUnknownKey + } + + want := sign(secret, orgID, expStr, keyID) + // hmac.Equal, not ==: string comparison short-circuits on the first differing + // byte and leaks how much of a guessed MAC was right. + if !hmac.Equal([]byte(want), []byte(mac)) { + return "", ErrBadMAC + } + if now.After(time.Unix(exp, 0).Add(SkewAllowance)) { + return "", fmt.Errorf("%w at %s (now %s)", ErrExpired, + time.Unix(exp, 0).UTC().Format(time.RFC3339), now.UTC().Format(time.RFC3339)) + } + return orgID, nil +} + +// ValidateField reports whether s may be used as an org id or key id. The MAC's +// signing input is pipe-separated and the token itself is dot-separated, so a +// value containing either character would make two different (org, exp, key) +// triples sign or parse identically. +func ValidateField(s string) error { + if s == "" { + return fmt.Errorf("scantoken: empty field") + } + if strings.ContainsAny(s, ".|") { + return fmt.Errorf("scantoken: %q contains '.' or '|', which the token format reserves as separators", s) + } + return nil +} + +// --------------------------------------------------------------------------- +// Key-file format +// --------------------------------------------------------------------------- +// +// ONE FILE, THREE READERS: the Production master (which mints), a Cluster proxy +// node (which also mints), and every scan node (which verifies). The installers +// write it; nobody edits it by hand. It lives at AIKEY_SCAN_TOKEN_KEY_FILE. +// +// current: k2: +// previous: k1: +// previous_valid_until: +// +// 🔴 WHY `previous` EXISTS AT ALL. The token is stateless with a 600s TTL and no +// revocation list. At the instant a rotation lands, tokens minted seconds ago +// with the old key are still in flight and still legitimately valid for up to +// 600 more seconds. A rotation that simply overwrote the secret would make every +// node answer `unauthorized` for ten minutes — indistinguishable, from the +// console, from the nodes being down. `previous_valid_until` is therefore +// now+TTL exactly: the last token mintable with the old key expires precisely +// then, so a longer window keeps a retired key alive for nobody and a shorter +// one strands tokens that were valid when they were issued. +// +// 🔴 LEGACY ONE-LINER IS STILL VALID and must stay that way: the first +// installers wrote a bare secret, then `:`. Those files are on +// deployed machines right now. Refusing them would turn an upgrade into an +// outage for every deployment that has not rotated yet, so a file with no +// `current:` line is read as the current key with no previous. +// +// Mirrored in Python by ai-compliance-workers/workers/scan_token.py +// (parse_keyset) and pinned across both by testdata/keyset_vectors.json. +const ( + fieldCurrent = "current" + fieldPrevious = "previous" + fieldPreviousValidUntil = "previous_valid_until" +) + +// ParseKeySet reads the key file's contents. +// +// Parsing is deliberately forgiving about blank lines and comments and +// deliberately STRICT about a malformed `current`: everything downstream — the +// endpoint that mints, the node that verifies — treats "no key" as "this +// deployment has no scan nodes", which is a silent, correct-looking state. A +// key file that exists but is unreadable is an operator mistake, and it must +// surface as an error rather than as a feature quietly switching itself off. +func ParseKeySet(raw []byte) (KeySet, error) { + var ks KeySet + var prevID, prevSecret string + sawCurrent := false + + for _, line := range strings.Split(string(raw), "\n") { + t := strings.TrimSpace(line) + if t == "" || strings.HasPrefix(t, "#") { + continue + } + name, value, ok := strings.Cut(t, ":") + if !ok { + continue + } + name = strings.TrimSpace(name) + value = strings.TrimSpace(value) + switch name { + case fieldCurrent: + id, secret := splitKeyValue(value) + if secret == "" { + return KeySet{}, fmt.Errorf("scantoken: key file has an empty `current` secret") + } + ks.Current = Key{ID: id, Secret: []byte(secret)} + sawCurrent = true + case fieldPrevious: + prevID, prevSecret = splitKeyValue(value) + case fieldPreviousValidUntil: + n, err := strconv.ParseInt(value, 10, 64) + if err != nil { + return KeySet{}, fmt.Errorf("scantoken: key file has an unparsable `previous_valid_until` %q: %w", value, err) + } + ks.PreviousValidUntil = time.Unix(n, 0) + } + } + + if !sawCurrent { + // Legacy form: the whole file is the key, either `:` or a + // bare secret. See the LEGACY note above — this path is load-bearing for + // every machine installed before rotation existed. + id, secret := splitKeyValue(strings.TrimSpace(string(raw))) + if secret == "" { + return KeySet{}, fmt.Errorf("scantoken: key file is empty") + } + ks.Current = Key{ID: id, Secret: []byte(secret)} + return ks, nil + } + + // A previous key with no grace instant would be accepted FOREVER, which is + // the opposite of what rotating is for. Drop it and say so, rather than + // quietly keeping a retired secret live. + if prevSecret != "" && !ks.PreviousValidUntil.IsZero() { + ks.Previous = &Key{ID: prevID, Secret: []byte(prevSecret)} + } else if prevSecret != "" { + return KeySet{}, fmt.Errorf("scantoken: key file has `previous` but no `previous_valid_until`, " + + "which would keep the retired key valid forever") + } + return ks, nil +} + +// splitKeyValue parses `:`, defaulting the id to DefaultKeyID +// when the value carries no id. A secret may itself contain ':' (base64 does +// not, but nothing stops an operator), so only the FIRST colon separates. +func splitKeyValue(v string) (id, secret string) { + v = strings.TrimSpace(v) + if v == "" { + return "", "" + } + if i, rest, ok := strings.Cut(v, ":"); ok && i != "" && strings.TrimSpace(rest) != "" { + return strings.TrimSpace(i), strings.TrimSpace(rest) + } + return DefaultKeyID, v +} + +// DefaultKeyID is the id assumed for a key file that names none. It is the id +// the very first installers minted with, so it must not change. +const DefaultKeyID = "k1" + +// FormatKeySet renders a KeySet back to the file format. The installers use it +// so the writer and the reader can never drift — the shape exists in exactly +// one place in Go, and is mirrored once in Python. +func FormatKeySet(ks KeySet) string { + var b strings.Builder + b.WriteString(fieldCurrent + ": " + ks.Current.ID + ":" + string(ks.Current.Secret) + "\n") + if ks.Previous != nil { + b.WriteString(fieldPrevious + ": " + ks.Previous.ID + ":" + string(ks.Previous.Secret) + "\n") + b.WriteString(fieldPreviousValidUntil + ": " + strconv.FormatInt(ks.PreviousValidUntil.Unix(), 10) + "\n") + } + return b.String() +} diff --git a/scantoken/scantoken_test.go b/scantoken/scantoken_test.go new file mode 100644 index 0000000..be472e9 --- /dev/null +++ b/scantoken/scantoken_test.go @@ -0,0 +1,199 @@ +package scantoken + +import ( + "encoding/json" + "os" + "path/filepath" + "strings" + "testing" + "time" +) + +var ( + keyA = Key{ID: "k1", Secret: []byte("0123456789abcdef0123456789abcdef")} + keyB = Key{ID: "k2", Secret: []byte("fedcba9876543210fedcba9876543210")} +) + +// TestScanToken_ExpiredRejected: the token is STATELESS — a node cannot ask +// anyone whether it was revoked, so the only thing standing between a leaked +// token and an attacker feeding a node content is the clock. 600s of validity +// plus a 30s backward skew allowance is the whole revocation story +// (design §4b.7), which makes this the fence that matters most in this file. +func TestScanToken_ExpiredRejected(t *testing.T) { + now := time.Unix(1_757_620_000, 0).UTC() + keys := KeySet{Current: keyA} + tok := Mint(keyA, "org_a", now) + + if org, err := Verify(tok, keys, now); err != nil || org != "org_a" { + t.Fatalf("fresh token must verify: org=%q err=%v", org, err) + } + // Still inside the 600s lifetime. + if _, err := Verify(tok, keys, now.Add(599*time.Second)); err != nil { + t.Errorf("token at t+599s must still verify: %v", err) + } + // Past expiry but inside the 30s skew allowance. + if _, err := Verify(tok, keys, now.Add(TTL+20*time.Second)); err != nil { + t.Errorf("token %v past expiry is inside the %v skew allowance and must verify: %v", + 20*time.Second, SkewAllowance, err) + } + // Past expiry AND past the skew allowance. + if _, err := Verify(tok, keys, now.Add(TTL+SkewAllowance+time.Second)); err == nil { + t.Error("an expired token was accepted — a leaked token would then be usable forever") + } +} + +// TestScanToken_PreviousKeyWithinGrace: rotation must not blackhole in-flight +// tokens. The installer writes a new key and the old one moves to `previous` +// with a valid_until 600s out; tokens minted just before rotation keep working +// until then, and not one second longer. +func TestScanToken_PreviousKeyWithinGrace(t *testing.T) { + now := time.Unix(1_757_620_000, 0).UTC() + old := Mint(keyA, "org_a", now) + + rotated := KeySet{ + Current: keyB, + Previous: &keyA, + PreviousValidUntil: now.Add(600 * time.Second), + } + if org, err := Verify(old, rotated, now.Add(time.Second)); err != nil || org != "org_a" { + t.Fatalf("token minted with the previous key must verify inside the grace: org=%q err=%v", org, err) + } + if _, err := Verify(old, rotated, now.Add(601*time.Second)); err == nil { + t.Error("the previous key was still accepted after PreviousValidUntil — the grace never ends") + } + // A token minted with the CURRENT key is unaffected by the grace window. + if _, err := Verify(Mint(keyB, "org_a", now), rotated, now.Add(time.Second)); err != nil { + t.Errorf("current-key token must verify: %v", err) + } + // An unknown key id is refused outright rather than falling back to current. + if _, err := Verify(Mint(Key{ID: "k9", Secret: keyA.Secret}, "org_a", now), rotated, now); err == nil { + t.Error("a token naming an unknown key id was accepted") + } +} + +// TestScanToken_TamperedMacRejected covers the case the format invites: the +// org id sits in PLAINTEXT in the token, so "just edit it" must fail the MAC. +func TestScanToken_TamperedMacRejected(t *testing.T) { + now := time.Unix(1_757_620_000, 0).UTC() + keys := KeySet{Current: keyA} + tok := Mint(keyA, "org_a", now) + + parts := strings.Split(tok, ".") + if len(parts) != 5 { + t.Fatalf("token shape changed: %q", tok) + } + for name, bad := range map[string]string{ + "org swapped": strings.Join([]string{parts[0], "org_b", parts[2], parts[3], parts[4]}, "."), + "expiry pushed": strings.Join([]string{parts[0], parts[1], "9999999999", parts[3], parts[4]}, "."), + "mac flipped": strings.Join([]string{parts[0], parts[1], parts[2], parts[3], flipLast(parts[4])}, "."), + "prefix swapped": strings.Join([]string{"sct2", parts[1], parts[2], parts[3], parts[4]}, "."), + "field dropped": strings.Join(parts[:4], "."), + } { + if _, err := Verify(bad, keys, now); err == nil { + t.Errorf("%s: tampered token was accepted (%q)", name, bad) + } + } +} + +// TestScanToken_VectorsStable pins fixed (key, org, exp) → token strings in +// testdata/vectors.json. +// +// WHY: ai-compliance-workers verifies these tokens in PYTHON (workers/scan_token.py). +// A Go-only test proves Go agrees with itself. These vectors are what stops the +// two implementations from drifting into "the node rejects every token and the +// proxy cannot tell why" — the failure mode that looks exactly like a network +// problem. Regenerate with UPDATE_VECTORS=1 only as a deliberate wire change. +func TestScanToken_VectorsStable(t *testing.T) { + path := filepath.Join("testdata", "vectors.json") + got := buildVectors() + + if os.Getenv("UPDATE_VECTORS") == "1" { + b, _ := json.MarshalIndent(got, "", " ") + if err := os.WriteFile(path, append(b, '\n'), 0o644); err != nil { + t.Fatalf("write vectors: %v", err) + } + t.Logf("vectors regenerated at %s", path) + return + } + raw, err := os.ReadFile(path) + if err != nil { + t.Fatalf("read vectors (regenerate with UPDATE_VECTORS=1): %v", err) + } + var want vectorFile + if err := json.Unmarshal(raw, &want); err != nil { + t.Fatalf("vectors.json is not JSON: %v", err) + } + if len(want.Vectors) != len(got.Vectors) { + t.Fatalf("vector count drifted: fixture %d, generator %d", len(want.Vectors), len(got.Vectors)) + } + for i := range got.Vectors { + if got.Vectors[i] != want.Vectors[i] { + t.Errorf("vector %d drifted — the Python mirror reads this file.\n got %+v\nwant %+v", + i, got.Vectors[i], want.Vectors[i]) + } + } + // The vectors must also VERIFY, so a regenerate can never bless a token that + // the library itself would reject. + for _, v := range want.Vectors { + keys := KeySet{Current: Key{ID: v.KeyID, Secret: []byte(v.SecretUTF8)}} + org, err := Verify(v.Token, keys, time.Unix(v.ExpUnix-1, 0).UTC()) + if err != nil || org != v.OrgID { + t.Errorf("fixture token does not verify: %+v org=%q err=%v", v, org, err) + } + } +} + +type vector struct { + KeyID string `json:"key_id"` + SecretUTF8 string `json:"secret_utf8"` + OrgID string `json:"org_id"` + ExpUnix int64 `json:"exp_unix"` + Token string `json:"token"` +} + +type vectorFile struct { + Note string `json:"note"` + TTLSecs int `json:"ttl_seconds"` + SkewSec int `json:"skew_allowance_seconds"` + Vectors []vector `json:"vectors"` +} + +func buildVectors() vectorFile { + cases := []struct { + key Key + org string + mint int64 + }{ + {keyA, "org_a", 1_757_620_000}, + {keyA, "org_with_underscores_and_digits_42", 1_700_000_000}, + {keyB, "org_b", 1_757_620_000}, + {Key{ID: "k-rotated-2026-09", Secret: []byte("a-32-byte-secret-with-dashes-01!")}, "org_c", 1_800_000_000}, + } + out := vectorFile{ + Note: "Cross-language fixture: Go pkg/scantoken and Python ai-compliance-workers/workers/scan_token.py must both produce and accept these byte-for-byte.", + TTLSecs: int(TTL / time.Second), + SkewSec: int(SkewAllowance / time.Second), + } + for _, c := range cases { + now := time.Unix(c.mint, 0).UTC() + out.Vectors = append(out.Vectors, vector{ + KeyID: c.key.ID, SecretUTF8: string(c.key.Secret), OrgID: c.org, + ExpUnix: c.mint + int64(TTL/time.Second), + Token: Mint(c.key, c.org, now), + }) + } + return out +} + +func flipLast(s string) string { + if s == "" { + return "x" + } + b := []byte(s) + if b[len(b)-1] == 'A' { + b[len(b)-1] = 'B' + } else { + b[len(b)-1] = 'A' + } + return string(b) +} diff --git a/scantoken/testdata/keyset_vectors.json b/scantoken/testdata/keyset_vectors.json new file mode 100644 index 0000000..9b5da81 --- /dev/null +++ b/scantoken/testdata/keyset_vectors.json @@ -0,0 +1,64 @@ +{ + "note": "Cross-language fixture for the AIKEY_SCAN_TOKEN_KEY_FILE format. Go pkg/scantoken.ParseKeySet and Python workers/scan_token.parse_keyset must agree on every case, including the two legacy shapes that are on deployed machines today.", + "cases": [ + { + "name": "rotated", + "raw": "current: k2:NEWSECRET-aaaaaaaaaaaaaaaaaaaaaa\nprevious: k1:OLDSECRET-bbbbbbbbbbbbbbbbbbbbbb\nprevious_valid_until: 1789200634\n", + "current_key_id": "k2", + "current_secret": "NEWSECRET-aaaaaaaaaaaaaaaaaaaaaa", + "previous_key_id": "k1", + "previous_secret": "OLDSECRET-bbbbbbbbbbbbbbbbbbbbbb", + "previous_valid_until": 1789200634 + }, + { + "name": "current only (rotated once, grace already reaped)", + "raw": "current: k3:ONLYSECRET-cccccccccccccccccccc\n", + "current_key_id": "k3", + "current_secret": "ONLYSECRET-cccccccccccccccccccc", + "previous_key_id": "", + "previous_secret": "", + "previous_valid_until": 0 + }, + { + "name": "legacy id:secret one-liner (pre-rotation installers)", + "raw": "k1:LEGACYSECRET-dddddddddddddddddd\n", + "current_key_id": "k1", + "current_secret": "LEGACYSECRET-dddddddddddddddddd", + "previous_key_id": "", + "previous_secret": "", + "previous_valid_until": 0 + }, + { + "name": "legacy bare secret (first installers)", + "raw": "BARESECRET-eeeeeeeeeeeeeeeeeeeeee\n", + "current_key_id": "k1", + "current_secret": "BARESECRET-eeeeeeeeeeeeeeeeeeeeee", + "previous_key_id": "", + "previous_secret": "", + "previous_valid_until": 0 + }, + { + "name": "comments and blank lines are ignored", + "raw": "# written by cluster-install.sh\n\ncurrent: k9:HASHSECRET-ffffffffffffffffff\n\n", + "current_key_id": "k9", + "current_secret": "HASHSECRET-ffffffffffffffffff", + "previous_key_id": "", + "previous_secret": "", + "previous_valid_until": 0 + } + ], + "reject_cases": [ + { + "name": "previous with no grace instant would live forever", + "raw": "current: k2:AAA\nprevious: k1:BBB\n" + }, + { + "name": "empty file", + "raw": "\n\n" + }, + { + "name": "unparsable previous_valid_until", + "raw": "current: k2:AAA\nprevious: k1:BBB\nprevious_valid_until: soon\n" + } + ] +} diff --git a/scantoken/testdata/vectors.json b/scantoken/testdata/vectors.json new file mode 100644 index 0000000..0ee0859 --- /dev/null +++ b/scantoken/testdata/vectors.json @@ -0,0 +1,35 @@ +{ + "note": "Cross-language fixture: Go pkg/scantoken and Python ai-compliance-workers/workers/scan_token.py must both produce and accept these byte-for-byte.", + "ttl_seconds": 600, + "skew_allowance_seconds": 30, + "vectors": [ + { + "key_id": "k1", + "secret_utf8": "0123456789abcdef0123456789abcdef", + "org_id": "org_a", + "exp_unix": 1757620600, + "token": "sct1.org_a.1757620600.k1.SfhPpi5rTc_ut6TTB980x31VMRHVjgo-9hKLj7r53rA" + }, + { + "key_id": "k1", + "secret_utf8": "0123456789abcdef0123456789abcdef", + "org_id": "org_with_underscores_and_digits_42", + "exp_unix": 1700000600, + "token": "sct1.org_with_underscores_and_digits_42.1700000600.k1.x6EjwBatBx_BAw4oBkX9RMFNbytceF_IIlZADKMBqls" + }, + { + "key_id": "k2", + "secret_utf8": "fedcba9876543210fedcba9876543210", + "org_id": "org_b", + "exp_unix": 1757620600, + "token": "sct1.org_b.1757620600.k2.cWJHOTe-fDEygNXm5Py8JSphHK0y-ETDoPyEp3O5vkc" + }, + { + "key_id": "k-rotated-2026-09", + "secret_utf8": "a-32-byte-secret-with-dashes-01!", + "org_id": "org_c", + "exp_unix": 1800000600, + "token": "sct1.org_c.1800000600.k-rotated-2026-09.F46Ufeahq3-0PElDRDlpwG6OsfDTQ_gU0aC0s33XyLI" + } + ] +} diff --git a/seatassign/seatassign_session_dist_test.go b/seatassign/seatassign_session_dist_test.go new file mode 100644 index 0000000..4c1bf70 --- /dev/null +++ b/seatassign/seatassign_session_dist_test.go @@ -0,0 +1,62 @@ +// SPIKE — openspec/changes/add-scan-node-deepscan tasks.md 1.5 +// +// Question: the scan-node design picks a node with weighted rendezvous hashing +// over the key `tenant_id‖content_sha256` (design §4b.2). seatassign.Primary is +// the chosen implementation. Does that key shape actually spread evenly over 4 +// equal-weight nodes, or does the shared tenant prefix bias one node? +// +// This is a measurement, not a fence: it prints the distribution so the number +// can be written into baseline-forensics.md §F5. +package seatassign + +import ( + "crypto/sha256" + "encoding/hex" + "fmt" + "math" + "testing" +) + +func TestSpikeContentKeyDistributionOverFourNodes(t *testing.T) { + nodes := []Account{ + {AccountID: "scan-1", Weight: 1}, + {AccountID: "scan-2", Weight: 1}, + {AccountID: "scan-3", Weight: 1}, + {AccountID: "scan-4", Weight: 1}, + } + const ( + keys = 10000 + tenant = "org_bosera_prod" + ) + counts := map[string]int{} + for i := 0; i < keys; i++ { + sum := sha256.Sum256([]byte(fmt.Sprintf("piece-%d", i))) + key := tenant + "‖" + hex.EncodeToString(sum[:]) + counts[Primary(key, nodes)]++ + } + + ideal := float64(keys) / float64(len(nodes)) + worst := 0.0 + for _, n := range nodes { + c := counts[n.AccountID] + dev := (float64(c) - ideal) / ideal * 100 + if math.Abs(dev) > math.Abs(worst) { + worst = dev + } + t.Logf("MEASURED %s: %d keys (%.2f%%), deviation from even share %+.2f%%", + n.AccountID, c, float64(c)/float64(keys)*100, dev) + } + t.Logf("MEASURED worst-node deviation over %d keys x %d nodes: %+.2f%%", keys, len(nodes), worst) + + // Same content must always pick the same node — the affinity the design + // relies on for "the node's own record becomes the de-facto shared layer". + sum := sha256.Sum256([]byte("stable-piece")) + key := tenant + "‖" + hex.EncodeToString(sum[:]) + first := Primary(key, nodes) + for i := 0; i < 100; i++ { + if got := Primary(key, nodes); got != first { + t.Fatalf("same key picked %s then %s — rendezvous is not deterministic", first, got) + } + } + t.Logf("MEASURED same tenant+content key -> %s on 100/100 draws", first) +}