From 2fe9b5ba0d397ef3f0e4ff657e3127b4b3f7861c Mon Sep 17 00:00:00 2001 From: Samuel Enocsson Date: Mon, 12 Oct 2026 15:18:51 +0200 Subject: [PATCH] fix(master): block inventory writes from revoked agents (#6) ApplySnapshot and UpsertService now take a FOR SHARE lock on the agent's row and refuse with model.ErrRevoked when it is revoked. RevokeAgent's UPDATE waits on that lock, so an in-flight snapshot or upsert can no longer bring a revoked agent's services back after the revoke cleanup. --- internal/master/model/model.go | 1 + .../master/store/agents_integration_test.go | 28 +++++++++++ internal/master/store/services.go | 46 +++++++++++++++++-- 3 files changed, 72 insertions(+), 3 deletions(-) diff --git a/internal/master/model/model.go b/internal/master/model/model.go index ed376be..e09bc4a 100644 --- a/internal/master/model/model.go +++ b/internal/master/model/model.go @@ -15,6 +15,7 @@ import ( var ( ErrNotFound = errors.New("not found") ErrConflict = errors.New("already exists") + ErrRevoked = errors.New("agent revoked") ) // Agent is a registered agent. Whether it is connected right now is held in diff --git a/internal/master/store/agents_integration_test.go b/internal/master/store/agents_integration_test.go index 399ee9f..568ffc1 100644 --- a/internal/master/store/agents_integration_test.go +++ b/internal/master/store/agents_integration_test.go @@ -110,3 +110,31 @@ func TestAgents_RevokeRemovesOwnComposeAndContainerServices(t *testing.T) { require.NoError(t, err) assert.Equal(t, []string{"gitea_gitea"}, names(got), "swarm rows are shared and must survive a revoke") } + +func TestRevokedAgent_InventoryWritesRefused(t *testing.T) { + s := openTestStore(t) + ctx := context.Background() + id := newAgent(t, s, "node-2") + require.NoError(t, s.RevokeAgent(ctx, id, t0)) + + assert.ErrorIs(t, s.ApplySnapshot(ctx, id, worker, []protocol.Service{slusk}, t0.Add(time.Minute)), model.ErrRevoked) + assert.ErrorIs(t, s.UpsertService(ctx, id, "", adhoc, t0.Add(time.Minute)), model.ErrRevoked) + + got, err := s.ListServices(ctx) + require.NoError(t, err) + assert.Empty(t, got, "a revoked agent must not write inventory") +} + +func TestRevokedAgent_UpsertDoesNotResurrectRemovedRow(t *testing.T) { + s := openTestStore(t) + ctx := context.Background() + id := newAgent(t, s, "node-2") + require.NoError(t, s.ApplySnapshot(ctx, id, worker, []protocol.Service{slusk}, t0)) + require.NoError(t, s.RevokeAgent(ctx, id, t0.Add(time.Hour))) + + assert.ErrorIs(t, s.UpsertService(ctx, id, "", slusk, t0.Add(2*time.Hour)), model.ErrRevoked) + + got, err := s.ListServices(ctx) + require.NoError(t, err) + assert.Empty(t, got, "the revoked agent's removed row must stay removed") +} diff --git a/internal/master/store/services.go b/internal/master/store/services.go index 1a6c6e5..6f858eb 100644 --- a/internal/master/store/services.go +++ b/internal/master/store/services.go @@ -2,9 +2,11 @@ package store import ( "context" + "errors" "fmt" "time" + "github.com/jackc/pgx/v5" "github.com/jackc/pgx/v5/pgconn" "github.com/shcizo/package-updater/internal/imageref" @@ -16,6 +18,25 @@ type execer interface { Exec(ctx context.Context, sql string, args ...any) (pgconn.CommandTag, error) } +// lockActiveAgent takes a share lock on the agent's row inside tx and fails +// with model.ErrRevoked if the agent is revoked or gone. RevokeAgent's +// UPDATE of the row waits for the lock, so an inventory write either commits +// before the revoke (whose service cleanup then sees and removes its rows) +// or runs after it and is refused. Without this, an in-flight write could +// bring a revoked agent's services back. +func lockActiveAgent(ctx context.Context, tx pgx.Tx, agentID int64) error { + var one int + err := tx.QueryRow(ctx, + `SELECT 1 FROM agents WHERE id = $1 AND revoked_at IS NULL FOR SHARE`, agentID).Scan(&one) + if errors.Is(err, pgx.ErrNoRows) { + return model.ErrRevoked + } + if err != nil { + return fmt.Errorf("lock agent: %w", err) + } + return nil +} + const upsertServiceSQL = ` INSERT INTO services ( kind, location_key, agent_id, compose_project, compose_service, working_dir, @@ -52,7 +73,8 @@ func upsertService(ctx context.Context, db execer, agentID int64, clusterID stri // ApplySnapshot makes the stored inventory for this agent's scope equal to // services: everything listed is upserted, everything in scope but not -// listed is marked removed. Scope is the agent's own Compose services and +// listed is marked removed. Returns model.ErrRevoked without writing if the +// agent is revoked. Scope is the agent's own Compose services and // containers, plus the cluster's Swarm services if the agent is a manager. // A manager without a cluster ID is treated as a non-manager, since its // Swarm rows could not be told apart from other agents' rows. @@ -62,6 +84,9 @@ func (s *Store) ApplySnapshot(ctx context.Context, agentID int64, h protocol.Hel return fmt.Errorf("begin snapshot: %w", err) } defer func() { _ = tx.Rollback(ctx) }() + if err := lockActiveAgent(ctx, tx, agentID); err != nil { + return err + } isManager := h.SwarmManager && h.SwarmClusterID != "" keys := make([]string, 0, len(services)) @@ -89,9 +114,24 @@ func (s *Store) ApplySnapshot(ctx context.Context, agentID int64, h protocol.Hel return nil } -// UpsertService stores one changed or new service. +// UpsertService stores one changed or new service. Returns +// model.ErrRevoked without writing if the agent is revoked. func (s *Store) UpsertService(ctx context.Context, agentID int64, clusterID string, svc protocol.Service, at time.Time) error { - return upsertService(ctx, s.pool, agentID, clusterID, svc, at) + tx, err := s.pool.Begin(ctx) + if err != nil { + return fmt.Errorf("begin upsert: %w", err) + } + defer func() { _ = tx.Rollback(ctx) }() + if err := lockActiveAgent(ctx, tx, agentID); err != nil { + return err + } + if err := upsertService(ctx, tx, agentID, clusterID, svc, at); err != nil { + return err + } + if err := tx.Commit(ctx); err != nil { + return fmt.Errorf("commit upsert: %w", err) + } + return nil } // RemoveService marks the service with the agent-relative key as removed.