package tools import ( "context" "encoding/json" "errors" "fmt" "time" "github.com/krow/krow-backend/go-api/internal/repo" ) // PostgresStore keeps pending confirmations in agent_confirmations. // // The store MemoryStore should have been. A confirmation is asked for in one // request and answered in another, and there is no reason those two requests // reach the same process — behind two replicas, an in-process store refuses // roughly half of all approvals, and refuses them in a way the person clicking // cannot act on and cannot even see the cause of. type PostgresStore struct { db repo.Querier } // NewPostgresStore builds a store over a pool or transaction. func NewPostgresStore(db repo.Querier) *PostgresStore { return &PostgresStore{db: db} } var _ Store = (*PostgresStore)(nil) // Issue records a pending confirmation. // // Fails loudly. The registry treats an Issue error as a reason to refuse the // tool outright, because showing somebody a question whose answer will be // discarded is worse than saying the tool is unavailable. func (s *PostgresStore) Issue(ctx context.Context, b binding, c *Confirmation) error { if b.OrgID == "" { // I5. There is no tenant-less confirmation, and writing one would put a // row in the table that no caller could ever legitimately resolve. return errors.New("tools: a confirmation needs an organization") } if c == nil || c.Token == "" { return errors.New("tools: a confirmation needs a token") } payload, err := json.Marshal(c) if err != nil { return fmt.Errorf("tools: encoding a confirmation: %w", err) } inputs := b.Inputs if len(inputs) == 0 { inputs = json.RawMessage("{}") } _, err = s.db.Exec(ctx, ` INSERT INTO agent_confirmations (token, org_id, user_id, run_id, tool, inputs_hash, inputs, agent_id, payload, expires_at) VALUES ($1, $2::uuid, $3, $4, $5, $6, $7::jsonb, $8, $9::jsonb, $10)`, c.Token, b.OrgID, nullableUUID(b.UserID), b.RunID, b.Tool, b.InputsHash, string(inputs), b.AgentID, payload, c.ExpiresAt) if err != nil { return fmt.Errorf("tools: recording a confirmation: %w", err) } return nil } // Resolve consumes a token if it authorises this exact call. // // The claim and the check are ONE statement, and they have to be. Reading the // row and then updating it would be the same logic with a race in the middle, // and the race is a duplicated write — precisely the failure this mechanism // exists to prevent. `consumed_at IS NULL` in the WHERE clause combined with // RETURNING means exactly one concurrent caller can take a token. // // The binding is in the same WHERE clause rather than compared afterwards, so a // token presented against a DIFFERENT call is not consumed. That matters when a // turn contains two write calls: spending the token on whichever was dispatched // first would refuse the one the person actually approved. // // run_id is stored but not matched — a resumed run has a new id by definition. // See the note in confirm.go. // // The hash comparison is SQL's rather than constant-time. The token is the // secret and it is unguessable; the hash is only reached by a caller who // already holds the token, so there is no oracle to protect here. func (s *PostgresStore) Resolve(ctx context.Context, token string, b binding) bool { if token == "" { return false } var claimed string err := s.db.QueryRow(ctx, ` UPDATE agent_confirmations SET consumed_at = now() WHERE token = $1 AND consumed_at IS NULL AND expires_at > now() AND tool = $2 AND inputs_hash = $3 AND org_id = $4::uuid AND user_id IS NOT DISTINCT FROM $5::uuid RETURNING token`, token, b.Tool, b.InputsHash, b.OrgID, nullableUUID(b.UserID), ).Scan(&claimed) // No rows is the ordinary case: no such token, already spent, expired, or // issued for a different call. Any other error is a database problem, and // the safe reading of "I could not verify this approval" is that it is not // approved. All of them refuse identically — a caller able to distinguish // "spent" from "never existed" could probe other people's approvals. return err == nil && claimed == token } // Sweep deletes confirmations that can no longer be answered. // // Housekeeping, not a security control: Resolve refuses an expired token // whether or not this has run. Returns how many rows it removed so a scheduled // caller can log something true. func (s *PostgresStore) Sweep(ctx context.Context, olderThan time.Duration) (int64, error) { tag, err := s.db.Exec(ctx, ` DELETE FROM agent_confirmations WHERE expires_at < now() - $1::interval`, olderThan.String()) if err != nil { return 0, fmt.Errorf("tools: sweeping confirmations: %w", err) } return tag.RowsAffected(), nil } // nullableUUID renders an empty principal id as SQL NULL. // // user_id is nullable and references users; an empty string would fail the cast // rather than storing "unknown", which is a real state — a run started by a // service principal has no user row behind it. func nullableUUID(s string) any { if s == "" { return nil } return s } // Redeem consumes a token and returns the call it authorised. // // One statement again, and for the same reason Resolve is: the claim and the // read have to be atomic or two concurrent redemptions both succeed. The // difference is what is checked — the caller and the expiry, but NOT the // arguments, because this is where the arguments come from. // // A row written before migration 000009 has a NULL `inputs`, and those cannot // be replayed. They are refused rather than replayed as `{}`: an empty argument // set is a different call from the one a person approved, and running it would // be worse than declining to. func (s *PostgresStore) Redeem(ctx context.Context, token string, p Principal) (Approved, bool) { if token == "" || p.OrgID == "" { return Approved{}, false } var ( tool string inputs []byte agentID string ) err := s.db.QueryRow(ctx, ` UPDATE agent_confirmations SET consumed_at = now() WHERE token = $1 AND consumed_at IS NULL AND expires_at > now() AND org_id = $2::uuid AND user_id IS NOT DISTINCT FROM $3::uuid AND inputs IS NOT NULL RETURNING tool, inputs, agent_id`, token, p.OrgID, nullableUUID(p.UserID), ).Scan(&tool, &inputs, &agentID) if err != nil { return Approved{}, false } return Approved{Tool: tool, Inputs: json.RawMessage(inputs), AgentID: agentID}, true }