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.
This commit is contained in:
@@ -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
|
||||
|
||||
@@ -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")
|
||||
}
|
||||
|
||||
@@ -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.
|
||||
|
||||
Reference in New Issue
Block a user