diff --git a/agent/capture_approval_test.go b/agent/capture_approval_test.go index 706fc2b..ab0c83b 100644 --- a/agent/capture_approval_test.go +++ b/agent/capture_approval_test.go @@ -70,7 +70,7 @@ func TestCaptureApprovalLifecycle(t *testing.T) { t.Error(err) return } - if err := ic.Decide(detail.ID, tc.approve); err != nil { + if err := ic.Decide(detail.ID, tc.approve, "", ""); err != nil { t.Error(err) } }) diff --git a/db/conversation.go b/db/conversation.go index 7db9afd..7804891 100644 --- a/db/conversation.go +++ b/db/conversation.go @@ -58,20 +58,24 @@ type Conversation struct { PinnedAt *time.Time `json:"pinned_at,omitempty"` CreatedAt time.Time `json:"created_at"` UpdatedAt time.Time `json:"updated_at"` + // TargetGate is the per-conversation opt-in of the approval gate for asset + // targets (a chat has no task). false (default) = no gate. + TargetGate bool `json:"target_gate"` } // ConversationPatch updates only the fields whose pointers are non-nil. type ConversationPatch struct { - Title *string - Pinned *bool + Title *string + Pinned *bool + TargetGate *bool } -const convCols = `id, agent_key, title, llm_profile_id, pinned_at, created_at, updated_at` +const convCols = `id, agent_key, title, llm_profile_id, pinned_at, created_at, updated_at, target_gate` func scanConv(row interface{ Scan(...any) error }) (Conversation, error) { var c Conversation var pinnedAt sql.NullTime - err := row.Scan(&c.ID, &c.AgentKey, &c.Title, &c.LLMProfileID, &pinnedAt, &c.CreatedAt, &c.UpdatedAt) + err := row.Scan(&c.ID, &c.AgentKey, &c.Title, &c.LLMProfileID, &pinnedAt, &c.CreatedAt, &c.UpdatedAt, &c.TargetGate) if pinnedAt.Valid { c.Pinned = true c.PinnedAt = &pinnedAt.Time @@ -82,6 +86,12 @@ func scanConv(row interface{ Scan(...any) error }) (Conversation, error) { // CreateConversation opens a new chat thread for agentKey with an initial title. // llmProfileID may be nil to use the globally active profile. func (d *DB) CreateConversation(agentKey, title string, llmProfileID *int64) (*Conversation, error) { + return d.CreateConversationGated(agentKey, title, llmProfileID, false) +} + +// CreateConversationGated is CreateConversation with the target-gate opt-in set +// at creation. +func (d *DB) CreateConversationGated(agentKey, title string, llmProfileID *int64, targetGate bool) (*Conversation, error) { tx, err := d.Begin() if err != nil { return nil, err @@ -92,8 +102,8 @@ func (d *DB) CreateConversation(agentKey, title string, llmProfileID *int64) (*C // by this transaction before it takes a profile lock, matching DeleteProfile's // child-row -> profile-row protocol. c, err := scanConv(tx.QueryRow(` -INSERT INTO conversations(agent_key, title, llm_profile_id) VALUES ($1, $2, NULL) -RETURNING `+convCols, agentKey, title)) +INSERT INTO conversations(agent_key, title, llm_profile_id, target_gate) VALUES ($1, $2, NULL, $3) +RETURNING `+convCols, agentKey, title, targetGate)) if err != nil { return nil, err } @@ -169,7 +179,7 @@ func (d *DB) GetConversation(id int64) (*Conversation, error) { return &c, nil } -// UpdateConversation applies a partial title/pin mutation and returns the updated +// UpdateConversation applies a partial title/pin/target-gate mutation and returns the updated // row. Pinning an already-pinned conversation preserves its original pin order. func (d *DB) UpdateConversation(id int64, patch ConversationPatch) (*Conversation, error) { c, err := scanConv(d.QueryRow(`UPDATE conversations SET @@ -178,9 +188,10 @@ func (d *DB) UpdateConversation(id int64, patch ConversationPatch) (*Conversatio WHEN $4::boolean IS NULL THEN pinned_at WHEN $4::boolean THEN COALESCE(pinned_at, now()) ELSE NULL - END + END, + target_gate = CASE WHEN $5::boolean THEN $6 ELSE target_gate END WHERE id=$1 -RETURNING `+convCols, id, patch.Title != nil, patch.Title, patch.Pinned)) +RETURNING `+convCols, id, patch.Title != nil, patch.Title, patch.Pinned, patch.TargetGate != nil, patch.TargetGate)) if err == sql.ErrNoRows { return nil, nil } @@ -203,19 +214,37 @@ func (d *DB) TouchConversation(id int64) error { return err } -// DeleteConversation removes a thread; its activities cascade via FK. +// DeleteConversation removes a thread; its activities cascade via FK. Its +// target pins are not rows of the thread: they are revoked in the same +// transaction and stay as history. func (d *DB) DeleteConversation(id int64) error { - _, err := d.Exec(`DELETE FROM conversations WHERE id=$1`, id) - return err + tx, err := d.Begin() + if err != nil { + return err + } + defer tx.Rollback() + if err := revokeConversationPins(tx, []int64{id}); err != nil { + return err + } + if _, err := tx.Exec(`DELETE FROM conversations WHERE id=$1`, id); err != nil { + return err + } + return tx.Commit() } // DeleteConversations removes existing threads in one statement and returns the -// ids that were actually present. Child activities and trigger runs cascade. +// ids that were actually present. Child activities and trigger runs cascade. The +// target pins of the removed threads are revoked in the same transaction. func (d *DB) DeleteConversations(ids []int64) ([]int64, error) { if len(ids) == 0 { return []int64{}, nil } - rows, err := d.Query(`DELETE FROM conversations WHERE id=ANY($1::bigint[]) RETURNING id`, ids) + tx, err := d.Begin() + if err != nil { + return nil, err + } + defer tx.Rollback() + rows, err := tx.Query(`DELETE FROM conversations WHERE id=ANY($1::bigint[]) RETURNING id`, ids) if err != nil { return nil, err } @@ -228,7 +257,14 @@ func (d *DB) DeleteConversations(ids []int64) ([]int64, error) { } deleted = append(deleted, id) } - return deleted, rows.Err() + if err := rows.Err(); err != nil { + return nil, err + } + rows.Close() + if err := revokeConversationPins(tx, deleted); err != nil { + return nil, err + } + return deleted, tx.Commit() } // AppendConvActivity records one step of a conversation (human message or an agent diff --git a/db/finding_retests.go b/db/finding_retests.go index 16e32cd..3d3600c 100644 --- a/db/finding_retests.go +++ b/db/finding_retests.go @@ -114,8 +114,11 @@ func (d *DB) CreateFindingRetest(ctx context.Context, findingID int64, notes str if runes := []rune(title); len(runes) > 100 { title = string(runes[:100]) } - c, err := scanConv(tx.QueryRowContext(ctx, `INSERT INTO conversations(agent_key,title) VALUES ($1,$2) RETURNING `+convCols, - FindingRetestAgentKey, fmt.Sprintf("Retest #%d · %s", findingID, title))) + // The retest conversation takes the target-gate opt-in of the finding's task. + // Pins are never copied: a pin belongs to the scope it was approved in. + c, err := scanConv(tx.QueryRowContext(ctx, `INSERT INTO conversations(agent_key,title,target_gate) VALUES ($1,$2, + COALESCE((SELECT t.target_gate FROM findings f JOIN tasks t ON t.id=f.task_id WHERE f.id=$3), false)) RETURNING `+convCols, + FindingRetestAgentKey, fmt.Sprintf("Retest #%d · %s", findingID, title), findingID)) if err != nil { return nil, nil, false, err } diff --git a/db/intercept_detail.go b/db/intercept_detail.go index 135270b..3bf7366 100644 --- a/db/intercept_detail.go +++ b/db/intercept_detail.go @@ -43,6 +43,35 @@ type InterceptAudit struct { Output string `json:"output,omitempty"` OutputTruncated bool `json:"output_truncated,omitempty"` ExecutionEndedAt *time.Time `json:"execution_ended_at,omitempty"` + // Targets is set on target approval requests only (target_key IS NOT NULL). + Targets []InterceptTarget `json:"targets,omitempty"` + // MergedAsk marks a target request that also carries the ask of a rule or of + // the judge: one row, one decision for both. + MergedAsk bool `json:"merged_ask,omitempty"` + // Approver and ApproverAddr identify the human behind a decision: the subject + // of the session token and the remote address of the request. decided_by keeps + // the actor class (human). + Approver string `json:"approver,omitempty"` + ApproverAddr string `json:"approver_addr,omitempty"` +} + +// InterceptTarget is one asset target of a target approval request: what was +// asked for, the addresses shown to the human, and the pin it relates to. +type InterceptTarget struct { + Host string `json:"host"` + Port int `json:"port"` // 0 = port not declared + Kind string `json:"kind,omitempty"` + ResolvedIPs []string `json:"resolved_ips"` + Resolver string `json:"resolver,omitempty"` + ResolvedAt time.Time `json:"resolved_at"` + PinID int64 `json:"pin_id,omitempty"` + PinVersion int `json:"pin_version,omitempty"` + PinnedIPs []string `json:"pinned_ips,omitempty"` + Reason string `json:"reason,omitempty"` // no_pin | new_ip | revoked + Outcome string `json:"outcome"` // requested | pinned | denied + // NameOnlyEffective: requested while the upstream egress proxy was set; a pin + // created from this request covers the name only. + NameOnlyEffective bool `json:"name_only_effective,omitempty"` } type InterceptDetail struct { @@ -77,6 +106,14 @@ func (d *DB) GetInterceptDetail(id int64) (*InterceptDetail, error) { // decidedBy records the settling actor (human | rule | model | system); it is // supplied by the caller because status alone cannot tell them apart. func (d *DB) ResolveIntercept(id int64, status, action, reason, decidedBy string) (bool, error) { + return d.ResolveInterceptBy(id, status, action, reason, decidedBy, "", "") +} + +// ResolveInterceptBy is ResolveIntercept for a decision whose author is known: +// approver (subject of the session token) and remoteAddr (address of the +// request) are recorded in the row's audit when non-empty. A row without an +// audit keeps none. +func (d *DB) ResolveInterceptBy(id int64, status, action, reason, decidedBy, approver, remoteAddr string) (bool, error) { switch decidedBy { case "human", "rule", "model", "system": default: @@ -86,9 +123,16 @@ func (d *DB) ResolveIntercept(id int64, status, action, reason, decidedBy string if action == "allow" { execution = "awaiting_result" } - patch, err := json.Marshal(map[string]any{ + fields := map[string]any{ "effective_action": action, "decision_reason": reason, "execution_status": execution, - }) + } + if approver != "" { + fields["approver"] = approver + } + if remoteAddr != "" { + fields["approver_addr"] = remoteAddr + } + patch, err := json.Marshal(fields) if err != nil { return false, err } diff --git a/db/schema.sql b/db/schema.sql index d849922..45ece30 100644 --- a/db/schema.sql +++ b/db/schema.sql @@ -558,6 +558,8 @@ CREATE TRIGGER trg_tasks_upd BEFORE UPDATE ON tasks -- planner heartbeat trigger interval (seconds); backfills old databases. Default 300s (5min). See docs/planner-trigger-impl-plan.md ALTER TABLE tasks ADD COLUMN IF NOT EXISTS plan_heartbeat_seconds INTEGER NOT NULL DEFAULT 300; ALTER TABLE tasks ADD COLUMN IF NOT EXISTS plan_quiet_mode BOOLEAN NOT NULL DEFAULT true; +-- Approval gate for asset targets, opt-in per task; backfills old databases. false (default) = no gate, behaviour unchanged. +ALTER TABLE tasks ADD COLUMN IF NOT EXISTS target_gate BOOLEAN NOT NULL DEFAULT false; ALTER TABLE tasks ADD COLUMN IF NOT EXISTS llm_rotate BOOLEAN NOT NULL DEFAULT false; -- Concurrency-limit hold state; backfills old databases. true=queued due to the concurrency limit, auto-starts when a slot frees up. ALTER TABLE tasks ADD COLUMN IF NOT EXISTS queued BOOLEAN NOT NULL DEFAULT false; @@ -995,6 +997,8 @@ CREATE TABLE IF NOT EXISTS conversations ( updated_at TIMESTAMPTZ NOT NULL DEFAULT now() ); ALTER TABLE conversations ADD COLUMN IF NOT EXISTS pinned_at TIMESTAMPTZ; +-- Approval gate for asset targets, opt-in per conversation (a chat has no task). false (default) = no gate. +ALTER TABLE conversations ADD COLUMN IF NOT EXISTS target_gate BOOLEAN NOT NULL DEFAULT false; CREATE INDEX IF NOT EXISTS idx_conversations_llm_profile ON conversations(llm_profile_id) WHERE llm_profile_id IS NOT NULL; CREATE INDEX IF NOT EXISTS idx_conversations_pinned ON conversations(pinned_at DESC) WHERE pinned_at IS NOT NULL; DROP TRIGGER IF EXISTS trg_conversations_upd ON conversations; @@ -1119,6 +1123,55 @@ UPDATE intercept_pending SET decision_source='model' -- existed) | human | rule | model | system. No backfill: readers fall back to -- decision_source for legacy rows instead of rewriting history. ALTER TABLE intercept_pending ADD COLUMN IF NOT EXISTS decided_by TEXT NOT NULL DEFAULT ''; +-- Target approval requests. target_key = ':' of the asset target the row asks for; +-- NULL = classic ask (and every row that predates the column). At most one open request per +-- scope (task or conversation) and target: concurrent calls to one host share one row. +ALTER TABLE intercept_pending ADD COLUMN IF NOT EXISTS target_key TEXT; +CREATE UNIQUE INDEX IF NOT EXISTS uq_intercept_pending_target + ON intercept_pending (COALESCE(task_id,''), COALESCE(conversation_id,0), target_key) + WHERE status = 'pending' AND target_key IS NOT NULL; + +-- Target pins: host + port approved by a human for one task or one conversation, with the +-- addresses shown at approval time. Written only by the decide endpoint. No foreign keys on +-- task_id, conversation_id and intercept_id: a pin outlives its task, its conversation and its +-- request row as history (revoked, never deleted). status and revoked_reason carry no CHECK +-- (validated in Go, same convention as findings.status). +CREATE TABLE IF NOT EXISTS target_pins ( + id BIGSERIAL PRIMARY KEY, + task_id BIGINT, + conversation_id BIGINT, + host TEXT NOT NULL, -- normalised name, IP literal or CIDR + port INTEGER NOT NULL DEFAULT 0, -- 0 = port not declared: covers every port of the host + version INTEGER NOT NULL DEFAULT 1, + pinned_ips TEXT[] NOT NULL DEFAULT '{}', -- canonical, sorted + resolver TEXT NOT NULL DEFAULT '', + resolved_at TIMESTAMPTZ NOT NULL DEFAULT now(), + status TEXT NOT NULL DEFAULT 'active', -- active | superseded | revoked + intercept_id BIGINT, -- the request row that carried the click + decided_by TEXT NOT NULL DEFAULT '', -- approver: subject of the session token + decided_from TEXT NOT NULL DEFAULT '', -- remote address of the click + -- Approved while the upstream egress proxy was set: the proxy chooses the address, so the + -- pin covers the name only and its addresses are not compared. Stays true if the proxy is + -- removed later; enforcing addresses again takes a new approval. + name_only_effective BOOLEAN NOT NULL DEFAULT false, + revoked_at TIMESTAMPTZ, + revoked_reason TEXT NOT NULL DEFAULT '', -- operator | task_deleted | conversation_deleted + created_at TIMESTAMPTZ NOT NULL DEFAULT now() +); +CREATE UNIQUE INDEX IF NOT EXISTS uq_target_pins_active + ON target_pins (COALESCE(task_id,0), COALESCE(conversation_id,0), host, port) + WHERE status = 'active'; +CREATE INDEX IF NOT EXISTS idx_target_pins_scope ON target_pins (task_id, conversation_id, status); + +-- Append-only history of the target-gate exemption list (settings key target_gate_exempt_hosts). +CREATE TABLE IF NOT EXISTS target_gate_exemption_log ( + id BIGSERIAL PRIMARY KEY, + changed_at TIMESTAMPTZ NOT NULL DEFAULT now(), + changed_by TEXT NOT NULL DEFAULT '', + before JSONB, + after JSONB, + digest TEXT NOT NULL DEFAULT '' +); -- ===================================================================== -- L. Finding persistence diff --git a/db/target_pins.go b/db/target_pins.go new file mode 100644 index 0000000..d11644b --- /dev/null +++ b/db/target_pins.go @@ -0,0 +1,393 @@ +package db + +import ( + "database/sql" + "encoding/json" + "errors" + "fmt" + "sort" + "strconv" + "time" +) + +// TargetScopeRef names what a target pin or a target approval request is bound +// to: one task or one conversation, never both. +type TargetScopeRef struct { + TaskID int64 + ConversationID int64 +} + +func (s TargetScopeRef) check() error { + if (s.TaskID > 0) == (s.ConversationID > 0) { + return fmt.Errorf("target scope must name one task or one conversation, got task %d and conversation %d", s.TaskID, s.ConversationID) + } + return nil +} + +// pendingScope is the scope as intercept_pending stores it: task_id is TEXT +// there, and the column of the other kind is NULL. +func (s TargetScopeRef) pendingScope() (taskID *string, convID *int64) { + if s.TaskID > 0 { + v := strconv.FormatInt(s.TaskID, 10) + return &v, nil + } + return nil, &s.ConversationID +} + +// TargetKey is the key of an asset target inside its scope: host and port. The +// port is always the part after the last colon, so the key stays unambiguous +// for IPv6 literals and CIDRs. +func TargetKey(host string, port int) string { + return host + ":" + strconv.Itoa(port) +} + +// Target pin states and revocation reasons. The columns carry no CHECK; these +// are the only values written. +const ( + TargetPinActive = "active" + TargetPinSuperseded = "superseded" + TargetPinRevoked = "revoked" + + TargetPinRevokedByOperator = "operator" + TargetPinRevokedTaskDeleted = "task_deleted" + TargetPinRevokedConversationDelete = "conversation_deleted" +) + +// TargetPin is one row of target_pins: a host and port a human approved for one +// scope, with the addresses that were shown at approval time. Port 0 means the +// port was not declared and covers every port of the host. +type TargetPin struct { + ID int64 `json:"id"` + TaskID *int64 `json:"task_id"` + ConversationID *int64 `json:"conversation_id"` + Host string `json:"host"` + Port int `json:"port"` + Version int `json:"version"` + PinnedIPs []string `json:"pinned_ips"` + Resolver string `json:"resolver"` + ResolvedAt time.Time `json:"resolved_at"` + Status string `json:"status"` + InterceptID *int64 `json:"intercept_id"` + DecidedBy string `json:"decided_by"` + DecidedFrom string `json:"decided_from"` + // NameOnlyEffective: approved while the upstream egress proxy was set, so the + // pin covers the name and its addresses are not compared. It stays set when + // the proxy is removed later. + NameOnlyEffective bool `json:"name_only_effective"` + RevokedAt *time.Time `json:"revoked_at,omitempty"` + RevokedReason string `json:"revoked_reason,omitempty"` + CreatedAt time.Time `json:"created_at"` +} + +const targetPinCols = `id, task_id, conversation_id, host, port, version, array_to_json(pinned_ips)::text, resolver, resolved_at, status, intercept_id, decided_by, decided_from, name_only_effective, revoked_at, revoked_reason, created_at` + +// targetPinScope matches the expressions of uq_target_pins_active. +const targetPinScope = `COALESCE(task_id,0)=$1 AND COALESCE(conversation_id,0)=$2` + +func scanTargetPin(row interface{ Scan(...any) error }) (*TargetPin, error) { + var p TargetPin + var ips string + if err := row.Scan(&p.ID, &p.TaskID, &p.ConversationID, &p.Host, &p.Port, &p.Version, &ips, &p.Resolver, &p.ResolvedAt, + &p.Status, &p.InterceptID, &p.DecidedBy, &p.DecidedFrom, &p.NameOnlyEffective, &p.RevokedAt, &p.RevokedReason, &p.CreatedAt); err != nil { + return nil, err + } + if err := json.Unmarshal([]byte(ips), &p.PinnedIPs); err != nil { + return nil, err + } + return &p, nil +} + +func queryTargetPins(q interface { + Query(string, ...any) (*sql.Rows, error) +}, query string, args ...any) ([]*TargetPin, error) { + rows, err := q.Query(query, args...) + if err != nil { + return nil, err + } + defer rows.Close() + out := []*TargetPin{} + for rows.Next() { + p, err := scanTargetPin(rows) + if err != nil { + return nil, err + } + out = append(out, p) + } + return out, rows.Err() +} + +// ActiveTargetPins returns the active pins that may cover the target +// (host, port) in scope, the pin on the exact port first. A target on a declared +// port has two candidates: the pin on that port and the pin on port 0. A target +// whose port is not declared (0) has one: a pin on a specific port never covers +// it. The address sets of two pins are never merged by the caller. +func (d *DB) ActiveTargetPins(scope TargetScopeRef, host string, port int) ([]*TargetPin, error) { + if err := scope.check(); err != nil { + return nil, err + } + return queryTargetPins(d, `SELECT `+targetPinCols+` FROM target_pins +WHERE `+targetPinScope+` AND host=$3 AND port IN ($4, 0) AND status='active' +ORDER BY port DESC`, scope.TaskID, scope.ConversationID, host, port) +} + +// ListTargetPins returns every pin of scope in any state, newest first. +func (d *DB) ListTargetPins(scope TargetScopeRef) ([]*TargetPin, error) { + if err := scope.check(); err != nil { + return nil, err + } + return queryTargetPins(d, `SELECT `+targetPinCols+` FROM target_pins +WHERE `+targetPinScope+` ORDER BY id DESC`, scope.TaskID, scope.ConversationID) +} + +// RevokeTargetPin revokes one active pin of scope on behalf of the operator and +// returns the pin as it is afterwards (nil when scope has no such pin). Revoking +// a pin that is no longer active changes nothing. +func (d *DB) RevokeTargetPin(scope TargetScopeRef, pinID int64) (*TargetPin, error) { + if err := scope.check(); err != nil { + return nil, err + } + if _, err := d.Exec(`UPDATE target_pins SET status='revoked', revoked_at=now(), revoked_reason=$4 +WHERE `+targetPinScope+` AND id=$3 AND status='active'`, scope.TaskID, scope.ConversationID, pinID, TargetPinRevokedByOperator); err != nil { + return nil, err + } + p, err := scanTargetPin(d.QueryRow(`SELECT `+targetPinCols+` FROM target_pins +WHERE `+targetPinScope+` AND id=$3`, scope.TaskID, scope.ConversationID, pinID)) + if errors.Is(err, sql.ErrNoRows) { + return nil, nil + } + return p, err +} + +// revokeTaskPins marks the active pins of a task that is being deleted as +// revoked. They stay as history; task ids are never reused, so a revoked pin +// can never match again. +func revokeTaskPins(tx *sql.Tx, taskID int64) error { + _, err := tx.Exec(`UPDATE target_pins SET status='revoked', revoked_at=now(), revoked_reason=$2 +WHERE task_id=$1 AND status='active'`, taskID, TargetPinRevokedTaskDeleted) + return err +} + +// revokeConversationPins is revokeTaskPins for deleted conversations. +func revokeConversationPins(tx *sql.Tx, convIDs []int64) error { + _, err := tx.Exec(`UPDATE target_pins SET status='revoked', revoked_at=now(), revoked_reason=$2 +WHERE conversation_id=ANY($1::bigint[]) AND status='active'`, convIDs, TargetPinRevokedConversationDelete) + return err +} + +// settleTaskTargetRequests closes the open target requests of a task that is +// being deleted: denied by the system, nothing pinned. Only target rows are +// touched. (A conversation needs no counterpart: its rows cascade.) +func settleTaskTargetRequests(tx *sql.Tx, taskID int64) error { + _, err := tx.Exec(`UPDATE intercept_pending SET status='denied', decided_at=NOW(), decided_by='system', + audit=CASE WHEN audit IS NULL THEN NULL ELSE audit || + '{"effective_action":"deny","decision_reason":"task deleted","execution_status":"not_executed"}'::jsonb END +WHERE COALESCE(task_id,'')=$1 AND status='pending' AND target_key IS NOT NULL`, strconv.FormatInt(taskID, 10)) + return err +} + +// TargetRequest is a new target approval request: the intercept_pending row a +// gated call asks for when one of its targets has no pin that covers it. +type TargetRequest struct { + Scope TargetScopeRef + Host string + Port int + RuleID int64 // owning rule of a merged ask; 0 otherwise + AgentName string + ToolName string + ToolInput []byte + Reason string + Audit *InterceptAudit +} + +// FindOrCreateTargetRequest returns the open request for (scope, host, port), +// creating it when there is none. At most one such row is open at a time +// (uq_intercept_pending_target), across goroutines and across instances that +// share the database, so concurrent calls to one host share one request. +// created reports whether this call wrote the row. id == 0 means the request +// that was open got decided in between: the caller goes back to the pin lookup. +func (d *DB) FindOrCreateTargetRequest(r TargetRequest) (id int64, created bool, err error) { + if err := r.Scope.check(); err != nil { + return 0, false, err + } + raw := json.RawMessage(r.ToolInput) + if len(raw) == 0 { + raw = json.RawMessage("{}") + } + var ruleID *int64 + if r.RuleID != 0 { + ruleID = &r.RuleID + } + taskID, convID := r.Scope.pendingScope() + key := TargetKey(r.Host, r.Port) + err = d.QueryRow(` +INSERT INTO intercept_pending(rule_id, conversation_id, task_id, agent_name, tool_name, tool_input, reason, decision_source, audit, target_key) +VALUES ($1, $2, $3, $4, $5, $6, $7, 'target', $8, $9) +ON CONFLICT (COALESCE(task_id,''), COALESCE(conversation_id,0), target_key) + WHERE status = 'pending' AND target_key IS NOT NULL +DO NOTHING RETURNING id`, + ruleID, convID, taskID, r.AgentName, r.ToolName, raw, r.Reason, firstAudit([]*InterceptAudit{r.Audit}), key).Scan(&id) + if err == nil { + return id, true, nil + } + if !errors.Is(err, sql.ErrNoRows) { + return 0, false, err + } + taskKey := "" + if taskID != nil { + taskKey = *taskID + } + err = d.QueryRow(`SELECT id FROM intercept_pending +WHERE COALESCE(task_id,'')=$1 AND COALESCE(conversation_id,0)=$2 AND target_key=$3 AND status='pending'`, + taskKey, r.Scope.ConversationID, key).Scan(&id) + if errors.Is(err, sql.ErrNoRows) { + return 0, false, nil + } + return id, false, err +} + +// TargetRequestKey returns the target key of an approval row: "" for a classic +// ask and for an id that does not exist. +func (d *DB) TargetRequestKey(id int64) (string, error) { + var key sql.NullString + err := d.QueryRow(`SELECT target_key FROM intercept_pending WHERE id=$1`, id).Scan(&key) + if errors.Is(err, sql.ErrNoRows) { + return "", nil + } + return key.String, err +} + +// ResolveTargetRequest settles an open target request with a human decision, in +// one transaction. Allowed pins the addresses the request displayed: the active +// pin of the same key, if any, becomes superseded and a new version holds its +// addresses plus the displayed ones (addresses leave only by revocation). +// Denied pins nothing. approver and remoteAddr go to the pin and to the row's +// audit; decided_by on the row stays the actor class, human. ok is false when +// the row is not an open target request. If the pin cannot be written the row +// stays pending. +func (d *DB) ResolveTargetRequest(id int64, allowed bool, approver, remoteAddr string) (pin *TargetPin, ok bool, err error) { + tx, err := d.Begin() + if err != nil { + return nil, false, err + } + defer tx.Rollback() + var taskID sql.NullString + var convID sql.NullInt64 + var raw []byte + err = tx.QueryRow(`SELECT task_id, conversation_id, audit FROM intercept_pending +WHERE id=$1 AND status='pending' AND target_key IS NOT NULL FOR UPDATE`, id).Scan(&taskID, &convID, &raw) + if errors.Is(err, sql.ErrNoRows) { + return nil, false, nil + } + if err != nil { + return nil, false, err + } + var audit InterceptAudit + if len(raw) > 0 { + if err := json.Unmarshal(raw, &audit); err != nil { + return nil, false, err + } + } + status, action, reason, execution := "denied", "deny", "user rejected the target; nothing was pinned", "not_executed" + if allowed { + if len(audit.Targets) != 1 { + return nil, false, fmt.Errorf("target request %d records %d targets, want 1", id, len(audit.Targets)) + } + scope := TargetScopeRef{ConversationID: convID.Int64} + if taskID.Valid { + if scope.TaskID, err = strconv.ParseInt(taskID.String, 10, 64); err != nil { + return nil, false, fmt.Errorf("target request %d: task id %q: %w", id, taskID.String, err) + } + } + if pin, err = writeTargetPin(tx, scope, id, &audit.Targets[0], approver, remoteAddr); err != nil { + return nil, false, err + } + status, action, reason = "allowed", "allow", "user approved the target; the displayed addresses were pinned" + // A plain target request is host-level: no single call owns its result. + execution = "unknown" + if audit.MergedAsk && audit.Correlation == "exact" { + execution = "awaiting_result" + } + } else if len(audit.Targets) == 1 { + audit.Targets[0].Outcome = "denied" + } + fields := map[string]any{ + "effective_action": action, "decision_reason": reason, "execution_status": execution, + "targets": audit.Targets, + } + if approver != "" { + fields["approver"] = approver + } + if remoteAddr != "" { + fields["approver_addr"] = remoteAddr + } + patch, err := json.Marshal(fields) + if err != nil { + return nil, false, err + } + if _, err := tx.Exec(`UPDATE intercept_pending SET status=$2, decided_at=NOW(), decided_by='human', + audit=COALESCE(audit,'{}'::jsonb) || $3::jsonb WHERE id=$1`, id, status, patch); err != nil { + return nil, false, err + } + if err := tx.Commit(); err != nil { + return nil, false, err + } + return pin, true, nil +} + +// writeTargetPin inserts the pin an approved request creates and fills the +// request's target with what was pinned. The only writer of target_pins rows. +func writeTargetPin(tx *sql.Tx, scope TargetScopeRef, interceptID int64, t *InterceptTarget, approver, remoteAddr string) (*TargetPin, error) { + if err := scope.check(); err != nil { + return nil, err + } + // Versions count every pin the key ever had, revoked ones included, so a + // version number names one pin of the key for good. + var version int + if err := tx.QueryRow(`SELECT COALESCE(max(version),0)+1 FROM target_pins +WHERE `+targetPinScope+` AND host=$3 AND port=$4`, scope.TaskID, scope.ConversationID, t.Host, t.Port).Scan(&version); err != nil { + return nil, err + } + ips := map[string]bool{} + for _, ip := range t.ResolvedIPs { + ips[ip] = true + } + prev, err := scanTargetPin(tx.QueryRow(`SELECT `+targetPinCols+` FROM target_pins +WHERE `+targetPinScope+` AND host=$3 AND port=$4 AND status='active' FOR UPDATE`, scope.TaskID, scope.ConversationID, t.Host, t.Port)) + switch { + case errors.Is(err, sql.ErrNoRows): + case err != nil: + return nil, err + default: + for _, ip := range prev.PinnedIPs { + ips[ip] = true + } + if _, err := tx.Exec(`UPDATE target_pins SET status='superseded' WHERE id=$1`, prev.ID); err != nil { + return nil, err + } + } + pinned := make([]string, 0, len(ips)) + for ip := range ips { + pinned = append(pinned, ip) + } + sort.Strings(pinned) + var taskID, convID *int64 + if scope.TaskID > 0 { + taskID = &scope.TaskID + } else { + convID = &scope.ConversationID + } + resolvedAt := t.ResolvedAt + if resolvedAt.IsZero() { + resolvedAt = time.Now().UTC() + } + pin, err := scanTargetPin(tx.QueryRow(` +INSERT INTO target_pins(task_id, conversation_id, host, port, version, pinned_ips, resolver, resolved_at, intercept_id, decided_by, decided_from, name_only_effective) +VALUES ($1, $2, $3, $4, $5, $6::text[], $7, $8, $9, $10, $11, $12) +RETURNING `+targetPinCols, + taskID, convID, t.Host, t.Port, version, pinned, t.Resolver, resolvedAt, interceptID, approver, remoteAddr, t.NameOnlyEffective)) + if err != nil { + return nil, err + } + t.Outcome, t.PinID, t.PinVersion, t.PinnedIPs = "pinned", pin.ID, pin.Version, pin.PinnedIPs + return pin, nil +} diff --git a/db/target_pins_test.go b/db/target_pins_test.go new file mode 100644 index 0000000..9cc0a36 --- /dev/null +++ b/db/target_pins_test.go @@ -0,0 +1,924 @@ +package db + +import ( + "bytes" + "encoding/json" + "errors" + "fmt" + "slices" + "strconv" + "strings" + "sync" + "testing" + "time" + + "github.com/jackc/pgx/v5/pgconn" +) + +func openTargetDB(t *testing.T) *DB { + t.Helper() + d, err := Open(testDSN(t)) + if err != nil { + t.Skipf("postgres unavailable (%v) - skipping", err) + } + t.Cleanup(func() { d.Close() }) + return d +} + +// targetTask creates a task and removes it, its pins and its approval rows at +// the end of the test (pins and approval rows outlive a task delete by design). +func targetTask(t *testing.T, d *DB, opts TaskCreateOptions) *Task { + t.Helper() + task, err := d.CreateTaskWithOptions("target gate store", "pins and requests", opts) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { + _ = d.DeleteTask(task.ID) + _, _ = d.Exec(`DELETE FROM target_pins WHERE task_id=$1`, task.ID) + _, _ = d.Exec(`DELETE FROM intercept_pending WHERE task_id=$1`, strconv.FormatInt(task.ID, 10)) + }) + return task +} + +func targetConversation(t *testing.T, d *DB) *Conversation { + t.Helper() + c, err := d.CreateConversation("assistant", fmt.Sprintf("target gate %d", time.Now().UnixNano()), nil) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { + _ = d.DeleteConversation(c.ID) + _, _ = d.Exec(`DELETE FROM target_pins WHERE conversation_id=$1`, c.ID) + }) + return c +} + +// openTargetRequest opens the request for (scope, host, port) showing ips. +func openTargetRequest(t *testing.T, d *DB, scope TargetScopeRef, host string, port int, ips ...string) (int64, bool) { + t.Helper() + id, created, err := d.FindOrCreateTargetRequest(TargetRequest{ + Scope: scope, Host: host, Port: port, AgentName: "worker", ToolName: "WebFetch", + ToolInput: []byte(`{"url":"https://` + host + `/"}`), Reason: "target " + host + " is not approved", + Audit: &InterceptAudit{ + Correlation: "unavailable", InitialAction: "ask", ExecutionStatus: "not_started", + Targets: []InterceptTarget{{Host: host, Port: port, Kind: "domain", ResolvedIPs: ips, Resolver: "system", + ResolvedAt: time.Now().UTC(), Reason: "no_pin", Outcome: "requested"}}, + }, + }) + if err != nil { + t.Fatal(err) + } + return id, created +} + +// approveTarget opens a request for the key and approves it. +func approveTarget(t *testing.T, d *DB, scope TargetScopeRef, host string, port int, ips ...string) *TargetPin { + t.Helper() + id, created := openTargetRequest(t, d, scope, host, port, ips...) + if !created { + t.Fatalf("a request for %s:%d was already open", host, port) + } + pin, ok, err := d.ResolveTargetRequest(id, true, "ARTEX", "198.51.100.9") + if err != nil || !ok || pin == nil { + t.Fatalf("approve %s:%d: pin=%+v ok=%v err=%v", host, port, pin, ok, err) + } + return pin +} + +func pendingTargetRows(t *testing.T, d *DB, scope TargetScopeRef) int { + t.Helper() + var n int + task := "" + if scope.TaskID > 0 { + task = strconv.FormatInt(scope.TaskID, 10) + } + if err := d.QueryRow(`SELECT count(*) FROM intercept_pending +WHERE COALESCE(task_id,'')=$1 AND COALESCE(conversation_id,0)=$2 AND status='pending' AND target_key IS NOT NULL`, + task, scope.ConversationID).Scan(&n); err != nil { + t.Fatal(err) + } + return n +} + +// The whole schema is applied on every start. Applying it again must change +// nothing, add no CHECK to the gate's tables, and leave the cascade of +// intercept_pending.conversation_id as it is. +func TestTargetGateSchemaIdempotent(t *testing.T) { + snapshot := func() string { + t.Helper() + d := openTargetDB(t) + var b strings.Builder + for _, q := range []string{ + `SELECT table_name||'.'||column_name||' '||data_type||' null='||is_nullable||' default='||COALESCE(column_default,'') +FROM information_schema.columns +WHERE table_schema=current_schema() AND (table_name IN ('target_pins','target_gate_exemption_log') + OR (table_name IN ('tasks','conversations') AND column_name='target_gate') + OR (table_name='intercept_pending' AND column_name='target_key')) +ORDER BY table_name, ordinal_position`, + `SELECT indexdef FROM pg_indexes WHERE schemaname=current_schema() +AND (tablename IN ('target_pins','target_gate_exemption_log') OR indexname='uq_intercept_pending_target') ORDER BY indexname`, + `SELECT conrelid::regclass::text||' '||conname||' '||contype::text||' '||confdeltype::text FROM pg_constraint +WHERE conrelid IN ('target_pins'::regclass,'target_gate_exemption_log'::regclass,'intercept_pending'::regclass) +AND contype IN ('c','f','p','u') ORDER BY 1`, + } { + rows, err := d.Query(q) + if err != nil { + t.Fatal(err) + } + for rows.Next() { + var line string + if err := rows.Scan(&line); err != nil { + t.Fatal(err) + } + b.WriteString(line + "\n") + } + if err := rows.Err(); err != nil { + t.Fatal(err) + } + rows.Close() + } + return b.String() + } + first := snapshot() // Open applied the schema + second := snapshot() // and applied it again + if first != second { + t.Fatalf("schema changed on the second application:\n--- first\n%s--- second\n%s", first, second) + } + for _, want := range []string{ + "tasks.target_gate boolean null=NO default=false", + "conversations.target_gate boolean null=NO default=false", + "intercept_pending.target_key text null=YES default=", + "target_pins.name_only_effective boolean null=NO default=false", + "target_gate_exemption_log.digest text", + "uq_intercept_pending_target", + "uq_target_pins_active", + "idx_target_pins_scope", + // The cascade stays: no constraint migration. + "intercept_pending intercept_pending_conversation_id_fkey f c", + } { + if !strings.Contains(first, want) { + t.Errorf("schema snapshot lacks %q:\n%s", want, first) + } + } + // Only primary keys on the new tables: no CHECK, no foreign key. + for _, line := range strings.Split(first, "\n") { + if (strings.HasPrefix(line, "target_pins ") || strings.HasPrefix(line, "target_gate_exemption_log ")) && !strings.Contains(line, "_pkey p ") { + t.Errorf("unexpected constraint on a target-gate table: %s", line) + } + } +} + +// Many callers, two instances, one key: exactly one open row, and every caller +// gets its id. This is what proves that ON CONFLICT infers the partial +// expression index. +func TestTargetRequestFindOrCreateIsIdempotent(t *testing.T) { + a, b := openTargetDB(t), openTargetDB(t) + scope := TargetScopeRef{TaskID: targetTask(t, a, TaskCreateOptions{}).ID} + const callers = 24 + ids := make([]int64, callers) + made := make([]bool, callers) + var wg sync.WaitGroup + for n := range callers { + inst := a + if n%2 == 1 { + inst = b + } + wg.Go(func() { + id, created, err := inst.FindOrCreateTargetRequest(TargetRequest{ + Scope: scope, Host: "app.example.com", Port: 443, ToolName: "WebFetch", Reason: "concurrent", + Audit: &InterceptAudit{Targets: []InterceptTarget{{Host: "app.example.com", Port: 443, ResolvedIPs: []string{"203.0.113.7"}}}}, + }) + if err != nil { + t.Error(err) + } + ids[n], made[n] = id, created + }) + } + wg.Wait() + created := 0 + for n := range callers { + if ids[n] == 0 || ids[n] != ids[0] { + t.Fatalf("callers got different requests: %v", ids) + } + if made[n] { + created++ + } + } + if created != 1 { + t.Fatalf("%d callers created the row, want 1", created) + } + if n := pendingTargetRows(t, a, scope); n != 1 { + t.Fatalf("%d open rows, want 1", n) + } + // Another port, another scope: other keys, other rows. + if id, made := openTargetRequest(t, a, scope, "app.example.com", 80, "203.0.113.7"); !made || id == ids[0] { + t.Fatalf("port 80 shared the row of port 443 (id=%d created=%v)", id, made) + } + conv := targetConversation(t, a) + if id, made := openTargetRequest(t, a, TargetScopeRef{ConversationID: conv.ID}, "app.example.com", 443, "203.0.113.7"); !made || id == ids[0] { + t.Fatalf("a conversation shared the row of a task (id=%d created=%v)", id, made) + } + var source string + var key *string + if err := a.QueryRow(`SELECT decision_source, target_key FROM intercept_pending WHERE id=$1`, ids[0]).Scan(&source, &key); err != nil { + t.Fatal(err) + } + if source != "target" || key == nil || *key != "app.example.com:443" { + t.Fatalf("row source=%q key=%v, want target and app.example.com:443", source, key) + } +} + +// A decided request does not block the key: the next call opens a new row. +// There is no sticky denial. +func TestTargetRequestNewRowAfterDecision(t *testing.T) { + d := openTargetDB(t) + scope := TargetScopeRef{TaskID: targetTask(t, d, TaskCreateOptions{}).ID} + first, _ := openTargetRequest(t, d, scope, "app.example.com", 443, "203.0.113.7") + if again, made := openTargetRequest(t, d, scope, "app.example.com", 443, "203.0.113.7"); made || again != first { + t.Fatalf("second call got id=%d created=%v, want to attach to %d", again, made, first) + } + if pin, ok, err := d.ResolveTargetRequest(first, false, "ARTEX", "198.51.100.9"); err != nil || !ok || pin != nil { + t.Fatalf("deny: pin=%+v ok=%v err=%v", pin, ok, err) + } + // A second decision on the same row is refused and changes nothing. + if _, ok, err := d.ResolveTargetRequest(first, true, "ARTEX", "198.51.100.9"); err != nil || ok { + t.Fatalf("second decision: ok=%v err=%v, want refused", ok, err) + } + if pins, err := d.ListTargetPins(scope); err != nil || len(pins) != 0 { + t.Fatalf("a denied request left pins: %+v err=%v", pins, err) + } + second, made := openTargetRequest(t, d, scope, "app.example.com", 443, "203.0.113.7") + if !made || second == first { + t.Fatalf("after a denial: id=%d created=%v, want a new row", second, made) + } + if key, err := d.TargetRequestKey(second); err != nil || key != "app.example.com:443" { + t.Fatalf("target key = %q, %v", key, err) + } + classic, err := d.CreateInterceptPending(0, 0, strconv.FormatInt(scope.TaskID, 10), "worker", "Bash", nil, "classic ask") + if err != nil { + t.Fatal(err) + } + if key, err := d.TargetRequestKey(classic); err != nil || key != "" { + t.Fatalf("classic ask has target key %q, %v", key, err) + } + // A classic ask is not a target request: it cannot be resolved as one. + if _, ok, err := d.ResolveTargetRequest(classic, true, "ARTEX", ""); err != nil || ok { + t.Fatalf("classic ask resolved as a target request: ok=%v err=%v", ok, err) + } +} + +// Settling the row, superseding the old pin and writing the new one are one +// transaction: when the pin cannot be written the row stays pending and the +// previous pin stays active. +func TestResolveTargetRequestIsAtomic(t *testing.T) { + d := openTargetDB(t) + scope := TargetScopeRef{TaskID: targetTask(t, d, TaskCreateOptions{}).ID} + const host = "pin-insert-fails.example.com" + v1 := approveTarget(t, d, scope, host, 443, "203.0.113.7") + + if _, err := d.Exec(`CREATE OR REPLACE FUNCTION test_refuse_target_pin() RETURNS trigger AS $$ +BEGIN + IF NEW.host = '` + host + `' THEN RAISE EXCEPTION 'pin insert refused by the test'; END IF; + RETURN NEW; +END $$ LANGUAGE plpgsql`); err != nil { + t.Fatal(err) + } + t.Cleanup(func() { + _, _ = d.Exec(`DROP TRIGGER IF EXISTS test_refuse_target_pin ON target_pins`) + _, _ = d.Exec(`DROP FUNCTION IF EXISTS test_refuse_target_pin()`) + }) + if _, err := d.Exec(`CREATE TRIGGER test_refuse_target_pin BEFORE INSERT ON target_pins +FOR EACH ROW EXECUTE FUNCTION test_refuse_target_pin()`); err != nil { + t.Fatal(err) + } + + id, _ := openTargetRequest(t, d, scope, host, 443, "203.0.113.99") + if pin, ok, err := d.ResolveTargetRequest(id, true, "ARTEX", "198.51.100.9"); err == nil || ok || pin != nil { + t.Fatalf("approval with a failing pin insert: pin=%+v ok=%v err=%v, want an error", pin, ok, err) + } + row, err := d.GetInterceptPending(id) + if err != nil || row == nil || row.Status != "pending" || row.DecidedBy != "" { + t.Fatalf("row after the failed approval = %+v, %v; want still pending", row, err) + } + active, err := d.ActiveTargetPins(scope, host, 443) + if err != nil || len(active) != 1 || active[0].ID != v1.ID || active[0].Version != 1 { + t.Fatalf("active pins after the failed approval = %+v, %v; want v1 untouched", active, err) + } + + // Without the failure the same row is approved. + if _, err := d.Exec(`DROP TRIGGER test_refuse_target_pin ON target_pins`); err != nil { + t.Fatal(err) + } + if pin, ok, err := d.ResolveTargetRequest(id, true, "ARTEX", "198.51.100.9"); err != nil || !ok || pin == nil || pin.Version != 2 { + t.Fatalf("approval after the failure was removed: pin=%+v ok=%v err=%v", pin, ok, err) + } +} + +// Approving a request for a pinned key extends the pin: the new version holds +// the old addresses plus the displayed ones. Addresses leave only by revocation. +func TestTargetPinVersioningExtends(t *testing.T) { + d := openTargetDB(t) + scope := TargetScopeRef{TaskID: targetTask(t, d, TaskCreateOptions{}).ID} + v1 := approveTarget(t, d, scope, "app.example.com", 443, "203.0.113.7") + if v1.Version != 1 || !slices.Equal(v1.PinnedIPs, []string{"203.0.113.7"}) || v1.Status != TargetPinActive || + v1.DecidedBy != "ARTEX" || v1.DecidedFrom != "198.51.100.9" || v1.Resolver != "system" || v1.InterceptID == nil || + v1.TaskID == nil || *v1.TaskID != scope.TaskID || v1.ConversationID != nil || v1.NameOnlyEffective { + t.Fatalf("v1 = %+v", v1) + } + v2 := approveTarget(t, d, scope, "app.example.com", 443, "203.0.113.9", "203.0.113.7", "203.0.113.8") + if v2.Version != 2 || !slices.Equal(v2.PinnedIPs, []string{"203.0.113.7", "203.0.113.8", "203.0.113.9"}) { + t.Fatalf("v2 = %+v, want version 2 with the union, sorted", v2) + } + pins, err := d.ListTargetPins(scope) + if err != nil || len(pins) != 2 || pins[0].ID != v2.ID || pins[0].Status != TargetPinActive || + pins[1].ID != v1.ID || pins[1].Status != TargetPinSuperseded { + t.Fatalf("pins = %+v, %v; want v2 active then v1 superseded", pins, err) + } + // The request row records what was pinned and by whom; decided_by on the row + // stays the actor class. + detail, err := d.GetInterceptDetail(*v2.InterceptID) + if err != nil || detail == nil || detail.Status != "allowed" || detail.DecidedBy != "human" || detail.DecisionSource != "target" || + detail.Audit == nil || detail.Audit.Approver != "ARTEX" || detail.Audit.ApproverAddr != "198.51.100.9" || + detail.Audit.EffectiveAction != "allow" || len(detail.Audit.Targets) != 1 { + t.Fatalf("request row = %+v, %v", detail, err) + } + if got := detail.Audit.Targets[0]; got.Outcome != "pinned" || got.PinID != v2.ID || got.PinVersion != 2 || len(got.PinnedIPs) != 3 { + t.Fatalf("recorded target = %+v", got) + } +} + +// One active pin per key is a database fact, across instances. +func TestTargetPinOneActivePerKey(t *testing.T) { + a, b := openTargetDB(t), openTargetDB(t) + scope := TargetScopeRef{TaskID: targetTask(t, a, TaskCreateOptions{}).ID} + approveTarget(t, a, scope, "app.example.com", 443, "203.0.113.7") + insert := func(port int, status string) error { + _, err := b.Exec(`INSERT INTO target_pins(task_id, host, port, version, pinned_ips, status) +VALUES ($1, 'app.example.com', $2, 9, '{203.0.113.99}', $3)`, scope.TaskID, port, status) + return err + } + var pgErr *pgconn.PgError + if err := insert(443, "active"); !errors.As(err, &pgErr) || pgErr.Code != "23505" { + t.Fatalf("second active pin on one key: %v, want a unique violation", err) + } + if err := insert(443, "superseded"); err != nil { + t.Fatalf("a non-active pin on the same key must be allowed: %v", err) + } + if err := insert(8443, "active"); err != nil { + t.Fatalf("an active pin on another port must be allowed: %v", err) + } + approveTarget(t, b, scope, "app.example.com", 443, "203.0.113.8") + var active int + if err := a.QueryRow(`SELECT count(*) FROM target_pins WHERE task_id=$1 AND host='app.example.com' AND port=443 AND status='active'`, + scope.TaskID).Scan(&active); err != nil || active != 1 { + t.Fatalf("%d active pins on the key after a second approval (%v), want 1", active, err) + } +} + +func TestTargetPinRevokeIdempotent(t *testing.T) { + d := openTargetDB(t) + scope := TargetScopeRef{TaskID: targetTask(t, d, TaskCreateOptions{}).ID} + other := TargetScopeRef{TaskID: targetTask(t, d, TaskCreateOptions{}).ID} + pin := approveTarget(t, d, scope, "app.example.com", 443, "203.0.113.7") + + // Another scope cannot see or revoke the pin. + if got, err := d.RevokeTargetPin(other, pin.ID); err != nil || got != nil { + t.Fatalf("revoke through another scope = %+v, %v; want nil", got, err) + } + if got, err := d.RevokeTargetPin(scope, pin.ID+1000000); err != nil || got != nil { + t.Fatalf("revoke of an unknown pin = %+v, %v; want nil", got, err) + } + if active, err := d.ActiveTargetPins(scope, "app.example.com", 443); err != nil || len(active) != 1 { + t.Fatalf("pin no longer active after foreign revokes: %+v, %v", active, err) + } + + first, err := d.RevokeTargetPin(scope, pin.ID) + if err != nil || first == nil || first.Status != TargetPinRevoked || first.RevokedReason != TargetPinRevokedByOperator || first.RevokedAt == nil { + t.Fatalf("revoke = %+v, %v", first, err) + } + second, err := d.RevokeTargetPin(scope, pin.ID) + if err != nil || second == nil || second.Status != TargetPinRevoked || !second.RevokedAt.Equal(*first.RevokedAt) { + t.Fatalf("second revoke = %+v, %v; want the first revocation unchanged", second, err) + } + if active, err := d.ActiveTargetPins(scope, "app.example.com", 443); err != nil || len(active) != 0 { + t.Fatalf("a revoked pin still matches: %+v, %v", active, err) + } + // A new approval starts from the displayed addresses only, under a new version. + again := approveTarget(t, d, scope, "app.example.com", 443, "203.0.113.50") + if again.Version != 2 || !slices.Equal(again.PinnedIPs, []string{"203.0.113.50"}) { + t.Fatalf("pin after a revocation = %+v, want version 2 holding only the new address", again) + } +} + +// A pin on port 0 (port not declared) covers every port of the host. A pin on +// a specific port covers that port only, and never a target whose port is not +// declared. +func TestActiveTargetPinPortPrecedence(t *testing.T) { + d := openTargetDB(t) + scope := TargetScopeRef{TaskID: targetTask(t, d, TaskCreateOptions{}).ID} + anyPort := approveTarget(t, d, scope, "both.example.com", 0, "203.0.113.1") + exact := approveTarget(t, d, scope, "both.example.com", 443, "203.0.113.2") + only443 := approveTarget(t, d, scope, "exact.example.com", 443, "203.0.113.3") + + for _, tc := range []struct { + host string + port int + want []int64 + }{ + {"both.example.com", 443, []int64{exact.ID, anyPort.ID}}, // the exact port first + {"both.example.com", 80, []int64{anyPort.ID}}, + {"both.example.com", 8443, []int64{anyPort.ID}}, + {"both.example.com", 0, []int64{anyPort.ID}}, // never the pin on 443 + {"exact.example.com", 443, []int64{only443.ID}}, + {"exact.example.com", 80, nil}, + {"exact.example.com", 0, nil}, // a specific port never covers "not declared" + {"unknown.example.com", 443, nil}, + } { + pins, err := d.ActiveTargetPins(scope, tc.host, tc.port) + if err != nil { + t.Fatal(err) + } + var got []int64 + for _, p := range pins { + got = append(got, p.ID) + } + if !slices.Equal(got, tc.want) { + t.Errorf("candidates for %s:%d = %v, want %v", tc.host, tc.port, got, tc.want) + } + } + // The address sets stay apart: each candidate holds its own. + pins, _ := d.ActiveTargetPins(scope, "both.example.com", 443) + if len(pins) != 2 || !slices.Equal(pins[0].PinnedIPs, []string{"203.0.113.2"}) || !slices.Equal(pins[1].PinnedIPs, []string{"203.0.113.1"}) { + t.Fatalf("candidates = %+v, want two pins with separate addresses", pins) + } +} + +// A hard delete keeps the gate's history: pins revoked, approval rows readable, +// no target request left open. +func TestTaskDeleteKeepsPinsAndApprovalRows(t *testing.T) { + d := openTargetDB(t) + task := targetTask(t, d, TaskCreateOptions{}) + other := targetTask(t, d, TaskCreateOptions{}) + scope := TargetScopeRef{TaskID: task.ID} + pin := approveTarget(t, d, scope, "app.example.com", 443, "203.0.113.7") + kept := approveTarget(t, d, TargetScopeRef{TaskID: other.ID}, "app.example.com", 443, "203.0.113.7") + open, _ := openTargetRequest(t, d, scope, "new.example.com", 443, "203.0.113.20") + classic, err := d.CreateInterceptPending(0, 0, strconv.FormatInt(task.ID, 10), "worker", "Bash", nil, "classic ask") + if err != nil { + t.Fatal(err) + } + + if err := d.DeleteTask(task.ID); err != nil { + t.Fatal(err) + } + pins, err := d.ListTargetPins(scope) + if err != nil || len(pins) != 1 || pins[0].ID != pin.ID || pins[0].Status != TargetPinRevoked || + pins[0].RevokedReason != TargetPinRevokedTaskDeleted || pins[0].RevokedAt == nil || + !slices.Equal(pins[0].PinnedIPs, []string{"203.0.113.7"}) || pins[0].DecidedBy != "ARTEX" { + t.Fatalf("pins of the deleted task = %+v, %v; want one revoked pin, still readable", pins, err) + } + if active, err := d.ActiveTargetPins(scope, "app.example.com", 443); err != nil || len(active) != 0 { + t.Fatalf("a pin of the deleted task still matches: %+v, %v", active, err) + } + if active, err := d.ActiveTargetPins(TargetScopeRef{TaskID: other.ID}, "app.example.com", 443); err != nil || len(active) != 1 || active[0].ID != kept.ID { + t.Fatalf("the pin of another task was touched: %+v, %v", active, err) + } + approved, err := d.GetInterceptDetail(*pin.InterceptID) + if err != nil || approved == nil || approved.Status != "allowed" || approved.Audit == nil || approved.Audit.Approver != "ARTEX" { + t.Fatalf("approval row of the deleted task = %+v, %v; want still readable", approved, err) + } + settled, err := d.GetInterceptDetail(open) + if err != nil || settled == nil || settled.Status != "denied" || settled.DecidedBy != "system" || + settled.Audit == nil || settled.Audit.EffectiveAction != "deny" || settled.Audit.DecisionReason != "task deleted" { + t.Fatalf("open target request after the delete = %+v, %v; want denied by the system", settled, err) + } + if n := pendingTargetRows(t, d, scope); n != 0 { + t.Fatalf("%d target requests still open after the delete", n) + } + // Classic asks are not the gate's: untouched, as before. + if row, err := d.GetInterceptPending(classic); err != nil || row == nil || row.Status != "pending" { + t.Fatalf("classic ask after the delete = %+v, %v; want untouched", row, err) + } +} + +// archiveTargetTask moves a task into the cold archive and returns the archive +// id and its snapshot. +func archiveTargetTask(t *testing.T, d *DB, taskID int64) (int64, *TaskArchiveSnapshot, *bytes.Buffer) { + t.Helper() + if err := d.EnsureLLMRecordsTable(); err != nil { + t.Fatal(err) + } + if err := d.EnsureLLMUsageTable(); err != nil { + t.Fatal(err) + } + if err := d.SetPaused(taskID, true); err != nil { + t.Fatal(err) + } + archive, err := d.QueueTaskArchive(taskID) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _, _ = d.Exec(`DELETE FROM task_archives WHERE id=$1`, archive.ID) }) + if claimed, err := d.ClaimTaskArchiveJob(t.Context()); err != nil || claimed == nil || claimed.ID != archive.ID { + t.Fatalf("archive claim = %+v, %v", claimed, err) + } + var llmRecords bytes.Buffer + snapshot, err := d.SnapshotTaskArchiveWithLLMRecords(taskID, &llmRecords) + if err != nil { + t.Fatal(err) + } + if err := d.CompleteTaskArchive(archive.ID, snapshot, "/tmp/target-gate-task.tar.zst", "abc", 100, 50); err != nil { + t.Fatal(err) + } + return archive.ID, snapshot, &llmRecords +} + +func TestArchiveStubDeleteRevokesPins(t *testing.T) { + d := openTargetDB(t) + task := targetTask(t, d, TaskCreateOptions{}) + scope := TargetScopeRef{TaskID: task.ID} + pin := approveTarget(t, d, scope, "app.example.com", 443, "203.0.113.7") + archiveID, _, _ := archiveTargetTask(t, d, task.ID) + if active, err := d.ActiveTargetPins(scope, "app.example.com", 443); err != nil || len(active) != 1 { + t.Fatalf("pin while the task is an archived stub = %+v, %v; want still active", active, err) + } + if _, err := d.QueueTaskArchiveDelete(archiveID); err != nil { + t.Fatal(err) + } + if claimed, err := d.ClaimTaskArchiveJob(t.Context()); err != nil || claimed == nil || claimed.State != Deleting { + t.Fatalf("delete claim = %+v, %v", claimed, err) + } + if err := d.DeleteTaskArchiveStub(archiveID); err != nil { + t.Fatal(err) + } + pins, err := d.ListTargetPins(scope) + if err != nil || len(pins) != 1 || pins[0].ID != pin.ID || pins[0].Status != TargetPinRevoked || pins[0].RevokedReason != TargetPinRevokedTaskDeleted { + t.Fatalf("pins after the stub delete = %+v, %v; want one pin revoked as task_deleted", pins, err) + } +} + +// The archive carries neither pins nor the flag: both stay in the live tables +// while the task is a stub and are there after the restore. +func TestArchiveRestoreKeepsPinsAndFlag(t *testing.T) { + d := openTargetDB(t) + on := true + task := targetTask(t, d, TaskCreateOptions{TargetGate: &on}) + scope := TargetScopeRef{TaskID: task.ID} + pin := approveTarget(t, d, scope, "app.example.com", 443, "203.0.113.7") + open, _ := openTargetRequest(t, d, scope, "new.example.com", 443, "203.0.113.20") + archiveID, snapshot, llmRecords := archiveTargetTask(t, d, task.ID) + + var flag bool + if err := d.QueryRow(`SELECT target_gate FROM tasks WHERE id=$1`, task.ID).Scan(&flag); err != nil || !flag { + t.Fatalf("flag of the archived stub = %v, %v; want kept", flag, err) + } + if _, err := d.QueueTaskArchiveRestore(archiveID); err != nil { + t.Fatal(err) + } + if claimed, err := d.ClaimTaskArchiveJob(t.Context()); err != nil || claimed == nil || claimed.State != Restoring { + t.Fatalf("restore claim = %+v, %v", claimed, err) + } + if _, err := d.RestoreTaskArchiveWithLLMRecords(archiveID, snapshot, 0, bytes.NewReader(llmRecords.Bytes())); err != nil { + t.Fatal(err) + } + if err := d.CompleteTaskArchiveRestore(archiveID); err != nil { + t.Fatal(err) + } + live, err := d.GetTask(task.ID) + if err != nil || live == nil || !live.TargetGate { + t.Fatalf("restored task = %+v, %v; want the gate still on", live, err) + } + active, err := d.ActiveTargetPins(scope, "app.example.com", 443) + if err != nil || len(active) != 1 || active[0].ID != pin.ID { + t.Fatalf("pins after the restore = %+v, %v; want the pin still active", active, err) + } + // The request that was open at archive time comes back closed (a timeout, + // never an approval) with its key, and the key is free again. + row, err := d.GetInterceptPending(open) + if err != nil || row == nil || row.Status != "timeout" || row.DecidedBy != "system" { + t.Fatalf("restored open request = %+v, %v; want a system timeout", row, err) + } + if key, err := d.TargetRequestKey(open); err != nil || key != "new.example.com:443" { + t.Fatalf("restored target key = %q, %v", key, err) + } + if id, made := openTargetRequest(t, d, scope, "new.example.com", 443, "203.0.113.20"); !made || id == open { + t.Fatalf("request after the restore: id=%d created=%v, want a new row", id, made) + } +} + +// Archives written before the gate have no target_key: the row restores with +// NULL, a classic ask. +func TestRestoreInterceptRowsWithoutTargetKey(t *testing.T) { + d := openTargetDB(t) + task := targetTask(t, d, TaskCreateOptions{}) + id, err := d.CreateDecidedIntercept(0, 0, strconv.FormatInt(task.ID, 10), "worker", "Bash", nil, "allowed", "[model] allow", + &InterceptAudit{InitialAction: "allow", ExecutionStatus: "succeeded"}) + if err != nil { + t.Fatal(err) + } + var raw []byte + if err := d.QueryRow(`SELECT row_to_json(ip) FROM intercept_pending ip WHERE id=$1`, id).Scan(&raw); err != nil { + t.Fatal(err) + } + var row map[string]any + if err := json.Unmarshal(raw, &row); err != nil { + t.Fatal(err) + } + if _, ok := row["target_key"]; !ok { + t.Fatal("control: the live row was expected to carry the target_key column") + } + delete(row, "target_key") + archived, _ := json.Marshal([]map[string]any{row}) + tx, err := d.Begin() + if err != nil { + t.Fatal(err) + } + defer tx.Rollback() + if _, err := tx.Exec(`DELETE FROM intercept_pending WHERE id=$1`, id); err != nil { + t.Fatal(err) + } + if err := restoreInterceptRows(tx, archived); err != nil { + t.Fatal(err) + } + var key *string + var status string + if err := tx.QueryRow(`SELECT target_key, status FROM intercept_pending WHERE id=$1`, id).Scan(&key, &status); err != nil { + t.Fatal(err) + } + if key != nil || status != "allowed" { + t.Fatalf("restored row target_key=%v status=%s, want NULL and allowed", key, status) + } +} + +// Deleting a conversation revokes its pins in the same transaction. The pins +// stay readable; the request rows go with the conversation (cascade). A delete +// that fails leaves the pins active. +func TestConversationDeleteRevokesPinsInSameTx(t *testing.T) { + d := openTargetDB(t) + conv, other := targetConversation(t, d), targetConversation(t, d) + scope := TargetScopeRef{ConversationID: conv.ID} + pin := approveTarget(t, d, scope, "app.example.com", 443, "203.0.113.7") + kept := approveTarget(t, d, TargetScopeRef{ConversationID: other.ID}, "app.example.com", 443, "203.0.113.7") + openTargetRequest(t, d, scope, "new.example.com", 443, "203.0.113.20") + if pin.ConversationID == nil || *pin.ConversationID != conv.ID || pin.TaskID != nil { + t.Fatalf("conversation pin = %+v", pin) + } + + // A reference without a delete rule makes the delete fail after the revoke + // statement has run inside the transaction. + if _, err := d.Exec(`CREATE TABLE IF NOT EXISTS test_conversation_delete_block (conversation_id BIGINT REFERENCES conversations(id))`); err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _, _ = d.Exec(`DROP TABLE IF EXISTS test_conversation_delete_block`) }) + if _, err := d.Exec(`INSERT INTO test_conversation_delete_block VALUES ($1)`, conv.ID); err != nil { + t.Fatal(err) + } + if err := d.DeleteConversation(conv.ID); err == nil { + t.Fatal("control: the delete was expected to fail on the blocking reference") + } + if active, err := d.ActiveTargetPins(scope, "app.example.com", 443); err != nil || len(active) != 1 || active[0].ID != pin.ID { + t.Fatalf("pins after a failed delete = %+v, %v; want the pin still active", active, err) + } + if got, err := d.GetConversation(conv.ID); err != nil || got == nil { + t.Fatalf("conversation after a failed delete = %+v, %v", got, err) + } + + if _, err := d.Exec(`DELETE FROM test_conversation_delete_block`); err != nil { + t.Fatal(err) + } + if err := d.DeleteConversation(conv.ID); err != nil { + t.Fatal(err) + } + pins, err := d.ListTargetPins(scope) + if err != nil || len(pins) != 1 || pins[0].ID != pin.ID || pins[0].Status != TargetPinRevoked || + pins[0].RevokedReason != TargetPinRevokedConversationDelete || pins[0].RevokedAt == nil || + !slices.Equal(pins[0].PinnedIPs, []string{"203.0.113.7"}) || pins[0].DecidedBy != "ARTEX" || pins[0].InterceptID == nil { + t.Fatalf("pins of the deleted conversation = %+v, %v; want one revoked pin, still readable", pins, err) + } + var rows int + if err := d.QueryRow(`SELECT count(*) FROM intercept_pending WHERE conversation_id=$1`, conv.ID).Scan(&rows); err != nil || rows != 0 { + t.Fatalf("%d request rows left after the delete (%v), want 0: the cascade stays", rows, err) + } + // The pin names a request row that is gone; it is kept as a plain number. + if gone, err := d.GetInterceptPending(*pins[0].InterceptID); err != nil || gone != nil { + t.Fatalf("request row of the pin = %+v, %v; want gone", gone, err) + } + if active, err := d.ActiveTargetPins(TargetScopeRef{ConversationID: other.ID}, "app.example.com", 443); err != nil || len(active) != 1 || active[0].ID != kept.ID { + t.Fatalf("the pin of another conversation was touched: %+v, %v", active, err) + } +} + +// The batch delete revokes the pins of the ids it actually removed, and only +// those. +func TestConversationBatchDeleteRevokesOnlyDeletedIds(t *testing.T) { + d := openTargetDB(t) + gone, stays := targetConversation(t, d), targetConversation(t, d) + approveTarget(t, d, TargetScopeRef{ConversationID: gone.ID}, "app.example.com", 443, "203.0.113.7") + approveTarget(t, d, TargetScopeRef{ConversationID: stays.ID}, "app.example.com", 443, "203.0.113.7") + // A pin whose conversation id is in the request but does not exist: the + // delete removes nothing for it, so the pin must not be revoked. + const unknown = int64(1) << 60 + if _, err := d.Exec(`INSERT INTO target_pins(conversation_id, host, port, pinned_ips) VALUES ($1, 'app.example.com', 443, '{203.0.113.7}')`, unknown); err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _, _ = d.Exec(`DELETE FROM target_pins WHERE conversation_id=$1`, unknown) }) + + deleted, err := d.DeleteConversations([]int64{gone.ID, unknown}) + if err != nil || !slices.Equal(deleted, []int64{gone.ID}) { + t.Fatalf("deleted = %v, %v; want only %d", deleted, err, gone.ID) + } + status := func(convID int64) (string, string) { + t.Helper() + pins, err := d.ListTargetPins(TargetScopeRef{ConversationID: convID}) + if err != nil || len(pins) != 1 { + t.Fatalf("pins of conversation %d = %+v, %v", convID, pins, err) + } + return pins[0].Status, pins[0].RevokedReason + } + if s, reason := status(gone.ID); s != TargetPinRevoked || reason != TargetPinRevokedConversationDelete { + t.Fatalf("pin of the deleted conversation: %s/%s, want revoked/conversation_deleted", s, reason) + } + if s, _ := status(stays.ID); s != TargetPinActive { + t.Fatalf("pin of a conversation that was not in the request: %s, want active", s) + } + if s, _ := status(unknown); s != TargetPinActive { + t.Fatalf("pin of an id the delete did not remove: %s, want active", s) + } +} + +func TestTaskTargetGateDefaultsOff(t *testing.T) { + d := openTargetDB(t) + task := targetTask(t, d, TaskCreateOptions{}) + if task.TargetGate { + t.Fatal("a task created without the option has the gate on") + } + got, err := d.GetTask(task.ID) + if err != nil || got == nil || got.TargetGate { + t.Fatalf("stored task = %+v, %v; want the gate off", got, err) + } + off := false + if explicit := targetTask(t, d, TaskCreateOptions{TargetGate: &off}); explicit.TargetGate { + t.Fatal("a task created with the option false has the gate on") + } + c := targetConversation(t, d) + if c.TargetGate { + t.Fatal("a conversation created without the flag has the gate on") + } +} + +func TestTaskTargetGatePatchRoundTrip(t *testing.T) { + d := openTargetDB(t) + on, off := true, false + task := targetTask(t, d, TaskCreateOptions{TargetGate: &on}) + if !task.TargetGate { + t.Fatal("create with the option true returned the gate off") + } + if got, err := d.GetTask(task.ID); err != nil || got == nil || !got.TargetGate { + t.Fatalf("stored task = %+v, %v; want the gate on", got, err) + } + // A patch of another field leaves the flag alone. + name := "renamed" + if got, err := d.UpdateTask(task.ID, TaskPatch{Name: &name}); err != nil || got == nil || !got.TargetGate || got.Name != name { + t.Fatalf("after a name patch = %+v, %v; want the gate still on", got, err) + } + if got, err := d.UpdateTask(task.ID, TaskPatch{TargetGate: &off}); err != nil || got == nil || got.TargetGate || got.Name != name { + t.Fatalf("after switching off = %+v, %v", got, err) + } + if got, err := d.UpdateTask(task.ID, TaskPatch{TargetGate: &on}); err != nil || got == nil || !got.TargetGate { + t.Fatalf("after switching on = %+v, %v", got, err) + } + listed, err := d.ListTasks() + if err != nil { + t.Fatal(err) + } + for _, l := range listed { + if l.ID == task.ID && !l.TargetGate { + t.Fatal("the task list reports the gate off") + } + } +} + +func TestConversationTargetGateRoundTrip(t *testing.T) { + d := openTargetDB(t) + c, err := d.CreateConversationGated("assistant", fmt.Sprintf("gated %d", time.Now().UnixNano()), nil, true) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = d.DeleteConversation(c.ID) }) + if !c.TargetGate { + t.Fatal("create with the flag returned the gate off") + } + if got, err := d.GetConversation(c.ID); err != nil || got == nil || !got.TargetGate { + t.Fatalf("stored conversation = %+v, %v; want the gate on", got, err) + } + title, on, off := "renamed", true, false + if got, err := d.UpdateConversation(c.ID, ConversationPatch{Title: &title}); err != nil || got == nil || !got.TargetGate || got.Title != title { + t.Fatalf("after a title patch = %+v, %v; want the gate still on", got, err) + } + if got, err := d.UpdateConversation(c.ID, ConversationPatch{TargetGate: &off}); err != nil || got == nil || got.TargetGate || got.Title != title { + t.Fatalf("after switching off = %+v, %v", got, err) + } + if got, err := d.UpdateConversation(c.ID, ConversationPatch{TargetGate: &on}); err != nil || got == nil || !got.TargetGate { + t.Fatalf("after switching on = %+v, %v", got, err) + } + listed, err := d.ListConversations() + if err != nil { + t.Fatal(err) + } + for _, l := range listed { + if l.ID == c.ID && !l.TargetGate { + t.Fatal("the conversation list reports the gate off") + } + } +} + +// A retest conversation takes the flag of the finding's task and none of its +// pins. +func TestRetestConversationCopiesFlagNotPins(t *testing.T) { + d := openTargetDB(t) + on := true + for _, tc := range []struct { + name string + opts TaskCreateOptions + want bool + }{ + {"gate on", TaskCreateOptions{TargetGate: &on}, true}, + {"gate off", TaskCreateOptions{}, false}, + } { + t.Run(tc.name, func(t *testing.T) { + task := targetTask(t, d, tc.opts) + approveTarget(t, d, TargetScopeRef{TaskID: task.ID}, "app.example.com", 443, "203.0.113.7") + fid, err := d.AddFinding(task.ID, 0, "retest-gate", "gated finding", "high", "summary", "evidence", "worker", nil) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { + _, _ = d.Exec(`DELETE FROM conversations WHERE id IN (SELECT conversation_id FROM finding_retests WHERE finding_id=$1)`, fid) + _, _ = d.DeleteFinding(fid) + }) + _, conv, created, err := d.CreateFindingRetest(t.Context(), fid, "") + if err != nil || !created || conv == nil { + t.Fatalf("retest: conv=%+v created=%v err=%v", conv, created, err) + } + if conv.TargetGate != tc.want { + t.Fatalf("retest conversation gate = %v, want %v", conv.TargetGate, tc.want) + } + if stored, err := d.GetConversation(conv.ID); err != nil || stored == nil || stored.TargetGate != tc.want { + t.Fatalf("stored retest conversation = %+v, %v", stored, err) + } + if pins, err := d.ListTargetPins(TargetScopeRef{ConversationID: conv.ID}); err != nil || len(pins) != 0 { + t.Fatalf("retest conversation inherited pins: %+v, %v", pins, err) + } + }) + } + // A finding without a task: the gate is off. + fid, err := d.AddFinding(0, 0, "retest-gate", "orphan finding", "high", "summary", "evidence", "worker", nil) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { + _, _ = d.Exec(`DELETE FROM conversations WHERE id IN (SELECT conversation_id FROM finding_retests WHERE finding_id=$1)`, fid) + _, _ = d.DeleteFinding(fid) + }) + if _, conv, _, err := d.CreateFindingRetest(t.Context(), fid, ""); err != nil || conv == nil || conv.TargetGate { + t.Fatalf("retest of a finding without a task: conv=%+v err=%v; want the gate off", conv, err) + } +} + +// Every human decision records who took it and from where, classic asks +// included. decided_by keeps the actor class. +func TestResolveInterceptByRecordsApprover(t *testing.T) { + d := openTargetDB(t) + task := targetTask(t, d, TaskCreateOptions{}) + create := func() int64 { + t.Helper() + id, err := d.CreateInterceptPending(0, 0, strconv.FormatInt(task.ID, 10), "worker", "Bash", nil, "classic ask", + &InterceptAudit{InitialAction: "ask", ExecutionStatus: "not_started", Correlation: "exact"}) + if err != nil { + t.Fatal(err) + } + return id + } + id := create() + if ok, err := d.ResolveInterceptBy(id, "allowed", "allow", "user approved execution", "human", "ARTEX", "198.51.100.9"); err != nil || !ok { + t.Fatalf("resolve: ok=%v err=%v", ok, err) + } + got, err := d.GetInterceptDetail(id) + if err != nil || got == nil || got.DecidedBy != "human" || got.Audit == nil || got.Audit.Approver != "ARTEX" || + got.Audit.ApproverAddr != "198.51.100.9" || got.Audit.EffectiveAction != "allow" || got.Audit.ExecutionStatus != "awaiting_result" { + t.Fatalf("decided classic ask = %+v, %v", got, err) + } + // A system decision has no author. + id = create() + if ok, err := d.ResolveIntercept(id, "timeout", "deny", "approval timed out", "system"); err != nil || !ok { + t.Fatalf("resolve: ok=%v err=%v", ok, err) + } + var raw []byte + if err := d.QueryRow(`SELECT audit FROM intercept_pending WHERE id=$1`, id).Scan(&raw); err != nil { + t.Fatal(err) + } + if strings.Contains(string(raw), "approver") { + t.Fatalf("a system decision recorded an approver: %s", raw) + } +} diff --git a/db/task_archives_restore.go b/db/task_archives_restore.go index c2e85bf..42a0492 100644 --- a/db/task_archives_restore.go +++ b/db/task_archives_restore.go @@ -766,6 +766,13 @@ WHERE archive.id=$1 FOR UPDATE OF archive,task`, archiveID).Scan(&taskID, &expID if !errors.Is(err, sql.ErrNoRows) { return err } + // As in DeleteTaskCascadePrepared: pins stay as revoked history. + if err := revokeTaskPins(tx, taskID); err != nil { + return err + } + if err := settleTaskTargetRequests(tx, taskID); err != nil { + return err + } if _, err := tx.Exec(`DELETE FROM tasks WHERE id=$1 AND archived_at IS NOT NULL`, taskID); err != nil { return err } diff --git a/db/tasks.go b/db/tasks.go index 413b7ef..3c4cce2 100644 --- a/db/tasks.go +++ b/db/tasks.go @@ -52,6 +52,9 @@ type Task struct { // PlanQuietMode gates the planner quiet-mode (#86): edge-triggered rounds // are skipped while work is in flight and the open-goal set is unchanged. PlanQuietMode bool `json:"plan_quiet_mode"` + // TargetGate is the per-task opt-in of the approval gate for asset targets. + // false (default) = no gate, behaviour unchanged. + TargetGate bool `json:"target_gate"` // CoverageEnabled is the master switch of the "asset coverage" feature (default true). When false: // do not compute/show test coverage, do not auto-accumulate task_scope (source=auto), do not open // add_task_scope/list_untested_assets to the agent, and do not inject the coverage block into the @@ -173,6 +176,8 @@ type TaskCreateOptions struct { // chain instead of the single active head. nil = off (default), preserving // the ordered-failover semantics. LLMRotate *bool + // TargetGate opts the task into the approval gate for asset targets; nil = off. + TargetGate *bool } // CreateTaskWithOptions creates an exploration, task, direct source relations, @@ -251,10 +256,11 @@ VALUES ($1, 'fact', $2, 0, 'origin', 'system')`, expID, string(originPayload)); if opts.LLMRotate != nil { rotate = *opts.LLMRotate } + t.TargetGate = opts.TargetGate != nil && *opts.TargetGate if err := tx.QueryRow(` -INSERT INTO tasks(name, category_id, description, goal, exploration_id, llm_profile_id, active_llm_profile_id, timeout_seconds, plan_heartbeat_seconds, plan_quiet_mode, coverage_enabled, llm_rotate) -VALUES ($1,$2,$3,$4,$5,$6,$6,$7,$8,$9,$10,$11) -RETURNING id, status, paused, created_at`, opts.Name, opts.CategoryID, description, goal, expID, active, opts.TimeoutSeconds, opts.PlanHeartbeatSeconds, quietMode, coverageEnabled, rotate).Scan(&t.ID, &t.Status, &t.Paused, &t.CreatedAt); err != nil { +INSERT INTO tasks(name, category_id, description, goal, exploration_id, llm_profile_id, active_llm_profile_id, timeout_seconds, plan_heartbeat_seconds, plan_quiet_mode, coverage_enabled, llm_rotate, target_gate) +VALUES ($1,$2,$3,$4,$5,$6,$6,$7,$8,$9,$10,$11,$12) +RETURNING id, status, paused, created_at`, opts.Name, opts.CategoryID, description, goal, expID, active, opts.TimeoutSeconds, opts.PlanHeartbeatSeconds, quietMode, coverageEnabled, rotate, t.TargetGate).Scan(&t.ID, &t.Status, &t.Paused, &t.CreatedAt); err != nil { return nil, err } if err := insertTaskRelations(tx, t.ID, opts.SourceTaskIDs); err != nil { @@ -370,11 +376,11 @@ func insertTaskLLMProfiles(tx *sql.Tx, taskID int64, profileIDs []int64) error { const taskCols = `id, COALESCE(name,''), category_id, COALESCE((SELECT category.name FROM task_categories category WHERE category.id=tasks.category_id),''), -description, goal, exploration_id, status, paused, queued, queued_at, COALESCE(queue_mode,''), llm_profile_id, active_llm_profile_id, COALESCE(parent_ref,''), pinned_at, created_at, completed_at, COALESCE(timeout_seconds,0), COALESCE(plan_heartbeat_seconds,300), COALESCE(plan_quiet_mode,true), COALESCE(coverage_enabled,true), COALESCE(llm_rotate,false), first_run_at, deadline_at` +description, goal, exploration_id, status, paused, queued, queued_at, COALESCE(queue_mode,''), llm_profile_id, active_llm_profile_id, COALESCE(parent_ref,''), pinned_at, created_at, completed_at, COALESCE(timeout_seconds,0), COALESCE(plan_heartbeat_seconds,300), COALESCE(plan_quiet_mode,true), COALESCE(coverage_enabled,true), COALESCE(llm_rotate,false), first_run_at, deadline_at, target_gate` func scanTask(sc interface{ Scan(...any) error }) (*Task, error) { var t Task - if err := sc.Scan(&t.ID, &t.Name, &t.CategoryID, &t.CategoryName, &t.Description, &t.Goal, &t.ExplorationID, &t.Status, &t.Paused, &t.Queued, &t.QueuedAt, &t.QueueMode, &t.LLMProfileID, &t.ActiveLLMProfileID, &t.ParentRef, &t.PinnedAt, &t.CreatedAt, &t.CompletedAt, &t.TimeoutSeconds, &t.PlanHeartbeatSeconds, &t.PlanQuietMode, &t.CoverageEnabled, &t.LLMRotate, &t.FirstRunAt, &t.DeadlineAt); err != nil { + if err := sc.Scan(&t.ID, &t.Name, &t.CategoryID, &t.CategoryName, &t.Description, &t.Goal, &t.ExplorationID, &t.Status, &t.Paused, &t.Queued, &t.QueuedAt, &t.QueueMode, &t.LLMProfileID, &t.ActiveLLMProfileID, &t.ParentRef, &t.PinnedAt, &t.CreatedAt, &t.CompletedAt, &t.TimeoutSeconds, &t.PlanHeartbeatSeconds, &t.PlanQuietMode, &t.CoverageEnabled, &t.LLMRotate, &t.FirstRunAt, &t.DeadlineAt, &t.TargetGate); err != nil { return nil, err } t.Pinned = t.PinnedAt != nil @@ -420,6 +426,8 @@ type TaskPatch struct { Pinned *bool // PlanQuietMode flips the per-task quiet-mode switch (#86); nil = leave unchanged. PlanQuietMode *bool + // TargetGate flips the per-task target-gate opt-in; nil = leave unchanged. + TargetGate *bool } // UpdateTask applies a partial task name/pin mutation and returns the updated row. @@ -432,9 +440,10 @@ func (d *DB) UpdateTask(id int64, patch TaskPatch) (*Task, error) { WHEN $4::boolean THEN COALESCE(pinned_at, now()) ELSE NULL END, - plan_quiet_mode = CASE WHEN $5::boolean THEN $6 ELSE plan_quiet_mode END + plan_quiet_mode = CASE WHEN $5::boolean THEN $6 ELSE plan_quiet_mode END, + target_gate = CASE WHEN $7::boolean THEN $8 ELSE target_gate END WHERE id=$1 AND deleted_at IS NULL -RETURNING `+taskCols, id, patch.Name != nil, patch.Name, patch.Pinned, patch.PlanQuietMode != nil, patch.PlanQuietMode)) +RETURNING `+taskCols, id, patch.Name != nil, patch.Name, patch.Pinned, patch.PlanQuietMode != nil, patch.PlanQuietMode, patch.TargetGate != nil, patch.TargetGate)) if err == sql.ErrNoRows { return nil, nil } @@ -673,6 +682,14 @@ DELETE FROM assets a USING deletable d WHERE a.id=d.id`, id, expID) } // llm_usage (the token metering ledger) is intentionally NOT deleted with the // task — it is kept as historical accounting even after the task is gone. + // Target pins and approval rows are kept as history too: the pins are revoked, + // the open target requests settled, in this transaction. + if err := revokeTaskPins(tx, id); err != nil { + return result, err + } + if err := settleTaskTargetRequests(tx, id); err != nil { + return result, err + } if _, err := tx.Exec(`DELETE FROM tasks WHERE id=$1`, id); err != nil { return result, err } diff --git a/intercept/intercept.go b/intercept/intercept.go index 9390431..036d598 100644 --- a/intercept/intercept.go +++ b/intercept/intercept.go @@ -87,11 +87,14 @@ const DefaultTargetWaitBound = 300 * time.Second // (TaskID) or one conversation (ConversationID), never both. Enabled reports the // scope's opt-in flag and is read on every gated call; nil means the gate is off. // WaitBound caps the wait for a human decision; 0 means never wait. +// Interceptor is the one that owns the gate of this session, so that code +// holding only a context (an HTTP redirect callback) can reach it. type TargetScope struct { TaskID string ConversationID int64 Enabled func() bool WaitBound time.Duration + Interceptor *Interceptor } // WithTargetScope returns a child context carrying the target-gate scope of the @@ -112,36 +115,52 @@ type compiledRule struct { re *regexp.Regexp } -// pendingManager tracks in-flight "ask" requests via per-request channels. +// pendingManager tracks in-flight approval requests via per-waiter channels. A +// classic ask has one waiter; a target request can have several (every call +// that names the same host waits on the same row). type pendingManager struct { mu sync.Mutex - ch map[int64]chan bool + ch map[int64][]chan bool } -func newPendingManager() *pendingManager { return &pendingManager{ch: map[int64]chan bool{}} } +func newPendingManager() *pendingManager { return &pendingManager{ch: map[int64][]chan bool{}} } +// add registers one more waiter for id and returns its channel. func (p *pendingManager) add(id int64) chan bool { ch := make(chan bool, 1) p.mu.Lock() - p.ch[id] = ch + p.ch[id] = append(p.ch[id], ch) p.mu.Unlock() return ch } +// resolve wakes every waiter of id with the decision. func (p *pendingManager) resolve(id int64, allowed bool) { p.mu.Lock() - ch, ok := p.ch[id] + waiters := p.ch[id] delete(p.ch, id) p.mu.Unlock() - if ok { + for _, ch := range waiters { ch <- allowed } } -func (p *pendingManager) remove(id int64) { +// remove takes one waiter's channel out; the other waiters of id stay. +func (p *pendingManager) remove(id int64, ch chan bool) { p.mu.Lock() - delete(p.ch, id) - p.mu.Unlock() + defer p.mu.Unlock() + waiters := p.ch[id] + for n, w := range waiters { + if w == ch { + waiters = append(waiters[:n:n], waiters[n+1:]...) + break + } + } + if len(waiters) == 0 { + delete(p.ch, id) + return + } + p.ch[id] = waiters } // Reviewer runs the LLM fallback judge for one tool call and returns its verdict @@ -159,6 +178,18 @@ type Interceptor struct { enabledTools map[string]bool // nil means not loaded yet pending *pendingManager reviewer Reviewer // nil = LLM fallback judge not wired + // upstreamProxy reports whether target traffic leaves through an upstream + // egress proxy (nil = no). Injected by the server layer. + upstreamProxy func() bool +} + +// SetUpstreamProxyCheck installs the function that reports whether an upstream +// egress proxy is configured. Behind one the proxy chooses the address it +// connects to, so a target approval given while it is set is name-only. +func (i *Interceptor) SetUpstreamProxyCheck(f func() bool) { + i.mu.Lock() + i.upstreamProxy = f + i.mu.Unlock() } // SetReviewer installs the LLM fallback judge callback. Passing nil disables it. @@ -633,7 +664,7 @@ func (i *Interceptor) HandleAsk(ctx context.Context, convID int64, dec Decision, i.bindResult(ctx, pendingID, audit) ch := i.pending.add(pendingID) - defer i.pending.remove(pendingID) + defer i.pending.remove(pendingID, ch) // Polling clients can decide after INSERT commits but before the channel is // registered. Re-read after registration so that decision cannot be lost; // later decisions will be delivered through ch. @@ -707,7 +738,11 @@ func (i *Interceptor) HandleAsk(ctx context.Context, convID int64, dec Decision, var ErrAlreadyDecided = errors.New("approval already decided or does not exist; please refresh the list") // Decide resolves a pending request. Called by the HTTP decide endpoint. -func (i *Interceptor) Decide(pendingID int64, allowed bool) error { +// approver (subject of the session token) and remoteAddr identify the human +// behind the decision; they are recorded in the row's audit. On a target +// request "allowed" also pins the addresses the request displayed, and both +// values go to the pin: a human decision is the only thing that creates a pin. +func (i *Interceptor) Decide(pendingID int64, allowed bool, approver, remoteAddr string) error { status := "denied" if allowed { status = "allowed" @@ -716,7 +751,16 @@ func (i *Interceptor) Decide(pendingID int64, allowed bool) error { if allowed { action, reason = "allow", "user approved execution" } - resolved, err := i.db.ResolveIntercept(pendingID, status, action, reason, "human") + targetKey, err := i.db.TargetRequestKey(pendingID) + if err != nil { + return err + } + var resolved bool + if targetKey != "" { + _, resolved, err = i.db.ResolveTargetRequest(pendingID, allowed, approver, remoteAddr) + } else { + resolved, err = i.db.ResolveInterceptBy(pendingID, status, action, reason, "human", approver, remoteAddr) + } if err != nil { return err } diff --git a/intercept/target_exempt.go b/intercept/target_exempt.go new file mode 100644 index 0000000..4df287e --- /dev/null +++ b/intercept/target_exempt.go @@ -0,0 +1,53 @@ +package intercept + +import ( + "fmt" + "net/netip" +) + +// metadataAddrs are instance-metadata addresses that lie outside the link-local +// ranges and would otherwise pass as private or public unicast. +var metadataAddrs = []netip.Addr{ + netip.MustParseAddr("fd00:ec2::254"), // AWS IMDS over IPv6 + netip.MustParseAddr("100.100.100.200"), // Alibaba Cloud +} + +// ValidateExemptionHost checks the host of an entry an operator adds to the +// target-gate exemption list. host is expected in normalised form. +// +// An exemption on a name is resolved at check time and is void when the name +// resolves to a private, loopback, link-local or unspecified address. An IP +// literal is never resolved, so that rule cannot protect it and the literal is +// checked here instead. A private unicast address is allowed: an internal +// service the operator names on purpose (10.0.0.5:8080). Loopback, link-local +// (169.254.0.0/16 holds the cloud metadata address 169.254.169.254; fe80::/10), +// unspecified and multicast addresses are never accepted, nor are the metadata +// addresses outside the link-local ranges. A host that is not an IP literal +// passes: names are the void rule's business. +func ValidateExemptionHost(host string) error { + addr, err := netip.ParseAddr(host) + if err != nil { + return nil + } + addr = addr.Unmap().WithZone("") + kind := "" + switch { + case addr.IsLoopback(): + kind = "a loopback" + case addr.IsUnspecified(): + kind = "the unspecified" + case addr.IsLinkLocalUnicast(): + kind = "a link-local (metadata range)" + case addr.IsMulticast(): + kind = "a multicast" + } + for _, m := range metadataAddrs { + if addr == m { + kind = "an instance-metadata" + } + } + if kind != "" { + return fmt.Errorf("%s is %s address and can never be exempt from target approval", host, kind) + } + return nil +} diff --git a/intercept/target_exempt_test.go b/intercept/target_exempt_test.go new file mode 100644 index 0000000..a551242 --- /dev/null +++ b/intercept/target_exempt_test.go @@ -0,0 +1,29 @@ +package intercept + +import "testing" + +// An operator may exempt a private unicast IP literal, never a loopback, +// link-local (cloud metadata), unspecified or multicast one. +func TestValidateExemptionHost(t *testing.T) { + for _, host := range []string{ + "10.0.0.5", "172.16.0.1", "192.168.1.1", "fd00::1", // private unicast + "203.0.113.7", "2001:db8::1", // public + "pypi.org", "internal.corp", "localhost", // names: the void rule's business + } { + if err := ValidateExemptionHost(host); err != nil { + t.Errorf("%s refused: %v", host, err) + } + } + for _, host := range []string{ + "169.254.169.254", "169.254.170.2", "169.254.0.1", "fe80::1", "fe80::1%eth0", // link-local, metadata + "::ffff:169.254.169.254", // the same address, IPv4-mapped + "fd00:ec2::254", "100.100.100.200", // metadata outside link-local + "127.0.0.1", "127.8.8.8", "::1", "::ffff:127.0.0.1", // loopback + "0.0.0.0", "::", // unspecified + "224.0.0.1", "239.255.255.250", "ff02::1", "ff05::2", // multicast + } { + if err := ValidateExemptionHost(host); err == nil { + t.Errorf("%s accepted, want it refused", host) + } + } +} diff --git a/intercept/target_request.go b/intercept/target_request.go new file mode 100644 index 0000000..559e71a --- /dev/null +++ b/intercept/target_request.go @@ -0,0 +1,374 @@ +package intercept + +import ( + "context" + "encoding/json" + "fmt" + "net" + "strconv" + "strings" + "time" + + "github.com/Autumn-27/artex/db" +) + +// Block codes a target request can end with. They are carried in the +// model-facing text of the blocked call. +const ( + // TargetApprovalPending: no decision within the wait bound. The request + // stays open, so a later approval still creates the pin. + TargetApprovalPending = "TARGET_APPROVAL_PENDING" + // TargetDenied: the request was closed without an approval. + TargetDenied = "TARGET_DENIED" + // TargetGateUnavailable: the request could not be written or read. + TargetGateUnavailable = "TARGET_GATE_UNAVAILABLE" +) + +// targetPollInterval is how often a waiter re-reads its request row. The poll +// covers a decision taken on another instance that shares the database. +var targetPollInterval = 2 * time.Second + +// TargetRequest is one target of a gated call that no pin covers. +type TargetRequest struct { + Host string // normalised name, IP literal or CIDR; with Port the key of the request + Port int // 0 = port not declared + Kind string + // ResolvedIPs are the addresses resolved for this request. They are what the + // human reads and what an approval pins. + ResolvedIPs []string + Resolver string + ResolvedAt time.Time + Reason string // no_pin | new_ip | revoked + Pin *db.TargetPin // the active pin an approval would extend (new_ip); nil otherwise +} + +// TargetResult is the outcome of RequestTargets. Code is empty when every +// request the call waited on was approved (or was decided in between): the +// caller then re-runs the pin lookup and executes only if every target is +// covered. Otherwise Code is the block code and Host, Port and PendingID name +// the request that stopped the call. +type TargetResult struct { + Code string + Host string + Port int + PendingID int64 + Reason string + // AskMerged reports that the ask passed in was carried by a request this call + // created, so its decision answered both. When false the caller still has to + // raise the ask on its own. + AskMerged bool +} + +func (s TargetScope) ref() (db.TargetScopeRef, error) { + ref := db.TargetScopeRef{ConversationID: s.ConversationID} + if s.TaskID != "" { + id, err := strconv.ParseInt(s.TaskID, 10, 64) + if err != nil { + return ref, fmt.Errorf("target scope: task id %q: %w", s.TaskID, err) + } + ref.TaskID = id + } + return ref, nil +} + +// scopeName is how a request text names what an approval is bound to. +func (s TargetScope) scopeName() string { + if s.TaskID != "" { + return "task " + s.TaskID + } + return "conversation " + strconv.FormatInt(s.ConversationID, 10) +} + +// targetWaiter is one request row a call waits on. +type targetWaiter struct { + target TargetRequest + id int64 // 0 = the open request was decided in between: nothing to wait for + created bool // this call wrote the row +} + +// RequestTargets asks a human to approve the targets of one gated call that no +// pin covers, and waits for the decisions up to the scope's wait bound. +// +// There is one open request per (scope, host, port). The first call writes the +// row and emits the approval card once; later calls for the same key attach to +// that row and wait on it. A request is durable and host-level: a call that +// stops waiting (bound elapsed, cancelled) leaves it open, and a later approval +// still creates the pin. Nothing here ever pins: only Decide does, on a click. +// +// ask is the pending ask of a rule or of the judge for the same call, nil when +// there is none. It is merged into the first row this call creates, so one +// click answers both. A merged row is bound to its creating call like a classic +// ask: when that call leaves, the system settles it, and a rule timeout never +// allows it, because a timeout cannot pin. +func (i *Interceptor) RequestTargets(ctx context.Context, scope TargetScope, toolName string, input []byte, targets []TargetRequest, ask *Decision) TargetResult { + ref, err := scope.ref() + if err != nil { + return TargetResult{Code: TargetGateUnavailable, Reason: err.Error()} + } + if scope.WaitBound <= 0 { + // A call that never waits cannot own a call-bound row. + ask = nil + } + + var res TargetResult + var merged *targetWaiter + var mergedAsk Decision + // Whatever the way out, a merged row never outlives its creating call. This + // is a no-op once the row is decided. + defer func() { + if merged != nil { + i.settleMerged(merged.id, "denied", "the call left before a decision") + } + }() + + waiters := make([]*targetWaiter, 0, len(targets)) + for _, t := range targets { + w, err := i.openTargetRequest(ctx, ref, scope, toolName, input, t, ask) + if err != nil { + res.Code, res.Host, res.Port, res.Reason = TargetGateUnavailable, t.Host, t.Port, err.Error() + return res + } + if w.created && ask != nil { + merged, mergedAsk, res.AskMerged = w, *ask, true + ask = nil + } + waiters = append(waiters, w) + } + if len(waiters) == 0 { + return res + } + + blocked := func(w *targetWaiter, code, reason string) TargetResult { + res.Code, res.Host, res.Port, res.PendingID, res.Reason = code, w.target.Host, w.target.Port, w.id, reason + return res + } + if scope.WaitBound <= 0 { + return blocked(waiters[0], TargetApprovalPending, "this call does not wait for approvals; the request stays open") + } + bound := time.NewTimer(scope.WaitBound) + defer bound.Stop() + + // The merged row first: its rule timeout runs from the creation of the row. + if merged != nil { + var ruleTimeout <-chan time.Time + if mergedAsk.TimeoutEnabled { + secs := mergedAsk.TimeoutSeconds + if secs <= 0 { + secs = 60 + } + timer := time.NewTimer(time.Duration(secs) * time.Second) + defer timer.Stop() + ruleTimeout = timer.C + } + status, reason := "", "" + switch st := i.awaitTarget(ctx, merged.id, bound.C, ruleTimeout); st { + case "allowed": + case waitRuleTimeout: + status, reason = "timeout", "approval timed out" + if mergedAsk.TimeoutAction == "allow" { + reason += `; the timeout action "allow" was overridden, because a timeout cannot pin a target` + } + case waitElapsed: + status, reason = "timeout", "no approval within the wait bound" + case waitCancelled: + status, reason = "denied", "work cancelled" + case waitGone: + return blocked(merged, TargetGateUnavailable, "the approval request no longer exists") + default: + return blocked(merged, TargetDenied, "the request was denied") + } + // A human decision that lands while the call is leaving still counts. + if status != "" && (i.settleMerged(merged.id, status, reason) || !i.targetAllowed(merged.id)) { + return blocked(merged, TargetDenied, reason) + } + } + + for _, w := range waiters { + if w == merged || w.id == 0 { + continue + } + st := i.awaitTarget(ctx, w.id, bound.C, nil) + if (st == "denied" || st == "timeout") && !w.created && i.wasMerged(w.id) { + // The row this call attached to was bound to another call and closed + // with it. That says nothing about the host: ask once more, with a + // plain request. + retry, err := i.openTargetRequest(ctx, ref, scope, toolName, input, w.target, nil) + if err != nil { + return blocked(w, TargetGateUnavailable, err.Error()) + } + if retry.id == 0 { + continue + } + w = retry + st = i.awaitTarget(ctx, w.id, bound.C, nil) + } + switch st { + case "allowed": + case waitElapsed: + return blocked(w, TargetApprovalPending, "no approval within the wait bound; the request stays open") + case waitCancelled: + return blocked(w, TargetApprovalPending, "work cancelled; the request stays open") + case waitGone: + return blocked(w, TargetGateUnavailable, "the approval request no longer exists") + default: + return blocked(w, TargetDenied, "the request was denied") + } + } + return res +} + +// openTargetRequest finds or creates the request row for one target. When this +// call writes the row it emits the approval card; an attaching call writes and +// emits nothing, and never rewrites the addresses a human may be reading. +func (i *Interceptor) openTargetRequest(ctx context.Context, ref db.TargetScopeRef, scope TargetScope, toolName string, input []byte, t TargetRequest, ask *Decision) (*targetWaiter, error) { + i.mu.RLock() + proxied := i.upstreamProxy + i.mu.RUnlock() + // Behind the upstream proxy the proxy chooses the address: the approval is + // for the name, and the pin created from this row says so. + nameOnly := proxied != nil && proxied() + + dec := Decision{Action: "ask"} + if ask != nil { + dec = *ask + } + dec.Message = targetReason(scope, t, ask, nameOnly) + audit := auditFor(ctx, dec, input, "pending") + audit.MergedAsk = ask != nil + target := db.InterceptTarget{ + Host: t.Host, Port: t.Port, Kind: t.Kind, + ResolvedIPs: t.ResolvedIPs, Resolver: t.Resolver, ResolvedAt: t.ResolvedAt, + Reason: t.Reason, Outcome: "requested", NameOnlyEffective: nameOnly, + } + if t.Pin != nil { + target.PinID, target.PinVersion, target.PinnedIPs = t.Pin.ID, t.Pin.Version, t.Pin.PinnedIPs + } + audit.Targets = []db.InterceptTarget{target} + + _, agentName := taskInfoFromCtx(ctx) + id, created, err := i.db.FindOrCreateTargetRequest(db.TargetRequest{ + Scope: ref, Host: t.Host, Port: t.Port, RuleID: dec.RuleID, + AgentName: agentName, ToolName: toolName, ToolInput: input, Reason: dec.Message, Audit: audit, + }) + if err != nil { + return nil, err + } + w := &targetWaiter{target: t, id: id, created: created} + if !created { + return w, nil + } + if ask != nil { + // Call-bound like a classic ask: the result of this call belongs to the row. + i.bindResult(ctx, id, audit) + } + detail, _ := json.Marshal(map[string]any{ + "pending_id": id, + "tool": toolName, + "input": json.RawMessage(input), + }) + activity := db.Activity{ + Kind: "intercept_request", + Summary: fmt.Sprintf("tool %s requests approval for target %s (#%d)", toolName, targetName(t), id), + Detail: string(detail), + } + if scope.ConversationID != 0 { + _, _ = i.db.AppendConvActivity(scope.ConversationID, activity) + } else if emit := taskEmitFromCtx(ctx); emit != nil { + emit(activity) + } + return w, nil +} + +// Ways a wait ends without a decision on the row. +const ( + waitElapsed = "wait_elapsed" // the scope's wait bound + waitRuleTimeout = "wait_rule_timeout" // the timeout of a merged ask + waitCancelled = "wait_cancelled" // the call's context + waitGone = "wait_gone" // the row no longer exists (its conversation was deleted) +) + +// awaitTarget waits for the decision on one request row and returns its status +// (allowed | denied | timeout) or the way the wait ended. It never settles the +// row. The row is re-read before each wait: that covers a decision taken before +// the waiter registered, and one taken on another instance. +func (i *Interceptor) awaitTarget(ctx context.Context, id int64, bound, ruleTimeout <-chan time.Time) string { + ch := i.pending.add(id) + defer i.pending.remove(id, ch) + poll := time.NewTicker(targetPollInterval) + defer poll.Stop() + for { + row, err := i.db.GetInterceptPending(id) + if err == nil && row == nil { + return waitGone + } + if err == nil && row.Status != "pending" { + return row.Status + } + select { + case <-ch: + case <-poll.C: + case <-bound: + return waitElapsed + case <-ruleTimeout: + return waitRuleTimeout + case <-ctx.Done(): + return waitCancelled + } + } +} + +// settleMerged closes a merged row on behalf of its creating call, which is +// leaving, and wakes the calls attached to it. It reports whether the row was +// still open; a merged row closed this way never pins. +func (i *Interceptor) settleMerged(id int64, status, reason string) bool { + resolved, err := i.db.ResolveIntercept(id, status, "deny", reason, "system") + if err != nil || !resolved { + return false + } + i.pending.resolve(id, false) + return true +} + +func (i *Interceptor) targetAllowed(id int64) bool { + row, err := i.db.GetInterceptPending(id) + return err == nil && row != nil && row.Status == "allowed" +} + +func (i *Interceptor) wasMerged(id int64) bool { + detail, err := i.db.GetInterceptDetail(id) + return err == nil && detail != nil && detail.Audit != nil && detail.Audit.MergedAsk +} + +func targetName(t TargetRequest) string { + if t.Port == 0 { + return t.Host + } + return net.JoinHostPort(t.Host, strconv.Itoa(t.Port)) +} + +// targetReason is the text a human reads before deciding: host, port, the +// addresses an approval pins, and what the approval is bound to. +func targetReason(scope TargetScope, t TargetRequest, ask *Decision, nameOnly bool) string { + var b strings.Builder + if ask != nil && ask.Message != "" { + b.WriteString(ask.Message) + b.WriteString("\nThis call also names a target that needs approval. ") + } + name, bound := targetName(t), scope.scopeName() + if t.Port == 0 { + name += " (port not declared: every port)" + } + ips := strings.Join(t.ResolvedIPs, ", ") + if t.Pin != nil { + fmt.Fprintf(&b, "Target %s is pinned for %s (pin v%d: %s) and now resolves to %s (resolver: %s).", + name, bound, t.Pin.Version, strings.Join(t.Pin.PinnedIPs, ", "), ips, t.Resolver) + } else { + fmt.Fprintf(&b, "Target %s is not approved for %s. It resolves to %s (resolver: %s).", name, bound, ips, t.Resolver) + } + fmt.Fprintf(&b, " Approving pins these addresses for %s until revoked.", bound) + if nameOnly { + b.WriteString(" Egress goes through the upstream proxy, which chooses the address it connects to." + + " The addresses listed are what the platform resolved. This approval is name-only.") + } + return b.String() +} diff --git a/intercept/target_request_test.go b/intercept/target_request_test.go new file mode 100644 index 0000000..5110254 --- /dev/null +++ b/intercept/target_request_test.go @@ -0,0 +1,754 @@ +package intercept + +import ( + "context" + "fmt" + "slices" + "strconv" + "strings" + "sync/atomic" + "testing" + "time" + + "github.com/Autumn-27/artex/db" +) + +func TestPendingWakesAllWaiters(t *testing.T) { + p := newPendingManager() + a, b, gone := p.add(1), p.add(1), p.add(1) + other := p.add(2) + p.remove(1, gone) + p.resolve(1, true) + for name, ch := range map[string]chan bool{"first": a, "second": b} { + select { + case allowed := <-ch: + if !allowed { + t.Fatalf("%s waiter got a denial", name) + } + default: + t.Fatalf("%s waiter of the request was not woken", name) + } + } + select { + case <-gone: + t.Fatal("a removed waiter was woken") + case <-other: + t.Fatal("a waiter of another request was woken") + default: + } + if len(p.ch) != 1 || len(p.ch[2]) != 1 { + t.Fatalf("waiters left = %v, want only the one of request 2", p.ch) + } + // Resolving a request nobody waits on, twice, blocks nobody. + p.resolve(1, true) + p.resolve(2, false) + if allowed := <-other; allowed { + t.Fatal("the waiter of request 2 got an approval") + } + p.remove(2, other) + if len(p.ch) != 0 { + t.Fatalf("waiters left = %v, want none", p.ch) + } +} + +func targetTestDB(t *testing.T) *db.DB { + t.Helper() + dsn, _, err := db.DSN() + if err != nil { + t.Skipf("no database config (%v) - skipping", err) + } + d, err := db.Open(dsn) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { d.Close() }) + return d +} + +// taskScope creates a task and returns its target scope on ic, with the gate +// on and the given wait bound. +func taskScope(t *testing.T, d *db.DB, ic *Interceptor, wait time.Duration) TargetScope { + t.Helper() + task, err := d.CreateTaskWithOptions("target requests", "request, wait, decide", db.TaskCreateOptions{}) + if err != nil { + t.Fatal(err) + } + id := strconv.FormatInt(task.ID, 10) + t.Cleanup(func() { + _ = d.DeleteTask(task.ID) + _, _ = d.Exec(`DELETE FROM target_pins WHERE task_id=$1`, task.ID) + _, _ = d.Exec(`DELETE FROM intercept_pending WHERE task_id=$1`, id) + }) + return TargetScope{TaskID: id, Enabled: func() bool { return true }, WaitBound: wait, Interceptor: ic} +} + +func (s TargetScope) dbRef(t *testing.T) db.TargetScopeRef { + t.Helper() + ref, err := s.ref() + if err != nil { + t.Fatal(err) + } + return ref +} + +func target(host string, port int, ips ...string) TargetRequest { + return TargetRequest{Host: host, Port: port, Kind: "domain", ResolvedIPs: ips, Resolver: "system", ResolvedAt: time.Now().UTC(), Reason: "no_pin"} +} + +var webFetchInput = []byte(`{"url":"https://app.example.com/"}`) + +// request runs RequestTargets in the background and returns the channel its +// result arrives on. +func request(ctx context.Context, ic *Interceptor, scope TargetScope, ask *Decision, targets ...TargetRequest) <-chan TargetResult { + out := make(chan TargetResult, 1) + go func() { out <- ic.RequestTargets(ctx, scope, "WebFetch", webFetchInput, targets, ask) }() + return out +} + +func result(t *testing.T, ch <-chan TargetResult) TargetResult { + t.Helper() + select { + case res := <-ch: + return res + case <-time.After(20 * time.Second): + t.Fatal("the call is still waiting") + return TargetResult{} + } +} + +func stillWaiting(t *testing.T, ch <-chan TargetResult) { + t.Helper() + select { + case res := <-ch: + t.Fatalf("the call returned %+v, want it still waiting", res) + case <-time.After(4 * targetPollInterval): + } +} + +func eventually(t *testing.T, what string, cond func() bool) { + t.Helper() + deadline := time.Now().Add(20 * time.Second) + for !cond() { + if time.Now().After(deadline) { + t.Fatalf("timed out waiting for %s", what) + } + time.Sleep(5 * time.Millisecond) + } +} + +// openRows returns the ids of the open target requests of scope, oldest first. +func openRows(t *testing.T, d *db.DB, scope TargetScope) []int64 { + t.Helper() + rows, err := d.Query(`SELECT id FROM intercept_pending +WHERE COALESCE(task_id,'')=$1 AND COALESCE(conversation_id,0)=$2 AND status='pending' AND target_key IS NOT NULL ORDER BY id`, + scope.TaskID, scope.ConversationID) + if err != nil { + t.Fatal(err) + } + defer rows.Close() + var ids []int64 + for rows.Next() { + var id int64 + if err := rows.Scan(&id); err != nil { + t.Fatal(err) + } + ids = append(ids, id) + } + return ids +} + +func allRows(t *testing.T, d *db.DB, scope TargetScope) int { + t.Helper() + var n int + if err := d.QueryRow(`SELECT count(*) FROM intercept_pending +WHERE COALESCE(task_id,'')=$1 AND COALESCE(conversation_id,0)=$2 AND target_key IS NOT NULL`, scope.TaskID, scope.ConversationID).Scan(&n); err != nil { + t.Fatal(err) + } + return n +} + +// waiters is how many calls wait on request id through ic. +func waiters(ic *Interceptor, id int64) int { + ic.pending.mu.Lock() + defer ic.pending.mu.Unlock() + return len(ic.pending.ch[id]) +} + +// oneOpenRow waits for the scope's single open request and returns its id. +func oneOpenRow(t *testing.T, d *db.DB, scope TargetScope) int64 { + t.Helper() + var ids []int64 + eventually(t, "one open target request", func() bool { + ids = openRows(t, d, scope) + return len(ids) == 1 + }) + return ids[0] +} + +func pins(t *testing.T, d *db.DB, scope TargetScope) []*db.TargetPin { + t.Helper() + out, err := d.ListTargetPins(scope.dbRef(t)) + if err != nil { + t.Fatal(err) + } + return out +} + +// Calls for one host share one request, in one process and across two +// instances on one database: one row, one card, one decision wakes them all, +// one pin. +func TestConcurrentCallsSameHostShareOneRequest(t *testing.T) { + d := targetTestDB(t) + first, second := New(d), New(targetTestDB(t)) + scope := taskScope(t, d, first, time.Minute) + var cards atomic.Int32 + ctx := WithTaskContext(context.Background(), scope.TaskID, "work#1", func(a db.Activity) { + if a.Kind == "intercept_request" { + cards.Add(1) + } + }) + tgt := target("app.example.com", 443, "203.0.113.7") + a := request(ctx, first, scope, nil, tgt) + id := oneOpenRow(t, d, scope) + b := request(ctx, first, scope, nil, tgt) + c := request(ctx, second, scope, nil, tgt) + eventually(t, "every call to wait on the request", func() bool { return waiters(first, id) == 2 && waiters(second, id) == 1 }) + if ids := openRows(t, d, scope); len(ids) != 1 || allRows(t, d, scope) != 1 { + t.Fatalf("open requests = %v, want the one row %d", ids, id) + } + if n := cards.Load(); n != 1 { + t.Fatalf("%d approval cards emitted, want 1", n) + } + stillWaiting(t, a) + + // Decided on the second instance: its own waiter is woken by the channel, + // the waiters of the first instance by their poll. + if err := second.Decide(id, true, "ARTEX", "198.51.100.9"); err != nil { + t.Fatal(err) + } + for name, ch := range map[string]<-chan TargetResult{"creator": a, "attached": b, "other instance": c} { + if res := result(t, ch); res.Code != "" { + t.Fatalf("%s call: %+v, want approved", name, res) + } + } + got := pins(t, d, scope) + if len(got) != 1 || got[0].Status != db.TargetPinActive || got[0].Host != "app.example.com" || got[0].Port != 443 || + !slices.Equal(got[0].PinnedIPs, []string{"203.0.113.7"}) { + t.Fatalf("pins = %+v, want one active pin", got) + } + if allRows(t, d, scope) != 1 || cards.Load() != 1 { + t.Fatalf("rows=%d cards=%d after the decision, want 1 and 1", allRows(t, d, scope), cards.Load()) + } + if err := first.Decide(id, false, "ARTEX", ""); err != ErrAlreadyDecided { + t.Fatalf("second decision: %v, want ErrAlreadyDecided", err) + } +} + +// A call that stops waiting leaves the request alone: it stays open, the next +// call attaches to it, and a later approval still pins. +func TestWaiterLeavingKeepsRequestOpen(t *testing.T) { + d := targetTestDB(t) + ic := New(d) + scope := taskScope(t, d, ic, time.Minute) + tgt := target("app.example.com", 443, "203.0.113.7") + ctx, cancel := context.WithCancel(context.Background()) + call := request(ctx, ic, scope, nil, tgt) + id := oneOpenRow(t, d, scope) + eventually(t, "the call to wait", func() bool { return waiters(ic, id) == 1 }) + cancel() + res := result(t, call) + if res.Code != TargetApprovalPending || res.PendingID != id || res.Host != "app.example.com" || res.Port != 443 { + t.Fatalf("cancelled call: %+v, want %s on request %d", res, TargetApprovalPending, id) + } + row, err := d.GetInterceptPending(id) + if err != nil || row == nil || row.Status != "pending" || row.DecidedBy != "" { + t.Fatalf("request after its waiter left = %+v, %v; want untouched", row, err) + } + if waiters(ic, id) != 0 { + t.Fatal("the waiter that left is still registered") + } + + again := request(context.Background(), ic, scope, nil, tgt) + eventually(t, "the next call to attach", func() bool { return waiters(ic, id) == 1 }) + if ids := openRows(t, d, scope); !slices.Equal(ids, []int64{id}) { + t.Fatalf("open requests = %v, want the same row %d", ids, id) + } + if err := ic.Decide(id, true, "ARTEX", ""); err != nil { + t.Fatal(err) + } + if res := result(t, again); res.Code != "" { + t.Fatalf("attached call: %+v, want approved", res) + } + if got := pins(t, d, scope); len(got) != 1 { + t.Fatalf("pins = %+v, want the late approval to pin", got) + } +} + +func TestWaitBoundReturnsPendingCode(t *testing.T) { + d := targetTestDB(t) + ic := New(d) + scope := taskScope(t, d, ic, 150*time.Millisecond) + tgt := target("app.example.com", 443, "203.0.113.7") + start := time.Now() + res := result(t, request(context.Background(), ic, scope, nil, tgt)) + if res.Code != TargetApprovalPending || res.PendingID == 0 { + t.Fatalf("bounded call: %+v, want %s", res, TargetApprovalPending) + } + if waited := time.Since(start); waited < 150*time.Millisecond { + t.Fatalf("the call returned after %v, before its wait bound", waited) + } + if ids := openRows(t, d, scope); !slices.Equal(ids, []int64{res.PendingID}) { + t.Fatalf("open requests = %v, want the request %d left open", ids, res.PendingID) + } + + // Bound 0 (the final settlement round): never wait, but still ask. + never := scope + never.WaitBound = 0 + other := target("other.example.com", 443, "203.0.113.8") + res = result(t, request(context.Background(), ic, never, nil, other)) + if res.Code != TargetApprovalPending || res.Host != "other.example.com" || res.PendingID == 0 { + t.Fatalf("zero-bound call: %+v, want an immediate %s", res, TargetApprovalPending) + } + if len(openRows(t, d, scope)) != 2 { + t.Fatalf("open requests = %v, want the zero-bound call to have created its request", openRows(t, d, scope)) + } + // An ask is never merged into a row the call does not wait on. + asked := ic.RequestTargets(context.Background(), never, "WebFetch", webFetchInput, + []TargetRequest{target("third.example.com", 443, "203.0.113.9")}, &Decision{Action: "ask", Message: "rule asks"}) + if asked.Code != TargetApprovalPending || asked.AskMerged { + t.Fatalf("zero-bound call with an ask: %+v, want a plain open request", asked) + } + if row, err := d.GetInterceptPending(asked.PendingID); err != nil || row == nil || row.Status != "pending" { + t.Fatalf("request of the zero-bound call = %+v, %v; want left open", row, err) + } +} + +func askRule(t *testing.T, d *db.DB, timeoutEnabled bool, seconds int, action string) Decision { + t.Helper() + rule, err := d.CreateInterceptRule(fmt.Sprintf("target merge %d", time.Now().UnixNano()), "tool_name", "string", + "NoSuchToolForTargetMerge", "ask", "", 0, false, timeoutEnabled, seconds, action) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = d.DeleteInterceptRule(rule.ID) }) + return Decision{Action: "ask", Message: "rule requires approval", RuleID: rule.ID, RuleName: rule.Name, + TimeoutEnabled: timeoutEnabled, TimeoutSeconds: seconds, TimeoutAction: action} +} + +// A merged row is bound to the call that created it: when that call leaves, the +// system closes the row, and nothing is pinned. +func TestMergedRequestSettledWhenCreatorLeaves(t *testing.T) { + d := targetTestDB(t) + ic := New(d) + scope := taskScope(t, d, ic, time.Minute) + ask := askRule(t, d, false, 0, "deny") + ctx, cancel := context.WithCancel(context.Background()) + call := request(ctx, ic, scope, &ask, target("app.example.com", 443, "203.0.113.7")) + id := oneOpenRow(t, d, scope) + detail, err := d.GetInterceptDetail(id) + if err != nil || detail == nil || detail.RuleID == nil || *detail.RuleID != ask.RuleID || detail.DecisionSource != "target" || + detail.Audit == nil || !detail.Audit.MergedAsk || detail.Audit.RuleName != ask.RuleName || len(detail.Audit.Targets) != 1 { + t.Fatalf("merged row = %+v, %v; want the rule and the target on one row", detail, err) + } + if !strings.Contains(detail.Reason, "rule requires approval") || !strings.Contains(detail.Reason, "app.example.com:443") || + !strings.Contains(detail.Reason, "203.0.113.7") || !strings.Contains(detail.Reason, "Approving pins these addresses for task "+scope.TaskID+" until revoked.") { + t.Fatalf("merged reason does not name both the ask and the target: %q", detail.Reason) + } + eventually(t, "the creating call to wait", func() bool { return waiters(ic, id) == 1 }) + cancel() + res := result(t, call) + if res.Code != TargetDenied || !res.AskMerged || res.PendingID != id { + t.Fatalf("creating call after cancel: %+v, want %s with the ask merged", res, TargetDenied) + } + row, err := d.GetInterceptDetail(id) + if err != nil || row == nil || row.Status != "denied" || row.DecidedBy != "system" || row.Audit.EffectiveAction != "deny" || + row.Audit.DecisionReason != "work cancelled" { + t.Fatalf("merged row after its creator left = %+v, %v; want denied by the system", row, err) + } + if got := pins(t, d, scope); len(got) != 0 { + t.Fatalf("a row closed by the system pinned: %+v", got) + } +} + +// A timeout never pins, whatever the rule says: timeout_action = allow is +// overridden on a merged row and the override is recorded. +func TestTimeoutNeverPins(t *testing.T) { + d := targetTestDB(t) + ic := New(d) + scope := taskScope(t, d, ic, time.Minute) + ask := askRule(t, d, true, 1, "allow") + res := result(t, request(context.Background(), ic, scope, &ask, target("app.example.com", 443, "203.0.113.7"))) + if res.Code != TargetDenied || !res.AskMerged { + t.Fatalf("call after the rule timeout: %+v, want %s", res, TargetDenied) + } + row, err := d.GetInterceptDetail(res.PendingID) + if err != nil || row == nil || row.Status != "timeout" || row.DecidedBy != "system" || row.Audit.EffectiveAction != "deny" || + row.Audit.ExecutionStatus != "not_executed" || !strings.Contains(row.Audit.DecisionReason, `timeout action "allow" was overridden`) { + t.Fatalf("timed-out merged row = %+v, %v; want a timeout that denies and records the override", row, err) + } + if got := pins(t, d, scope); len(got) != 0 { + t.Fatalf("a timeout pinned: %+v", got) + } + if len(openRows(t, d, scope)) != 0 { + t.Fatal("the merged row is still open after its creator left") + } +} + +// What is displayed is what is pinned: the addresses stored when the row was +// created. A call that attaches with other addresses never rewrites them. +func TestDecidePinsDisplayedIPs(t *testing.T) { + d := targetTestDB(t) + ic := New(d) + scope := taskScope(t, d, ic, time.Minute) + creator := request(context.Background(), ic, scope, nil, target("app.example.com", 443, "203.0.113.8", "203.0.113.7")) + id := oneOpenRow(t, d, scope) + attached := request(context.Background(), ic, scope, nil, target("app.example.com", 443, "198.51.100.66")) + eventually(t, "both calls to wait", func() bool { return waiters(ic, id) == 2 }) + shown, err := d.GetInterceptDetail(id) + if err != nil || shown == nil || len(shown.Audit.Targets) != 1 || + !slices.Equal(shown.Audit.Targets[0].ResolvedIPs, []string{"203.0.113.8", "203.0.113.7"}) || strings.Contains(shown.Reason, "198.51.100.66") { + t.Fatalf("displayed request = %+v, %v; want the creator's addresses only", shown, err) + } + if err := ic.Decide(id, true, "ARTEX", "198.51.100.9"); err != nil { + t.Fatal(err) + } + if res := result(t, creator); res.Code != "" { + t.Fatalf("creator: %+v", res) + } + if res := result(t, attached); res.Code != "" { + t.Fatalf("attached: %+v", res) + } + got := pins(t, d, scope) + if len(got) != 1 || !slices.Equal(got[0].PinnedIPs, []string{"203.0.113.7", "203.0.113.8"}) || got[0].InterceptID == nil || *got[0].InterceptID != id { + t.Fatalf("pins = %+v, want exactly the displayed addresses", got) + } +} + +func TestDeniedRequestPinsNothing(t *testing.T) { + d := targetTestDB(t) + ic := New(d) + scope := taskScope(t, d, ic, time.Minute) + tgt := target("app.example.com", 443, "203.0.113.7") + call := request(context.Background(), ic, scope, nil, tgt) + id := oneOpenRow(t, d, scope) + if err := ic.Decide(id, false, "ARTEX", "198.51.100.9"); err != nil { + t.Fatal(err) + } + if res := result(t, call); res.Code != TargetDenied || res.PendingID != id { + t.Fatalf("denied call: %+v, want %s", res, TargetDenied) + } + if got := pins(t, d, scope); len(got) != 0 { + t.Fatalf("a denial pinned: %+v", got) + } + row, err := d.GetInterceptDetail(id) + if err != nil || row == nil || row.Status != "denied" || row.DecidedBy != "human" || row.Audit.Targets[0].Outcome != "denied" || + row.Audit.Approver != "ARTEX" || row.Audit.ApproverAddr != "198.51.100.9" { + t.Fatalf("denied row = %+v, %v", row, err) + } + // No sticky denial: the next call asks again, on a new row. + next := request(context.Background(), ic, scope, nil, tgt) + if fresh := oneOpenRow(t, d, scope); fresh == id { + t.Fatal("the next call reused the denied row") + } else if err := ic.Decide(fresh, true, "ARTEX", ""); err != nil { + t.Fatal(err) + } + if res := result(t, next); res.Code != "" { + t.Fatalf("call after the denial: %+v, want approved", res) + } +} + +// A call attached to a merged row that closes without an approval is not +// blocked for the host: the row was about another call. It asks once more with +// a plain request, then blocks. +func TestAttachedWaiterRequeuesAfterMergedDenial(t *testing.T) { + d := targetTestDB(t) + for _, approveRetry := range []bool{true, false} { + t.Run(fmt.Sprintf("retry approved=%v", approveRetry), func(t *testing.T) { + ic := New(d) + scope := taskScope(t, d, ic, time.Minute) + ask := askRule(t, d, false, 0, "deny") + tgt := target("app.example.com", 443, "203.0.113.7") + creator := request(context.Background(), ic, scope, &ask, tgt) + merged := oneOpenRow(t, d, scope) + attached := request(context.Background(), ic, scope, nil, tgt) + eventually(t, "both calls to wait", func() bool { return waiters(ic, merged) == 2 }) + if err := ic.Decide(merged, false, "ARTEX", ""); err != nil { + t.Fatal(err) + } + if res := result(t, creator); res.Code != TargetDenied || !res.AskMerged { + t.Fatalf("creator of the merged row: %+v, want %s", res, TargetDenied) + } + var retry int64 + eventually(t, "the attached call to ask again", func() bool { + ids := openRows(t, d, scope) + if len(ids) == 1 && ids[0] != merged { + retry = ids[0] + } + return retry != 0 + }) + plain, err := d.GetInterceptDetail(retry) + if err != nil || plain == nil || plain.RuleID != nil || plain.Audit.MergedAsk { + t.Fatalf("retry row = %+v, %v; want a plain target request", plain, err) + } + stillWaiting(t, attached) + if err := ic.Decide(retry, approveRetry, "ARTEX", ""); err != nil { + t.Fatal(err) + } + res := result(t, attached) + if approveRetry && (res.Code != "" || len(pins(t, d, scope)) != 1) { + t.Fatalf("attached call after its retry was approved: %+v pins=%+v", res, pins(t, d, scope)) + } + if !approveRetry && (res.Code != TargetDenied || res.PendingID != retry || len(pins(t, d, scope)) != 0) { + t.Fatalf("attached call after its retry was denied: %+v, want %s", res, TargetDenied) + } + // One retry, then block: never a third request. + if n := allRows(t, d, scope); n != 2 { + t.Fatalf("%d request rows, want the merged one and one retry", n) + } + }) + } +} + +// The request rows of a conversation go with it (cascade). A waiter that finds +// its row gone stops with a block. +func TestWaiterStopsWhenRowDeleted(t *testing.T) { + d := targetTestDB(t) + ic := New(d) + conv, err := d.CreateConversation("assistant", fmt.Sprintf("target wait %d", time.Now().UnixNano()), nil) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { + _ = d.DeleteConversation(conv.ID) + _, _ = d.Exec(`DELETE FROM target_pins WHERE conversation_id=$1`, conv.ID) + }) + scope := TargetScope{ConversationID: conv.ID, Enabled: func() bool { return true }, WaitBound: time.Minute, Interceptor: ic} + call := request(WithConvID(context.Background(), conv.ID), ic, scope, nil, target("app.example.com", 443, "203.0.113.7")) + id := oneOpenRow(t, d, scope) + eventually(t, "the call to wait", func() bool { return waiters(ic, id) == 1 }) + // The card went to the conversation stream. + var cards int + if err := d.QueryRow(`SELECT count(*) FROM conversation_activities WHERE conversation_id=$1 AND kind='intercept_request'`, conv.ID).Scan(&cards); err != nil || cards != 1 { + t.Fatalf("%d approval cards in the conversation (%v), want 1", cards, err) + } + if err := d.DeleteConversation(conv.ID); err != nil { + t.Fatal(err) + } + if res := result(t, call); res.Code != TargetGateUnavailable || res.PendingID != id { + t.Fatalf("call whose request was deleted: %+v, want %s", res, TargetGateUnavailable) + } +} + +// Every human decision records who took it and from where: on the pin and in +// the audit of a target row, in the audit of a classic ask. decided_by on the +// row stays the actor class. +func TestDecideRecordsApproverAndAddress(t *testing.T) { + d := targetTestDB(t) + ic := New(d) + scope := taskScope(t, d, ic, time.Minute) + call := request(context.Background(), ic, scope, nil, target("app.example.com", 443, "203.0.113.7")) + id := oneOpenRow(t, d, scope) + if err := ic.Decide(id, true, "ARTEX", "198.51.100.9"); err != nil { + t.Fatal(err) + } + if res := result(t, call); res.Code != "" { + t.Fatalf("approved call: %+v", res) + } + got := pins(t, d, scope) + if len(got) != 1 || got[0].DecidedBy != "ARTEX" || got[0].DecidedFrom != "198.51.100.9" { + t.Fatalf("pin = %+v, want the approver and the address", got) + } + row, err := d.GetInterceptDetail(id) + if err != nil || row == nil || row.DecidedBy != "human" || row.Audit.Approver != "ARTEX" || row.Audit.ApproverAddr != "198.51.100.9" { + t.Fatalf("target row = %+v, %v", row, err) + } + + // A classic ask goes through the same Decide and records the same. + ctx := WithTaskContext(context.Background(), scope.TaskID, "work#1", nil) + allowed := make(chan bool, 1) + go func() { + allowed <- ic.HandleAsk(ctx, 0, Decision{Action: "ask", Message: "classic ask"}, "Bash", []byte(`{"command":"id"}`)) + }() + var classic int64 + eventually(t, "the classic ask", func() bool { + return d.QueryRow(`SELECT id FROM intercept_pending WHERE task_id=$1 AND status='pending' AND target_key IS NULL`, scope.TaskID).Scan(&classic) == nil + }) + eventually(t, "the classic ask to wait", func() bool { return waiters(ic, classic) == 1 }) + if err := ic.Decide(classic, true, "ARTEX", "198.51.100.10"); err != nil { + t.Fatal(err) + } + select { + case ok := <-allowed: + if !ok { + t.Fatal("the approved classic ask was refused") + } + case <-time.After(20 * time.Second): + t.Fatal("the classic ask is still waiting") + } + row, err = d.GetInterceptDetail(classic) + if err != nil || row == nil || row.Status != "allowed" || row.DecidedBy != "human" || row.Audit == nil || + row.Audit.Approver != "ARTEX" || row.Audit.ApproverAddr != "198.51.100.10" || len(row.Audit.Targets) != 0 { + t.Fatalf("classic row = %+v, %v; want the approver and the address in its audit", row, err) + } + if len(pins(t, d, scope)) != 1 { + t.Fatal("approving a classic ask pinned something") + } +} + +// Behind the upstream proxy the proxy chooses the address: the request says the +// approval is name-only and the pin carries the mark. The mark stays when the +// proxy is removed; addresses are enforced again only after a new approval. +func TestPinUnderGlobalProxyIsNameOnly(t *testing.T) { + d := targetTestDB(t) + ic := New(d) + var proxied atomic.Bool + ic.SetUpstreamProxyCheck(proxied.Load) + scope := taskScope(t, d, ic, time.Minute) + const sentence = "This approval is name-only." + approve := func(host string) (*db.TargetPin, string) { + t.Helper() + call := request(context.Background(), ic, scope, nil, target(host, 443, "203.0.113.7")) + id := oneOpenRow(t, d, scope) + row, err := d.GetInterceptDetail(id) + if err != nil || row == nil { + t.Fatalf("request row: %+v, %v", row, err) + } + if err := ic.Decide(id, true, "ARTEX", ""); err != nil { + t.Fatal(err) + } + if res := result(t, call); res.Code != "" { + t.Fatalf("approved call: %+v", res) + } + active, err := d.ActiveTargetPins(scope.dbRef(t), host, 443) + if err != nil || len(active) != 1 { + t.Fatalf("active pins of %s = %+v, %v", host, active, err) + } + return active[0], row.Reason + } + + direct, reason := approve("direct.example.com") + if direct.NameOnlyEffective || strings.Contains(reason, sentence) { + t.Fatalf("without a proxy: pin=%+v reason=%q, want an address pin", direct, reason) + } + + proxied.Store(true) + named, reason := approve("proxied.example.com") + if !named.NameOnlyEffective || !strings.Contains(reason, sentence) || !strings.Contains(reason, "upstream proxy") { + t.Fatalf("behind the proxy: pin=%+v reason=%q, want a name-only pin and a request that says so", named, reason) + } + if !slices.Equal(named.PinnedIPs, []string{"203.0.113.7"}) { + t.Fatalf("name-only pin addresses = %v, want what the platform resolved", named.PinnedIPs) + } + + // The proxy is removed: the pin stays name-only, the address pin stays one. + proxied.Store(false) + ref := scope.dbRef(t) + if still, err := d.ActiveTargetPins(ref, "proxied.example.com", 443); err != nil || len(still) != 1 || !still[0].NameOnlyEffective { + t.Fatalf("name-only pin after the proxy was removed = %+v, %v; want still name-only", still, err) + } + if still, err := d.ActiveTargetPins(ref, "direct.example.com", 443); err != nil || len(still) != 1 || still[0].NameOnlyEffective { + t.Fatalf("address pin = %+v, %v", still, err) + } + // Enforcing addresses again takes a revocation and a new approval. + if _, err := d.RevokeTargetPin(ref, named.ID); err != nil { + t.Fatal(err) + } + if again, _ := approve("proxied.example.com"); again.NameOnlyEffective || again.Version != 2 { + t.Fatalf("pin approved after the proxy was removed = %+v, want an address pin, version 2", again) + } +} + +// A new_ip request shows the old version, the pinned addresses and the new +// ones, and its approval extends the pin. +func TestNewIPRequestShowsOldAndNewAddresses(t *testing.T) { + d := targetTestDB(t) + ic := New(d) + scope := taskScope(t, d, ic, time.Minute) + first := request(context.Background(), ic, scope, nil, target("app.example.com", 0, "203.0.113.7")) + id := oneOpenRow(t, d, scope) + if row, err := d.GetInterceptDetail(id); err != nil || row == nil || !strings.Contains(row.Reason, "app.example.com (port not declared: every port)") { + t.Fatalf("port-0 request = %+v, %v; want the reason to say the port is not declared", row, err) + } + if err := ic.Decide(id, true, "ARTEX", ""); err != nil { + t.Fatal(err) + } + if res := result(t, first); res.Code != "" { + t.Fatalf("first call: %+v", res) + } + v1 := pins(t, d, scope)[0] + + moved := target("app.example.com", 0, "203.0.113.99") + moved.Reason, moved.Pin = "new_ip", v1 + second := request(context.Background(), ic, scope, nil, moved) + id = oneOpenRow(t, d, scope) + row, err := d.GetInterceptDetail(id) + if err != nil || row == nil || !strings.Contains(row.Reason, "pin v1: 203.0.113.7") || !strings.Contains(row.Reason, "now resolves to 203.0.113.99") { + t.Fatalf("new_ip request = %+v, %v; want the old version, the pinned and the new addresses", row, err) + } + if got := row.Audit.Targets[0]; got.Reason != "new_ip" || got.PinID != v1.ID || got.PinVersion != 1 || !slices.Equal(got.PinnedIPs, []string{"203.0.113.7"}) { + t.Fatalf("recorded target = %+v", got) + } + if err := ic.Decide(id, true, "ARTEX", ""); err != nil { + t.Fatal(err) + } + if res := result(t, second); res.Code != "" { + t.Fatalf("second call: %+v", res) + } + active, err := d.ActiveTargetPins(scope.dbRef(t), "app.example.com", 0) + if err != nil || len(active) != 1 || active[0].Version != 2 || !slices.Equal(active[0].PinnedIPs, []string{"203.0.113.7", "203.0.113.99"}) { + t.Fatalf("pin after the new_ip approval = %+v, %v; want v2 with both addresses", active, err) + } +} + +// One call, several targets: one request per host and port, and the call +// passes only when all of them are approved. +func TestSeveralTargetsNeedEveryApproval(t *testing.T) { + d := targetTestDB(t) + ic := New(d) + scope := taskScope(t, d, ic, time.Minute) + ask := askRule(t, d, false, 0, "deny") + call := request(context.Background(), ic, scope, &ask, + target("a.example.com", 443, "203.0.113.1"), target("b.example.com", 443, "203.0.113.2")) + var ids []int64 + eventually(t, "two open requests", func() bool { + ids = openRows(t, d, scope) + return len(ids) == 2 + }) + // The ask is merged into the first row the call created; the other is plain. + for n, id := range ids { + row, err := d.GetInterceptDetail(id) + if err != nil || row == nil || row.Audit.MergedAsk != (n == 0) || (row.RuleID != nil) != (n == 0) { + t.Fatalf("row %d = %+v, %v; want the ask merged into the first row only", n, row, err) + } + } + if err := ic.Decide(ids[0], true, "ARTEX", ""); err != nil { + t.Fatal(err) + } + stillWaiting(t, call) + if err := ic.Decide(ids[1], false, "ARTEX", ""); err != nil { + t.Fatal(err) + } + res := result(t, call) + if res.Code != TargetDenied || res.Host != "b.example.com" || !res.AskMerged { + t.Fatalf("call with one target denied: %+v, want %s on b.example.com", res, TargetDenied) + } + if got := pins(t, d, scope); len(got) != 1 || got[0].Host != "a.example.com" { + t.Fatalf("pins = %+v, want only the approved host", got) + } +} + +// A scope that names no task and no conversation, or a task id that is not a +// number, never reaches the store. +func TestRequestTargetsRejectsBadScope(t *testing.T) { + ic := New(nil) + for name, scope := range map[string]TargetScope{ + "empty": {WaitBound: time.Minute}, + "both": {TaskID: "7", ConversationID: 9, WaitBound: time.Minute}, + "task not int": {TaskID: "seven", WaitBound: time.Minute}, + } { + res := ic.RequestTargets(context.Background(), scope, "WebFetch", webFetchInput, []TargetRequest{target("app.example.com", 443, "203.0.113.7")}, nil) + if res.Code != TargetGateUnavailable { + t.Errorf("%s scope: %+v, want %s", name, res, TargetGateUnavailable) + } + } +} diff --git a/intercept/testmain_test.go b/intercept/testmain_test.go new file mode 100644 index 0000000..c362b9c --- /dev/null +++ b/intercept/testmain_test.go @@ -0,0 +1,17 @@ +package intercept + +import ( + "os" + "testing" + "time" + + "github.com/Autumn-27/artex/db" +) + +func TestMain(m *testing.M) { + // Own database per package (#80), as in db, agent, server and evidence. + db.SuffixTestDatabase("intercept") + // A short poll keeps the cross-instance tests fast; production uses 2 s. + targetPollInterval = 50 * time.Millisecond + os.Exit(m.Run()) +} diff --git a/server/auth.go b/server/auth.go index a932332..b76a1d2 100644 --- a/server/auth.go +++ b/server/auth.go @@ -178,6 +178,28 @@ func verifyJWT(tokenStr string, key []byte) bool { return err == nil && t.Valid } +// jwtSubject returns the subject of tokenStr when it is a valid, non-expired +// HS256 token, and "" otherwise. +func jwtSubject(tokenStr string, key []byte) string { + var claims jwt.RegisteredClaims + t, err := jwt.ParseWithClaims(tokenStr, &claims, func(t *jwt.Token) (any, error) { + if _, ok := t.Method.(*jwt.SigningMethodHMAC); !ok { + return nil, fmt.Errorf("unexpected signing method") + } + return key, nil + }) + if err != nil || !t.Valid { + return "" + } + return claims.Subject +} + +// approver is who a human decision is recorded for: the subject of the +// request's session token. +func (s *Server) approver(r *http.Request) string { + return jwtSubject(extractToken(r), s.jwtKey) +} + // extractToken reads the JWT from Authorization: Bearer header, // artex_token cookie, or ?token= query param (for SSE connections). func extractToken(r *http.Request) string { diff --git a/server/conversations.go b/server/conversations.go index fae6e29..9aa9b28 100644 --- a/server/conversations.go +++ b/server/conversations.go @@ -79,6 +79,7 @@ func (s *Server) pgCreateConversation(w http.ResponseWriter, r *http.Request) { AgentKey string `json:"agent_key"` Title string `json:"title"` LLMProfileID *int64 `json:"llm_profile_id"` + TargetGate bool `json:"target_gate"` // approval gate for asset targets, opt-in } if !decodeConversationRequest(w, r, &req) { return @@ -115,7 +116,7 @@ func (s *Server) pgCreateConversation(w http.ResponseWriter, r *http.Request) { writeErr(w, 400, fmt.Sprintf("Title must be at most %d characters", maxConversationTitleRunes)) return } - c, err := pg.CreateConversation(req.AgentKey, title, req.LLMProfileID) + c, err := pg.CreateConversationGated(req.AgentKey, title, req.LLMProfileID, req.TargetGate) if err != nil { writeErr(w, 500, err.Error()) return @@ -176,14 +177,15 @@ func (s *Server) pgRenameConversation(w http.ResponseWriter, r *http.Request) { return } var req struct { - Title *string `json:"title"` - Pinned *bool `json:"pinned"` + Title *string `json:"title"` + Pinned *bool `json:"pinned"` + TargetGate *bool `json:"target_gate"` } if !decodeConversationRequest(w, r, &req) { return } - if req.Title == nil && req.Pinned == nil { - writeErr(w, 400, "Provide at least title or pinned") + if req.Title == nil && req.Pinned == nil && req.TargetGate == nil { + writeErr(w, 400, "Provide at least title, pinned or target_gate") return } if req.Title != nil { @@ -198,7 +200,7 @@ func (s *Server) pgRenameConversation(w http.ResponseWriter, r *http.Request) { } req.Title = &title } - updated, err := pg.UpdateConversation(c.ID, db.ConversationPatch{Title: req.Title, Pinned: req.Pinned}) + updated, err := pg.UpdateConversation(c.ID, db.ConversationPatch{Title: req.Title, Pinned: req.Pinned, TargetGate: req.TargetGate}) if err != nil { writeErr(w, 500, err.Error()) return @@ -559,8 +561,11 @@ func (s *Server) runConversationTurn(ctx context.Context, cancel context.CancelC } } // Approval scope of this turn's tool calls for the target gate (chat and retester). - // Conversations have no opt-in flag yet, so Enabled stays nil: the gate is off. - gctx := intercept.WithTargetScope(ctx, intercept.TargetScope{ConversationID: c.ID, WaitBound: intercept.DefaultTargetWaitBound}) + // The conversation's opt-in flag is the one read when the turn started: a + // switch through the API applies from the next turn. + gate := c.TargetGate + gctx := intercept.WithTargetScope(ctx, intercept.TargetScope{ConversationID: c.ID, Enabled: func() bool { return gate }, + WaitBound: intercept.DefaultTargetWaitBound, Interceptor: s.m.interceptor}) // On a manual stop ctx is cancelled; Chat already emits a clean "stopped by the user" // step, so skip the raw-error entry — only surface genuine failures. if _, err := ca.Chat(gctx, c.AgentKey, sessionID, msg, maxTurns, maxDuration, webSearch, emit); err != nil { diff --git a/server/dto.go b/server/dto.go index 282cf80..8d0c79b 100644 --- a/server/dto.go +++ b/server/dto.go @@ -63,6 +63,7 @@ type TaskDTO struct { CompanyIDs []int64 `json:"company_ids"` CoverageEnabled bool `json:"coverage_enabled"` // asset coverage feature switch (fixed at creation) PlanQuietMode bool `json:"plan_quiet_mode"` // planner quiet-mode switch (#86), editable mid-run + TargetGate bool `json:"target_gate"` // approval gate for asset targets, opt-in, editable mid-run } // FindingSeverityDTO is the per-severity finding count shown in the task list (critical/high/medium/low). @@ -138,6 +139,7 @@ func taskDTO(t *Task, status string) TaskDTO { CompanyIDs: lifecycle.CompanyIDs, CoverageEnabled: t.CoverageEnabled, PlanQuietMode: t.PlanQuietMode, + TargetGate: t.targetGate.Load(), } } diff --git a/server/intercept.go b/server/intercept.go index adc9f42..20d8cd7 100644 --- a/server/intercept.go +++ b/server/intercept.go @@ -328,9 +328,9 @@ func interceptFilterParams(q url.Values) (db.InterceptApprovalFilter, error) { return filter, fmt.Errorf("status must be pending, allowed, denied or timeout") } switch filter.DecisionSource { - case "", "model", "rule", "unknown": + case "", "model", "rule", "target", "unknown": default: - return filter, fmt.Errorf("decision_source must be model, rule or unknown") + return filter, fmt.Errorf("decision_source must be model, rule, target or unknown") } switch filter.DecidedBy { case "", "human", "rule", "model", "system": @@ -372,7 +372,9 @@ func (s *Server) interceptDecide(w http.ResponseWriter, r *http.Request) { writeErr(w, 400, "decision must be allowed or denied") return } - if err := s.m.interceptor.Decide(id, req.Decision == "allowed"); err != nil { + // On a target request "allowed" pins the addresses the request displayed. + // Every human decision records who took it and from where. + if err := s.m.interceptor.Decide(id, req.Decision == "allowed", s.approver(r), clientIP(r, s.trustedProxies)); err != nil { if errors.Is(err, intercept.ErrAlreadyDecided) { writeErr(w, 409, err.Error()) return @@ -383,6 +385,84 @@ func (s *Server) interceptDecide(w http.ResponseWriter, r *http.Request) { writeJSON(w, 200, map[string]any{"ok": true}) } +// --- target pins (hosts a human approved for one task or one conversation) --- + +// taskPinScope resolves {id} to the pin scope of an existing task. +func (s *Server) taskPinScope(w http.ResponseWriter, r *http.Request) (db.TargetScopeRef, bool) { + id, ok := pathInt(r, "id") + if !ok || id <= 0 { + writeErr(w, 400, "bad task id") + return db.TargetScopeRef{}, false + } + if _, ok := s.m.Task(strconv.FormatInt(id, 10)); !ok { + writeErr(w, 404, "task not found") + return db.TargetScopeRef{}, false + } + return db.TargetScopeRef{TaskID: id}, true +} + +func (s *Server) taskListTargetPins(w http.ResponseWriter, r *http.Request) { + if scope, ok := s.taskPinScope(w, r); ok { + s.listTargetPins(w, scope) + } +} + +func (s *Server) taskRevokeTargetPin(w http.ResponseWriter, r *http.Request) { + if scope, ok := s.taskPinScope(w, r); ok { + s.revokeTargetPin(w, r, scope) + } +} + +func (s *Server) convListTargetPins(w http.ResponseWriter, r *http.Request) { + if _, c, ok := s.convByID(w, r); ok { + s.listTargetPins(w, db.TargetScopeRef{ConversationID: c.ID}) + } +} + +func (s *Server) convRevokeTargetPin(w http.ResponseWriter, r *http.Request) { + if _, c, ok := s.convByID(w, r); ok { + s.revokeTargetPin(w, r, db.TargetScopeRef{ConversationID: c.ID}) + } +} + +// listTargetPins returns every pin of the scope, in any state, newest first. +func (s *Server) listTargetPins(w http.ResponseWriter, scope db.TargetScopeRef) { + pg := s.pg(w) + if pg == nil { + return + } + pins, err := pg.ListTargetPins(scope) + if err != nil { + writeErr(w, 500, err.Error()) + return + } + writeJSON(w, 200, map[string]any{"pins": pins}) +} + +// revokeTargetPin revokes pin {pid} of the scope. The next call to that host +// asks again. Revoking a pin that is no longer active changes nothing. +func (s *Server) revokeTargetPin(w http.ResponseWriter, r *http.Request, scope db.TargetScopeRef) { + pg := s.pg(w) + if pg == nil { + return + } + pid, ok := pathInt(r, "pid") + if !ok { + writeErr(w, 400, "bad pin id") + return + } + pin, err := pg.RevokeTargetPin(scope, pid) + if err != nil { + writeErr(w, 500, err.Error()) + return + } + if pin == nil { + writeErr(w, 404, "pin not found") + return + } + writeJSON(w, 200, pin) +} + // --- tool-config (global tool interception scope) --- // interceptGetToolConfig returns the list of tool names that are currently diff --git a/server/manager.go b/server/manager.go index 59c5443..433d237 100644 --- a/server/manager.go +++ b/server/manager.go @@ -13,6 +13,7 @@ import ( "strconv" "strings" "sync" + "sync/atomic" "time" "github.com/Autumn-27/artex/agent" @@ -72,6 +73,10 @@ type Task struct { notify chan struct{} lifecycleMu sync.RWMutex llmMu sync.RWMutex + // targetGate is the task's opt-in to the approval gate for asset targets. + // Atomic: it is switched through the API while tool calls read it. + targetGate atomic.Bool + interceptor *intercept.Interceptor // owner of the target gate, handed to the scope // pendingTriggers accumulates the concrete changes (worker done / finding) that // fired planning rounds since the last one consumed them. The debounce coalesces @@ -397,6 +402,8 @@ func NewManager(dir, proxyAddr string) (*Manager, error) { log.Printf("[llmusage] create table: %v", err) } m := &Manager{dir: dir, pg: pg, assets: pg.Assets(), tasks: map[string]*Task{}, interceptor: intercept.New(pg)} + // Behind the global egress proxy a target approval is name-only. + m.interceptor.SetUpstreamProxyCheck(func() bool { return m.GlobalProxy() != "" }) if proxyAddr != "" { tr, err := traffic.Open(filepath.Join(dir, "traffic"), proxyAddr) if err != nil { @@ -852,7 +859,7 @@ func unixNanoOrZero(t *time.Time) int64 { } func taskFromPG(pt *pgdb.Task, store *pgdb.ExplorationStore, ic *intercept.Interceptor) *Task { - return &Task{ + t := &Task{ ID: strconv.FormatInt(pt.ID, 10), ExpID: pt.ExplorationID, Name: pt.Name, CategoryID: cloneInt64Ptr(pt.CategoryID), CategoryName: pt.CategoryName, @@ -871,14 +878,18 @@ func taskFromPG(pt *pgdb.Task, store *pgdb.ExplorationStore, ic *intercept.Inter CoverageEnabled: pt.CoverageEnabled, FirstRunAt: unixOrZero(pt.FirstRunAt), DeadlineAt: unixOrZero(pt.DeadlineAt), Store: store, Guard: guard.NewWithInterceptor(ic), TargetGuard: guard.NewTargetOnly(ic), notify: make(chan struct{}, 1), + interceptor: ic, } + t.targetGate.Store(pt.TargetGate) + return t } // targetScope is the approval scope of this task's tool calls for the target -// gate. wait bounds how long a gated call may wait for a human decision. Tasks -// have no opt-in flag yet, so Enabled stays nil: the gate is off. +// gate. wait bounds how long a gated call may wait for a human decision. +// Enabled reads the task's opt-in flag on every call, so a switch through the +// API takes effect on the next tool call. func (t *Task) targetScope(wait time.Duration) intercept.TargetScope { - return intercept.TargetScope{TaskID: t.ID, WaitBound: wait} + return intercept.TargetScope{TaskID: t.ID, Enabled: t.targetGate.Load, WaitBound: wait, Interceptor: t.interceptor} } // UpdateTaskMetadata changes list-only task metadata without interrupting any @@ -907,6 +918,7 @@ func (m *Manager) UpdateTaskMetadata(taskID string, patch pgdb.TaskPatch) (*Task // Live switch (#86): runRound reads the flag per round, so the next // planning round (edge or heartbeat) already respects the new value. task.PlanQuietMode = updated.PlanQuietMode + task.targetGate.Store(updated.TargetGate) return task, nil } diff --git a/server/server.go b/server/server.go index a1f40d9..0a3f3fe 100644 --- a/server/server.go +++ b/server/server.go @@ -715,6 +715,8 @@ func (s *Server) Handler() http.Handler { mux.HandleFunc("PUT /api/tasks/{id}/intercept-rules/{rid}", s.taskInterceptUpdateRule) mux.HandleFunc("DELETE /api/tasks/{id}/intercept-rules/{rid}", s.taskInterceptDeleteRule) mux.HandleFunc("POST /api/tasks/{id}/intercept-rules/{rid}/toggle", s.taskInterceptToggleRule) + mux.HandleFunc("GET /api/tasks/{id}/target-pins", s.taskListTargetPins) + mux.HandleFunc("POST /api/tasks/{id}/target-pins/{pid}/revoke", s.taskRevokeTargetPin) mux.HandleFunc("POST /api/tasks/control/batch", s.controlTasksBatch) mux.HandleFunc("GET /api/task-archives", s.listTaskArchives) mux.HandleFunc("GET /api/task-archives/{id}", s.getTaskArchive) @@ -860,6 +862,8 @@ func (s *Server) Handler() http.Handler { mux.HandleFunc("POST /api/conversations/{id}/messages", s.pgSendConversationMessage) mux.HandleFunc("POST /api/conversations/{id}/stop", s.pgStopConversation) mux.HandleFunc("GET /api/conversations/{id}/messages/{seq}", s.pgConversationMsgDetail) + mux.HandleFunc("GET /api/conversations/{id}/target-pins", s.convListTargetPins) + mux.HandleFunc("POST /api/conversations/{id}/target-pins/{pid}/revoke", s.convRevokeTargetPin) mux.HandleFunc("GET /api/agents", s.pgListAgents) mux.HandleFunc("POST /api/agents", s.pgCreateAgent) @@ -1584,6 +1588,7 @@ type createTaskReq struct { LLMRotate *bool `json:"llm_rotate"` // task-level timeout (seconds); 0/omitted = no limit PlanHeartbeatSeconds int `json:"plan_heartbeat_seconds"` // planner heartbeat interval (seconds); 0/omitted = default 600 (10min); floor = default = 600, lower values are raised to 600 PlanQuietMode *bool `json:"plan_quiet_mode,omitempty"` // planner quiet-mode switch (#86); omitted = default on; false = edge rounds always plan + TargetGate *bool `json:"target_gate,omitempty"` // approval gate for asset targets, opt-in; omitted/null = off SeedFirstIntent *bool `json:"seed_first_intent,omitempty"` // issue one seed intent at creation (content = description + goal) so a worker can start right away without waiting for the first planner round; omitted/null = default off, the standard plan-then-execute flow. Pass true explicitly to enable (for CTF tasks often solved in one work, this skips the pre-start planner round). CoverageEnabled *bool `json:"coverage_enabled,omitempty"` // asset coverage feature; omitted/null = default on (true). false = disables coverage calculation/display/auto-accumulated scope + hides add_task_scope/list_untested_assets. company links are unaffected. // InterceptRules: task-level asset interception/allow rules (entered at creation, stored in task_intercept_rules, never the global table). @@ -1660,6 +1665,7 @@ func (s *Server) createTask(w http.ResponseWriter, r *http.Request) { CoverageEnabled: req.CoverageEnabled, InterceptRules: interceptRules, LLMRotate: req.LLMRotate, + TargetGate: req.TargetGate, }) if err != nil { if errors.Is(err, db.ErrTaskCategoryInvalid) || errors.Is(err, db.ErrTaskCategoryNotFound) { diff --git a/server/target_gate_flag_test.go b/server/target_gate_flag_test.go new file mode 100644 index 0000000..9031298 --- /dev/null +++ b/server/target_gate_flag_test.go @@ -0,0 +1,553 @@ +package server + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "net/http/httptest" + "slices" + "strconv" + "strings" + "testing" + "time" + + "github.com/Autumn-27/artex/agent" + "github.com/Autumn-27/artex/db" + "github.com/Autumn-27/artex/intercept" +) + +// Target-gate opt-in flag, pins and decisions over the real mux and the real +// auth. The server context is already cancelled, so a task created through +// these tests never starts an engine. + +type gateServer struct { + s *Server + m *Manager + token string +} + +func newGateServer(t *testing.T) *gateServer { + t.Helper() + skipWithoutDSN(t) + m, err := NewManager(t.TempDir(), "") + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = m.Close() }) + ctx, cancel := context.WithCancel(context.Background()) + cancel() + s := New(ctx, m, t.TempDir(), t.TempDir(), t.TempDir()) + token, err := signJWT(s.jwtKey) + if err != nil { + t.Fatal(err) + } + return &gateServer{s: s, m: m, token: token} +} + +// do sends an authenticated request. httptest sets the peer to 192.0.2.1. +func (g *gateServer) do(t *testing.T, method, path, body string) (int, map[string]any) { + t.Helper() + req := httptest.NewRequest(method, path, strings.NewReader(body)) + req.Header.Set("Authorization", "Bearer "+g.token) + rec := httptest.NewRecorder() + g.s.Handler().ServeHTTP(rec, req) + var out map[string]any + _ = json.Unmarshal(rec.Body.Bytes(), &out) + return rec.Code, out +} + +// newTask creates a task through the manager and removes it, its pins and its +// approval rows when the test ends. +func (g *gateServer) newTask(t *testing.T, opts db.TaskCreateOptions) *Task { + t.Helper() + created, err := g.m.CreateTaskWithOptions("target gate flag", "flag, pins and decisions", opts) + if err != nil { + t.Fatal(err) + } + g.cleanupTask(t, created.ID) + task, ok := g.m.Task(created.ID) + if !ok { + t.Fatal("task not found after creation") + } + return task +} + +func (g *gateServer) cleanupTask(t *testing.T, id string) { + t.Helper() + t.Cleanup(func() { + n, _ := strconv.ParseInt(id, 10, 64) + _ = g.m.pg.DeleteTask(n) + _, _ = g.m.pg.Exec(`DELETE FROM target_pins WHERE task_id=$1`, n) + _, _ = g.m.pg.Exec(`DELETE FROM intercept_pending WHERE task_id=$1`, id) + }) +} + +func taskRef(t *testing.T, task *Task) db.TargetScopeRef { + t.Helper() + id, err := strconv.ParseInt(task.ID, 10, 64) + if err != nil { + t.Fatal(err) + } + return db.TargetScopeRef{TaskID: id} +} + +// openTargetRow opens a target approval request showing ips. +func openTargetRow(t *testing.T, pg *db.DB, scope db.TargetScopeRef, host string, ips ...string) int64 { + t.Helper() + id, created, err := pg.FindOrCreateTargetRequest(db.TargetRequest{ + Scope: scope, Host: host, Port: 443, AgentName: "worker", ToolName: "WebFetch", Reason: "target " + host + " is not approved", + Audit: &db.InterceptAudit{Correlation: "unavailable", InitialAction: "ask", ExecutionStatus: "not_started", + Targets: []db.InterceptTarget{{Host: host, Port: 443, ResolvedIPs: ips, Resolver: "system", ResolvedAt: time.Now().UTC(), Outcome: "requested"}}}, + }) + if err != nil || !created { + t.Fatalf("open target request: id=%d created=%v err=%v", id, created, err) + } + return id +} + +// The flag is read on every tool call: a switch through the API is seen by a +// scope that was handed out before it. +func TestTargetGateLiveSwitch(t *testing.T) { + g := newGateServer(t) + task := g.newTask(t, db.TaskCreateOptions{}) + scope := task.targetScope(intercept.DefaultTargetWaitBound) + if scope.Enabled == nil || scope.Enabled() { + t.Fatal("a new task must have the gate off") + } + if scope.Interceptor != g.m.interceptor { + t.Fatal("the task scope does not carry the manager's interceptor") + } + if code, body := g.do(t, "GET", "/api/tasks/"+task.ID, ""); code != 200 || body["target_gate"] != false { + t.Fatalf("task detail: %d %v, want target_gate false", code, body) + } + + code, body := g.do(t, "PATCH", "/api/tasks/"+task.ID, `{"target_gate":true}`) + if code != 200 || body["target_gate"] != true { + t.Fatalf("switch on: %d %v", code, body) + } + if !scope.Enabled() { + t.Fatal("the scope handed out before the switch still reports the gate off") + } + id, _ := strconv.ParseInt(task.ID, 10, 64) + if stored, err := g.m.pg.GetTask(id); err != nil || stored == nil || !stored.TargetGate { + t.Fatalf("stored task = %+v, %v; want the gate on", stored, err) + } + // Another metadata patch leaves the flag alone. + if code, body := g.do(t, "PATCH", "/api/tasks/"+task.ID, `{"name":"renamed"}`); code != 200 || body["target_gate"] != true || !scope.Enabled() { + t.Fatalf("rename: %d %v, want the gate still on", code, body) + } + if code, body := g.do(t, "PATCH", "/api/tasks/"+task.ID, `{"target_gate":false}`); code != 200 || body["target_gate"] != false || scope.Enabled() { + t.Fatalf("switch off: %d %v enabled=%v", code, body, scope.Enabled()) + } + if code, body := g.do(t, "PATCH", "/api/tasks/"+task.ID, `{}`); code != 400 || !strings.Contains(fmt.Sprint(body["error"]), "target_gate") { + t.Fatalf("empty patch: %d %v, want 400 naming target_gate", code, body) + } +} + +const gateTestGoal = "Probe the lab target at http://127.0.0.1:9/v1 for unauthenticated data exposure; record one record as evidence and stop." + +// The flag is set on the existing create path. Omitted means off, and a task +// loaded from the database keeps what was stored. +func TestCreateTaskAcceptsTargetGate(t *testing.T) { + g := newGateServer(t) + for _, tc := range []struct { + field string + want bool + }{ + {`,"target_gate":true`, true}, + {`,"target_gate":false`, false}, + {``, false}, + } { + code, body := g.do(t, "POST", "/api/tasks", `{"description":"gate create","goal":"`+gateTestGoal+`"`+tc.field+`}`) + id, _ := body["id"].(string) + if code != 201 || id == "" { + t.Fatalf("create %q: %d %v", tc.field, code, body) + } + g.cleanupTask(t, id) + task, ok := g.m.Task(id) + if !ok || body["target_gate"] != tc.want || task.targetScope(0).Enabled() != tc.want { + t.Fatalf("create %q: target_gate=%v, want %v", tc.field, body["target_gate"], tc.want) + } + n, _ := strconv.ParseInt(id, 10, 64) + stored, err := g.m.pg.GetTask(n) + if err != nil || stored == nil || stored.TargetGate != tc.want { + t.Fatalf("create %q: stored %+v, %v; want %v", tc.field, stored, err, tc.want) + } + if loaded := taskFromPG(stored, g.m.pg.Exploration(stored.ExplorationID), g.m.interceptor); loaded.targetScope(0).Enabled() != tc.want { + t.Fatalf("create %q: a task loaded from the database has the gate %v", tc.field, !tc.want) + } + } +} + +// No inheritance: a task spawned by the orchestration agent starts with the +// gate off, even from a gated parent. +func TestSpawnedTaskStartsWithGateOff(t *testing.T) { + g := newGateServer(t) + on := true + parent := g.newTask(t, db.TaskCreateOptions{TargetGate: &on}) + if !parent.targetScope(0).Enabled() { + t.Fatal("control: the parent was expected to have the gate on") + } + in, _ := json.Marshal(map[string]any{"description": "spawned child", "goal": gateTestGoal, "parent_ref": parent.ID, "source_task_ids": []string{parent.ID}}) + res, err := g.s.toolSpawnTask().Call(context.Background(), in, nil) + if err != nil { + t.Fatal(err) + } + childID, found := strings.CutPrefix(res.Flatten(), "task created: ") + if !found { + t.Fatalf("spawn_task result: %q", res.Flatten()) + } + g.cleanupTask(t, childID) + child, ok := g.m.Task(childID) + if !ok || child.ParentRef != parent.ID { + t.Fatalf("spawned task = %+v (found=%v), want a child of %s", child, ok, parent.ID) + } + if child.targetScope(0).Enabled() { + t.Fatal("the spawned task inherited the gate") + } + n, _ := strconv.ParseInt(childID, 10, 64) + if stored, err := g.m.pg.GetTask(n); err != nil || stored == nil || stored.TargetGate { + t.Fatalf("stored child = %+v, %v; want the gate off", stored, err) + } + if pins, err := g.m.pg.ListTargetPins(db.TargetScopeRef{TaskID: n}); err != nil || len(pins) != 0 { + t.Fatalf("the spawned task has pins: %+v, %v", pins, err) + } +} + +// The conversation flag is set on create and on patch, and the turn that starts +// after a switch runs with the new value. +func TestConversationTargetGateOverHTTP(t *testing.T) { + g := newGateServer(t) + turnScope := func(id int64) intercept.TargetScope { + t.Helper() + conv, err := g.m.pg.GetConversation(id) + if err != nil || conv == nil { + t.Fatalf("conversation %d: %+v, %v", id, conv, err) + } + probe := newScopeProbeProvider() + g.s.cfgMu.Lock() + g.s.chatAgent = agent.NewChatAgent(probe, "test-model", g.m.dir, nil, 10000) + g.s.cfgMu.Unlock() + // The server context of this fixture is cancelled; the turn gets its own. + ctx, cancel := context.WithCancelCause(intercept.WithConvID(context.Background(), id)) + g.s.runConversationTurn(ctx, cancel, conv, "hello", g.s.convBusyKey(id)) + got := probe.first(t) + if !got.ok || got.scope.ConversationID != id || got.scope.Enabled == nil { + t.Fatalf("turn scope = %+v (present=%v)", got.scope, got.ok) + } + return got.scope + } + create := func(body string) int64 { + t.Helper() + code, out := g.do(t, "POST", "/api/conversations", body) + id, _ := out["id"].(float64) + if code != 200 || id == 0 { + t.Fatalf("create conversation: %d %v", code, out) + } + t.Cleanup(func() { _ = g.m.pg.DeleteConversation(int64(id)) }) + return int64(id) + } + + plain := create(`{"agent_key":"auto","title":"gate off"}`) + if scope := turnScope(plain); scope.Enabled() { + t.Fatal("a conversation created without the flag runs with the gate on") + } + gated := create(`{"agent_key":"auto","title":"gate on","target_gate":true}`) + scope := turnScope(gated) + if !scope.Enabled() || scope.Interceptor != g.m.interceptor { + t.Fatalf("gated conversation: enabled=%v, want the gate on and the manager's interceptor", scope.Enabled()) + } + path := "/api/conversations/" + strconv.FormatInt(gated, 10) + if code, out := g.do(t, "PATCH", path, `{"target_gate":false}`); code != 200 || out["target_gate"] != false { + t.Fatalf("switch off: %d %v", code, out) + } + if turnScope(gated).Enabled() { + t.Fatal("the turn after the switch still runs with the gate on") + } + if code, out := g.do(t, "PATCH", path, `{"title":"renamed","target_gate":true}`); code != 200 || out["target_gate"] != true || out["title"] != "renamed" { + t.Fatalf("switch on with a rename: %d %v", code, out) + } + if code, out := g.do(t, "PATCH", path, `{}`); code != 400 || !strings.Contains(fmt.Sprint(out["error"]), "target_gate") { + t.Fatalf("empty patch: %d %v, want 400 naming target_gate", code, out) + } +} + +// The decide endpoint keeps its request and response shape. On a target row +// "allowed" pins the displayed addresses. Every human decision, on a target row +// or on a classic ask, records the subject of the session token and the remote +// address; decided_by stays the actor class. +func TestInterceptDecideRecordsApproverOverHTTP(t *testing.T) { + g := newGateServer(t) + task := g.newTask(t, db.TaskCreateOptions{}) + scope := taskRef(t, task) + pg := g.m.pg + decide := func(id int64, decision string) int { + t.Helper() + code, _ := g.do(t, "POST", fmt.Sprintf("/api/intercept/pending/%d/decide", id), `{"decision":"`+decision+`"}`) + return code + } + + target := openTargetRow(t, pg, scope, "app.example.com", "203.0.113.7") + if code := decide(target, "allowed"); code != 200 { + t.Fatalf("decide target row: %d", code) + } + pins, err := pg.ListTargetPins(scope) + if err != nil || len(pins) != 1 || pins[0].DecidedBy != "ARTEX" || pins[0].DecidedFrom != "192.0.2.1" || + !slices.Equal(pins[0].PinnedIPs, []string{"203.0.113.7"}) || pins[0].InterceptID == nil || *pins[0].InterceptID != target { + t.Fatalf("pins = %+v, %v; want one pin with the approver and the address", pins, err) + } + row, err := pg.GetInterceptDetail(target) + if err != nil || row == nil || row.Status != "allowed" || row.DecidedBy != "human" || row.Audit.Approver != "ARTEX" || row.Audit.ApproverAddr != "192.0.2.1" { + t.Fatalf("target row = %+v, %v", row, err) + } + if code := decide(target, "denied"); code != 409 { + t.Fatalf("second decision: %d, want 409", code) + } + + denied := openTargetRow(t, pg, scope, "other.example.com", "203.0.113.8") + if code := decide(denied, "denied"); code != 200 { + t.Fatalf("deny target row: %d", code) + } + if pins, err := pg.ListTargetPins(scope); err != nil || len(pins) != 1 { + t.Fatalf("pins after a denial = %+v, %v; want still one", pins, err) + } + + classic, err := pg.CreateInterceptPending(0, 0, task.ID, "worker", "Bash", []byte(`{"command":"id"}`), "classic ask", + &db.InterceptAudit{Correlation: "exact", InitialAction: "ask", ExecutionStatus: "not_started"}) + if err != nil { + t.Fatal(err) + } + if code := decide(classic, "allowed"); code != 200 { + t.Fatalf("decide classic ask: %d", code) + } + row, err = pg.GetInterceptDetail(classic) + if err != nil || row == nil || row.Status != "allowed" || row.DecidedBy != "human" || row.Audit == nil || + row.Audit.Approver != "ARTEX" || row.Audit.ApproverAddr != "192.0.2.1" || row.Audit.EffectiveAction != "allow" { + t.Fatalf("classic row = %+v, %v; want the approver and the address in its audit", row, err) + } + if pins, err := pg.ListTargetPins(scope); err != nil || len(pins) != 1 { + t.Fatalf("pins after a classic approval = %+v, %v; want still one", pins, err) + } + + // Target rows are listed under their own decision source. + code, list := g.do(t, "GET", "/api/intercept/task/"+task.ID+"?decision_source=target&page=1&size=50", "") + if items, _ := list["items"].([]any); code != 200 || len(items) != 2 { + t.Fatalf("decision_source=target: %d %v, want the two target rows", code, list) + } + if code, _ := g.do(t, "GET", "/api/intercept/history?decision_source=nonsense&page=1", ""); code != 400 { + t.Fatalf("unknown decision_source: %d, want 400", code) + } +} + +func TestJWTSubject(t *testing.T) { + key := []byte("0123456789abcdef0123456789abcdef") + token, err := signJWT(key) + if err != nil { + t.Fatal(err) + } + if got := jwtSubject(token, key); got != "ARTEX" { + t.Fatalf("subject = %q, want ARTEX", got) + } + if got := jwtSubject(token, []byte("another key, another key, another")); got != "" { + t.Fatalf("subject of a token signed with another key = %q, want none", got) + } + if got := jwtSubject("", key); got != "" { + t.Fatalf("subject of no token = %q", got) + } +} + +// Pin routes: list and revoke, scoped to one existing task or conversation. +func TestTargetPinRoutes(t *testing.T) { + g := newGateServer(t) + pg := g.m.pg + task, other := g.newTask(t, db.TaskCreateOptions{}), g.newTask(t, db.TaskCreateOptions{}) + conv, err := pg.CreateConversation("auto", "pin routes", nil) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { + _ = pg.DeleteConversation(conv.ID) + _, _ = pg.Exec(`DELETE FROM target_pins WHERE conversation_id=$1`, conv.ID) + }) + pin := func(scope db.TargetScopeRef, host string) int64 { + t.Helper() + p, ok, err := pg.ResolveTargetRequest(openTargetRow(t, pg, scope, host, "203.0.113.7"), true, "ARTEX", "192.0.2.1") + if err != nil || !ok || p == nil { + t.Fatalf("pin %s: %+v ok=%v err=%v", host, p, ok, err) + } + return p.ID + } + taskPin := pin(taskRef(t, task), "app.example.com") + otherPin := pin(taskRef(t, other), "app.example.com") + convPin := pin(db.TargetScopeRef{ConversationID: conv.ID}, "chat.example.com") + + listed := func(path string) []any { + t.Helper() + code, body := g.do(t, "GET", path, "") + pins, ok := body["pins"].([]any) + if code != 200 || !ok { + t.Fatalf("GET %s: %d %v", path, code, body) + } + return pins + } + for _, tc := range []struct { + base string + id int64 + host string + }{ + {"/api/tasks/" + task.ID, taskPin, "app.example.com"}, + {"/api/conversations/" + strconv.FormatInt(conv.ID, 10), convPin, "chat.example.com"}, + } { + pins := listed(tc.base + "/target-pins") + if len(pins) != 1 { + t.Fatalf("%s: %d pins, want 1", tc.base, len(pins)) + } + got := pins[0].(map[string]any) + if got["id"] != float64(tc.id) || got["host"] != tc.host || got["status"] != "active" || got["decided_by"] != "ARTEX" || got["name_only_effective"] != false { + t.Fatalf("%s: listed pin = %v", tc.base, got) + } + revoke := fmt.Sprintf("%s/target-pins/%d/revoke", tc.base, tc.id) + code, body := g.do(t, "POST", revoke, "") + if code != 200 || body["status"] != "revoked" || body["revoked_reason"] != "operator" { + t.Fatalf("POST %s: %d %v", revoke, code, body) + } + at := body["revoked_at"] + // Revoking again changes nothing. + if code, body := g.do(t, "POST", revoke, ""); code != 200 || body["status"] != "revoked" || body["revoked_at"] != at { + t.Fatalf("second POST %s: %d %v", revoke, code, body) + } + if pins := listed(tc.base + "/target-pins"); len(pins) != 1 || pins[0].(map[string]any)["status"] != "revoked" { + t.Fatalf("%s after the revoke: %v, want the pin kept as revoked", tc.base, pins) + } + // A pin of another scope is not reachable through this one. + if code, _ := g.do(t, "POST", fmt.Sprintf("%s/target-pins/%d/revoke", tc.base, otherPin), ""); code != 404 { + t.Fatalf("revoke of a foreign pin through %s: %d, want 404", tc.base, code) + } + } + if active, err := pg.ActiveTargetPins(taskRef(t, other), "app.example.com", 443); err != nil || len(active) != 1 { + t.Fatalf("the pin of the other task = %+v, %v; want still active", active, err) + } + for _, path := range []string{"/api/tasks/999999999/target-pins", "/api/conversations/999999999/target-pins"} { + if code, _ := g.do(t, "GET", path, ""); code != 404 { + t.Fatalf("GET %s: %d, want 404", path, code) + } + } + if code, _ := g.do(t, "POST", "/api/tasks/"+task.ID+"/target-pins/notanumber/revoke", ""); code != 400 { + t.Fatalf("revoke with a bad pin id: %d, want 400", code) + } +} + +// A worker that loses its claim while it waits in the gate stops waiting, and +// the request stays open: it is host-level, not owned by the call. The worker +// that took the intent over attaches to the same request. +func TestClaimLostEndsGateWait(t *testing.T) { + g := newGateServer(t) + task := g.newTask(t, db.TaskCreateOptions{}) + pg, ic := g.m.pg, g.m.interceptor + store := task.Store + intent, err := store.AddIntent(map[string]any{"summary": "leased"}, 1, nil, "planner") + if err != nil { + t.Fatal(err) + } + token, err := store.ClaimIntentToken(intent, "work#1") + if err != nil || token == "" { + t.Fatalf("claim: token=%q err=%v", token, err) + } + ctx, cancel := context.WithCancelCause(intercept.WithTaskContext(context.Background(), task.ID, "work#1", nil)) + defer cancel(nil) + stop := startClaimHeartbeat(store, 20*time.Millisecond, "work#1", intent, token, cancel) + defer stop() + + scope := task.targetScope(time.Minute) + targets := []intercept.TargetRequest{{Host: "app.example.com", Port: 443, ResolvedIPs: []string{"203.0.113.7"}, Resolver: "system", ResolvedAt: time.Now().UTC(), Reason: "no_pin"}} + input := []byte(`{"url":"https://app.example.com/"}`) + stale := make(chan intercept.TargetResult, 1) + go func() { stale <- ic.RequestTargets(ctx, scope, "WebFetch", input, targets, nil) }() + openRows := func() []int64 { + t.Helper() + rows, err := pg.Query(`SELECT id FROM intercept_pending WHERE task_id=$1 AND status='pending' AND target_key IS NOT NULL ORDER BY id`, task.ID) + if err != nil { + t.Fatal(err) + } + defer rows.Close() + var ids []int64 + for rows.Next() { + var id int64 + if err := rows.Scan(&id); err != nil { + t.Fatal(err) + } + ids = append(ids, id) + } + return ids + } + waitUntil(t, "the gate request", func() bool { return len(openRows()) == 1 }) + request := openRows()[0] + select { + case res := <-stale: + t.Fatalf("the call returned %+v while its claim was live", res) + case <-time.After(200 * time.Millisecond): + } + + // Takeover, as in TestClaimHeartbeatRenewsThenCancelsOnTakeover. + var reclaimed int64 + for attempt := 0; attempt < 200 && reclaimed == 0; attempt++ { + if _, err := pg.Exec(`UPDATE exploration_nodes SET last_heartbeat_at = now() - $2 * INTERVAL '1 second' WHERE id=$1`, + intent, (db.IntentClaimTTL + time.Second).Seconds()); err != nil { + t.Fatal(err) + } + if reclaimed, err = store.ResetRunningIntents(); err != nil { + t.Fatal(err) + } + } + if reclaimed != 1 { + t.Fatalf("could not take the claim over: reclaimed=%d", reclaimed) + } + if other, err := store.ClaimIntentToken(intent, "work#2"); err != nil || other == "" { + t.Fatalf("re-claim: token=%q err=%v", other, err) + } + + var res intercept.TargetResult + select { + case res = <-stale: + case <-time.After(10 * time.Second): + t.Fatal("the gate wait did not end after the claim was lost") + } + if cause := context.Cause(ctx); !errors.Is(cause, agent.AbortWorkClaimLost) { + t.Fatalf("cancel cause=%v, want AbortWorkClaimLost", cause) + } + if res.Code != intercept.TargetApprovalPending || res.PendingID != request { + t.Fatalf("stale call: %+v, want %s on request %d", res, intercept.TargetApprovalPending, request) + } + if row, err := pg.GetInterceptPending(request); err != nil || row == nil || row.Status != "pending" || row.DecidedBy != "" { + t.Fatalf("request after the claim was lost = %+v, %v; want still open", row, err) + } + + // The second claimer re-runs the intent, hits the same host and attaches. + next := make(chan intercept.TargetResult, 1) + go func() { + next <- ic.RequestTargets(intercept.WithTaskContext(context.Background(), task.ID, "work#2", nil), scope, "WebFetch", input, targets, nil) + }() + select { + case res := <-next: + t.Fatalf("the second claimer returned %+v before a decision", res) + case <-time.After(200 * time.Millisecond): + } + if ids := openRows(); !slices.Equal(ids, []int64{request}) { + t.Fatalf("open requests = %v, want the second claimer on the same request %d", ids, request) + } + if err := ic.Decide(request, true, "ARTEX", "192.0.2.1"); err != nil { + t.Fatal(err) + } + select { + case res := <-next: + if res.Code != "" { + t.Fatalf("second claimer after the approval: %+v", res) + } + case <-time.After(10 * time.Second): + t.Fatal("the second claimer is still waiting after the approval") + } + if pins, err := pg.ListTargetPins(taskRef(t, task)); err != nil || len(pins) != 1 { + t.Fatalf("pins = %+v, %v; want one", pins, err) + } +} diff --git a/server/target_gate_scope_test.go b/server/target_gate_scope_test.go index ac75f8b..a8d0552 100644 --- a/server/target_gate_scope_test.go +++ b/server/target_gate_scope_test.go @@ -335,8 +335,11 @@ func TestEveryAgentKindCarriesTargetScope(t *testing.T) { if got.scope.WaitBound != site.wait { t.Fatalf("wait bound = %v, want %v", got.scope.WaitBound, site.wait) } - if got.scope.Enabled != nil { - t.Fatal("the gate must be off: tasks have no opt-in flag yet") + if got.scope.Enabled == nil || got.scope.Enabled() { + t.Fatal("the gate must be off: the task did not opt in") + } + if got.scope.Interceptor != f.m.interceptor { + t.Fatal("the scope does not carry the interceptor that owns the gate") } }) } @@ -348,8 +351,11 @@ func TestEveryAgentKindCarriesTargetScope(t *testing.T) { if got.scope.WaitBound != intercept.DefaultTargetWaitBound { t.Fatalf("wait bound = %v, want %v", got.scope.WaitBound, intercept.DefaultTargetWaitBound) } - if got.scope.Enabled != nil { - t.Fatal("the gate must be off: conversations have no opt-in flag yet") + if got.scope.Enabled == nil || got.scope.Enabled() { + t.Fatal("the gate must be off: the conversation did not opt in") + } + if got.scope.Interceptor != f.m.interceptor { + t.Fatal("the scope does not carry the interceptor that owns the gate") } }) } diff --git a/server/task_metadata.go b/server/task_metadata.go index 0cd4d21..7bcd738 100644 --- a/server/task_metadata.go +++ b/server/task_metadata.go @@ -28,6 +28,7 @@ func (s *Server) updateTaskMetadata(w http.ResponseWriter, r *http.Request) { Name *string `json:"name"` Pinned *bool `json:"pinned"` PlanQuietMode *bool `json:"plan_quiet_mode"` + TargetGate *bool `json:"target_gate"` } if err := decode(r, &request); err != nil { var tooLarge *http.MaxBytesError @@ -38,8 +39,8 @@ func (s *Server) updateTaskMetadata(w http.ResponseWriter, r *http.Request) { } return } - if request.Name == nil && request.Pinned == nil && request.PlanQuietMode == nil { - writeErr(w, http.StatusBadRequest, "at least one of name, pinned or plan_quiet_mode must be provided") + if request.Name == nil && request.Pinned == nil && request.PlanQuietMode == nil && request.TargetGate == nil { + writeErr(w, http.StatusBadRequest, "at least one of name, pinned, plan_quiet_mode or target_gate must be provided") return } if request.Name != nil { @@ -54,7 +55,7 @@ func (s *Server) updateTaskMetadata(w http.ResponseWriter, r *http.Request) { } request.Name = &name } - task, err := s.m.UpdateTaskMetadata(taskID, db.TaskPatch{Name: request.Name, Pinned: request.Pinned, PlanQuietMode: request.PlanQuietMode}) + task, err := s.m.UpdateTaskMetadata(taskID, db.TaskPatch{Name: request.Name, Pinned: request.Pinned, PlanQuietMode: request.PlanQuietMode, TargetGate: request.TargetGate}) if err != nil { writeErr(w, http.StatusInternalServerError, err.Error()) return