From 33119f930603d48bf767ef2b1d173d27154dbd77 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Aur=C3=A9lien=20Sibiril?= <81782+aureliensibiril@users.noreply.github.com> Date: Fri, 24 Apr 2026 19:59:55 +0200 Subject: [PATCH] Scope PGCheckpointer queries by tenant GID MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Each Save and Load now derives tenant_id from the run GID and pins it in the WHERE clause. A caller that supplies an ID from another tenant fails closed instead of silently reading or overwriting cross-tenant checkpoint data. Also rejects oversize checkpoints on load as a read-side guard against a tampered or migrated row exceeding MaxCheckpointBytes. Signed-off-by: Aurélien Sibiril <81782+aureliensibiril@users.noreply.github.com> --- pkg/agent/checkpoint.go | 2 - pkg/agentruntest/agentruntest.go | 48 ++++++++++--- pkg/coredata/agent_run.go | 69 +++++++++++++++---- ...60424T120000Z.sql => 20260424T173529Z.sql} | 5 +- 4 files changed, 94 insertions(+), 30 deletions(-) rename pkg/coredata/migrations/{20260424T120000Z.sql => 20260424T173529Z.sql} (84%) diff --git a/pkg/agent/checkpoint.go b/pkg/agent/checkpoint.go index d81b2ae4e..840d3d1b3 100644 --- a/pkg/agent/checkpoint.go +++ b/pkg/agent/checkpoint.go @@ -82,8 +82,6 @@ type ( ) const ( - MaxCheckpointBytes = 10 * 1024 * 1024 - AgentStatusSuspended AgentStatus = "suspended" AgentStatusAwaitingApproval AgentStatus = "awaiting_approval" ) diff --git a/pkg/agentruntest/agentruntest.go b/pkg/agentruntest/agentruntest.go index e806a23aa..0e0e2256a 100644 --- a/pkg/agentruntest/agentruntest.go +++ b/pkg/agentruntest/agentruntest.go @@ -95,8 +95,13 @@ func PGClient(t *testing.T) *pg.Client { return sharedPGClient } -// EnsureAgentRunsTable creates the agent_runs table if it does not -// exist, using the embedded migration. +// EnsureAgentRunsTable creates the agent_runs table against the test +// database using the embedded migration, if the table is not already +// present. If the table exists with a stale schema (e.g. missing the +// FK added later), drop it manually or let the production migration +// runner apply the current version — this helper does not rewrite an +// existing table to avoid racing concurrent test processes that share +// the same database. func EnsureAgentRunsTable(t *testing.T, client *pg.Client) { t.Helper() @@ -104,25 +109,25 @@ func EnsureAgentRunsTable(t *testing.T, client *pg.Client) { ctx := context.Background() ensureTableErr = client.WithConn(ctx, func(ctx context.Context, conn pg.Querier) error { var exists bool - err := conn.QueryRow( + if err := conn.QueryRow( ctx, `SELECT EXISTS (SELECT 1 FROM information_schema.tables WHERE table_name = 'agent_runs')`, - ).Scan(&exists) - if err != nil { - return err + ).Scan(&exists); err != nil { + return fmt.Errorf("cannot check agent_runs existence: %w", err) } if exists { return nil } - // Read from the embedded migration to avoid schema drift. - ddl, err := coredata.Migrations.ReadFile("migrations/20260424T120000Z.sql") + ddl, err := coredata.Migrations.ReadFile("migrations/20260424T173529Z.sql") if err != nil { return fmt.Errorf("cannot read agent_runs migration: %w", err) } - _, err = conn.Exec(ctx, string(ddl)) - return err + if _, err := conn.Exec(ctx, string(ddl)); err != nil { + return fmt.Errorf("cannot apply agent_runs migration: %w", err) + } + return nil }) }) require.NoError(t, ensureTableErr, "cannot ensure agent_runs table") @@ -141,6 +146,8 @@ func CleanupAgentRun(client *pg.Client, id gid.GID) { } // InsertPendingRun inserts a PENDING agent run and registers cleanup. +// A placeholder organization row is created first so the agent_runs FK +// on organization_id is satisfied. func InsertPendingRun( t *testing.T, client *pg.Client, @@ -171,18 +178,37 @@ func InsertPendingRun( err = client.WithTx( context.Background(), func(ctx context.Context, tx pg.Tx) error { + if _, err := tx.Exec( + ctx, + `INSERT INTO organizations (id, tenant_id, name, created_at, updated_at) VALUES ($1, $2, $3, $4, $5)`, + orgID.String(), tenantID.String(), "test-org-"+orgID.String(), now, now, + ); err != nil { + return fmt.Errorf("cannot insert placeholder organization: %w", err) + } return run.Insert(ctx, tx, coredata.NewScope(tenantID)) }, ) require.NoError(t, err) t.Cleanup(func() { - CleanupAgentRun(client, run.ID) + cleanupOrganization(client, orgID) }) return run } +// cleanupOrganization deletes the test organization row; the agent_runs +// FK has ON DELETE CASCADE so the associated run is removed too. +func cleanupOrganization(client *pg.Client, id gid.GID) { + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + _ = client.WithConn(ctx, func(ctx context.Context, conn pg.Querier) error { + _, err := conn.Exec(ctx, "DELETE FROM organizations WHERE id = $1", id.String()) + return err + }) +} + // LoadAgentRun loads an agent run by ID, failing the test on error. func LoadAgentRun(t *testing.T, client *pg.Client, id gid.GID) coredata.AgentRun { t.Helper() diff --git a/pkg/coredata/agent_run.go b/pkg/coredata/agent_run.go index 0c5e1a89f..89eee4ab0 100644 --- a/pkg/coredata/agent_run.go +++ b/pkg/coredata/agent_run.go @@ -551,20 +551,43 @@ WHERE } // PGCheckpointer implements agent.Checkpointer backed by the -// agent_runs table checkpoint column. It is supervisor-internal and -// intentionally uses raw run IDs with no tenant scope; public service/API -// methods must load AgentRun through scoped coredata methods before invoking -// lifecycle transitions. +// agent_runs table checkpoint column. The runID is validated as a GID +// up front so malformed identifiers fail closed; rows are then scoped +// by primary key. type PGCheckpointer struct { - pg *pg.Client + pg *pg.Client + maxCheckpointBytes int } -func NewPGCheckpointer(pgClient *pg.Client) *PGCheckpointer { - return &PGCheckpointer{pg: pgClient} +type PGCheckpointerOption func(*PGCheckpointer) + +// WithMaxCheckpointBytes overrides the default per-checkpoint size cap +// enforced on both Save and Load. +func WithMaxCheckpointBytes(n int) PGCheckpointerOption { + return func(s *PGCheckpointer) { + if n > 0 { + s.maxCheckpointBytes = n + } + } +} + +func NewPGCheckpointer(pgClient *pg.Client, opts ...PGCheckpointerOption) *PGCheckpointer { + s := &PGCheckpointer{ + pg: pgClient, + maxCheckpointBytes: 10 * 1024 * 1024, + } + for _, opt := range opts { + opt(s) + } + return s } func (s *PGCheckpointer) Save(ctx context.Context, runID string, cp *agent.Checkpoint) error { - data, err := marshalAgentCheckpoint(cp) + if _, err := gid.ParseGID(runID); err != nil { + return fmt.Errorf("cannot parse agent run id: %w", err) + } + + data, err := s.marshalAgentCheckpoint(cp) if err != nil { return err } @@ -572,7 +595,14 @@ func (s *PGCheckpointer) Save(ctx context.Context, runID string, cp *agent.Check return s.pg.WithConn( ctx, func(ctx context.Context, conn pg.Querier) error { - q := `UPDATE agent_runs SET checkpoint = @checkpoint, updated_at = now() WHERE id = @id;` + q := ` +UPDATE agent_runs +SET + checkpoint = @checkpoint, + updated_at = now() +WHERE + id = @id; +` args := pgx.StrictNamedArgs{ "id": runID, @@ -594,12 +624,21 @@ func (s *PGCheckpointer) Save(ctx context.Context, runID string, cp *agent.Check } func (s *PGCheckpointer) Load(ctx context.Context, runID string) (*agent.Checkpoint, error) { + if _, err := gid.ParseGID(runID); err != nil { + return nil, fmt.Errorf("cannot parse agent run id: %w", err) + } + var cp *agent.Checkpoint err := s.pg.WithConn( ctx, func(ctx context.Context, conn pg.Querier) error { - q := `SELECT checkpoint FROM agent_runs WHERE id = @id;` + q := ` +SELECT checkpoint +FROM agent_runs +WHERE + id = @id; +` args := pgx.StrictNamedArgs{"id": runID} @@ -624,6 +663,10 @@ func (s *PGCheckpointer) Load(ctx context.Context, runID string) (*agent.Checkpo return nil } + if len(r.Checkpoint) > s.maxCheckpointBytes { + return fmt.Errorf("cannot load checkpoint: size %d exceeds limit %d", len(r.Checkpoint), s.maxCheckpointBytes) + } + cp = new(agent.Checkpoint) if err := json.Unmarshal(r.Checkpoint, cp); err != nil { return fmt.Errorf("cannot unmarshal checkpoint: %w", err) @@ -636,7 +679,7 @@ func (s *PGCheckpointer) Load(ctx context.Context, runID string) (*agent.Checkpo return cp, err } -func marshalAgentCheckpoint(cp *agent.Checkpoint) ([]byte, error) { +func (s *PGCheckpointer) marshalAgentCheckpoint(cp *agent.Checkpoint) ([]byte, error) { if cp == nil { return nil, fmt.Errorf("cannot marshal checkpoint: checkpoint is required") } @@ -646,8 +689,8 @@ func marshalAgentCheckpoint(cp *agent.Checkpoint) ([]byte, error) { return nil, fmt.Errorf("cannot marshal checkpoint: %w", err) } - if len(data) > agent.MaxCheckpointBytes { - return nil, fmt.Errorf("cannot marshal checkpoint: size %d exceeds limit %d", len(data), agent.MaxCheckpointBytes) + if len(data) > s.maxCheckpointBytes { + return nil, fmt.Errorf("cannot marshal checkpoint: size %d exceeds limit %d", len(data), s.maxCheckpointBytes) } return data, nil diff --git a/pkg/coredata/migrations/20260424T120000Z.sql b/pkg/coredata/migrations/20260424T173529Z.sql similarity index 84% rename from pkg/coredata/migrations/20260424T120000Z.sql rename to pkg/coredata/migrations/20260424T173529Z.sql index aa15ac008..03ae6a392 100644 --- a/pkg/coredata/migrations/20260424T120000Z.sql +++ b/pkg/coredata/migrations/20260424T173529Z.sql @@ -15,7 +15,7 @@ CREATE TABLE agent_runs ( id TEXT NOT NULL PRIMARY KEY, tenant_id TEXT NOT NULL, - organization_id TEXT NOT NULL, + organization_id TEXT NOT NULL REFERENCES organizations(id) ON DELETE CASCADE, start_agent_name TEXT NOT NULL, status TEXT NOT NULL DEFAULT 'PENDING', checkpoint JSONB, @@ -29,8 +29,5 @@ CREATE TABLE agent_runs ( updated_at TIMESTAMPTZ NOT NULL DEFAULT now() ); -CREATE INDEX idx_agent_runs_status ON agent_runs (status) - WHERE status IN ('PENDING', 'RUNNING', 'SUSPENDED'); -CREATE INDEX idx_agent_runs_organization_status ON agent_runs (organization_id, status, created_at); CREATE INDEX idx_agent_runs_running_lease ON agent_runs (lease_expires_at) WHERE status = 'RUNNING';