agent build
This commit is contained in:
178
go-api/internal/tools/confirm_store.go
Normal file
178
go-api/internal/tools/confirm_store.go
Normal file
@@ -0,0 +1,178 @@
|
||||
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
|
||||
}
|
||||
Reference in New Issue
Block a user