krow_backend/go-api/internal/tools/confirm_store.go

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
}