fix: address final review findings for master/agent phase 1 (#6)

- inventory: resolve the image reference via ContainerInspect when a moved
  tag makes ContainerList report the image ID
- store: allow reusing a revoked agent's name (partial unique index) and
  mark the revoked agent's compose/container services removed
- web: publish a services event on revoke; no-store on the one-time token page
- webauth: reject empty OIDC state; accept the CSRF form token from the body only
- agent conn: bound the WebSocket handshake to 30s
- hub: log agent_replaced when a connection replaces an older one

internal/agent/conn/conn.go

@@ -22,6 +22,10 @@ import (
 // is too small for the master's future job messages; matches the master.
 const readLimit = 1 << 20
 
+// dialTimeout bounds the TLS + WebSocket handshake so a master that accepts
+// TCP but never answers does not hang the agent forever.
+const dialTimeout = 30 * time.Second
+
 // Source supplies what the agent reports. *inventory.Collector satisfies it.
 type Source interface {
 	Hello(ctx context.Context) (protocol.Hello, error)
@@ -92,10 +96,14 @@ func (a *Agent) Run(ctx context.Context) error {
 // hello was delivered, so Run resets its backoff only after real progress.
 func (a *Agent) session(parent context.Context) (connected bool, err error) {
 	httpClient := &http.Client{Transport: &http.Transport{TLSClientConfig: pin.ClientConfig(a.cfg.Fingerprint)}}
-	c, _, err := websocket.Dial(parent, a.cfg.URL, &websocket.DialOptions{
+	// The dial context only bounds the handshake; coder/websocket does not
+	// tie the established connection to it, so cancelling afterwards is safe.
+	dialCtx, cancelDial := context.WithTimeout(parent, dialTimeout)
+	c, _, err := websocket.Dial(dialCtx, a.cfg.URL, &websocket.DialOptions{
 		HTTPClient: httpClient,
 		HTTPHeader: http.Header{"Authorization": []string{"Bearer " + a.cfg.Token}},
 	})
+	cancelDial()
 	if err != nil {
 		return false, fmt.Errorf("dial master: %w", err)
 	}

internal/agent/inventory/collector.go

@@ -28,6 +28,7 @@ const (
 // *client.Client satisfies it.
 type DockerClient interface {
 	ContainerList(ctx context.Context, opts container.ListOptions) ([]container.Summary, error)
+	ContainerInspect(ctx context.Context, containerID string) (container.InspectResponse, error)
 	ImageList(ctx context.Context, opts image.ListOptions) ([]image.Summary, error)
 	ServiceList(ctx context.Context, opts swarm.ServiceListOptions) ([]swarm.Service, error)
 	Info(ctx context.Context) (system.Info, error)
@@ -83,7 +84,11 @@ func (c *Collector) Snapshot(ctx context.Context) ([]protocol.Service, error) {
 		if _, isTask := ct.Labels[labelSwarmTask]; isTask {
 			continue
 		}
-		svc := c.fromContainer(ct, repoDigests[ct.ImageID])
+		ref, err := c.imageRef(ctx, ct)
+		if err != nil {
+			return nil, err
+		}
+		svc := c.fromContainer(ct, ref, repoDigests[ct.ImageID])
 		if prev, ok := byKey[svc.Key()]; ok {
 			svc = mergeReplicas(prev, svc)
 		}
@@ -112,10 +117,28 @@ func (c *Collector) Snapshot(ctx context.Context) ([]protocol.Service, error) {
 	return out, nil
 }
 
-func (c *Collector) fromContainer(ct container.Summary, repoDigests []string) protocol.Service {
+// imageRef returns the image reference the container was created from.
+// When the tag has since moved to another image, ContainerList reports the
+// image ID ("sha256:...") instead, so the reference is read from the
+// container's config.
+func (c *Collector) imageRef(ctx context.Context, ct container.Summary) (string, error) {
+	if !strings.HasPrefix(ct.Image, "sha256:") {
+		return ct.Image, nil
+	}
+	insp, err := c.cli.ContainerInspect(ctx, ct.ID)
+	if err != nil {
+		return "", fmt.Errorf("docker container inspect %s: %w", ct.ID, err)
+	}
+	if insp.Config == nil {
+		return ct.Image, nil
+	}
+	return insp.Config.Image, nil
+}
+
+func (c *Collector) fromContainer(ct container.Summary, ref string, repoDigests []string) protocol.Service {
 	svc := protocol.Service{
-		Image:  ct.Image,
+		Image:  ref,
-		Digest: digestFor(ct.Image, repoDigests),
+		Digest: digestFor(ref, repoDigests),
 		Status: containerStatus(ct.State),
 		Self:   c.selfID != "" && strings.HasPrefix(ct.ID, c.selfID),
 	}

internal/agent/inventory/collector_test.go

@@ -22,12 +22,18 @@ type fakeDocker struct {
 	info             system.Info
 	containerErr     error
 	serviceListCalls int
+	inspect          map[string]container.InspectResponse
+	inspectErr       error
 }
 
 func (f *fakeDocker) ContainerList(context.Context, container.ListOptions) ([]container.Summary, error) {
 	return f.containers, f.containerErr
 }
 
+func (f *fakeDocker) ContainerInspect(_ context.Context, id string) (container.InspectResponse, error) {
+	return f.inspect[id], f.inspectErr
+}
+
 func (f *fakeDocker) ImageList(context.Context, image.ListOptions) ([]image.Summary, error) {
 	return f.images, nil
 }
@@ -72,6 +78,32 @@ func TestSnapshot_ComposeContainerWithRepoDigest(t *testing.T) {
 	}}, got)
 }
 
+func TestSnapshot_MovedTagUsesConfiguredReference(t *testing.T) {
+	ct := composeContainer("abc123", "slusk", "web", "sha256:img1", "sha256:img1", "running")
+	cli := &fakeDocker{
+		containers: []container.Summary{ct},
+		images:     []image.Summary{{ID: "sha256:img1", RepoDigests: []string{"gitea.shcizo.se/shcizo/slusk@sha256:d1"}}},
+		inspect: map[string]container.InspectResponse{
+			"abc123": {Config: &container.Config{Image: "gitea.shcizo.se/shcizo/slusk:1.4.2"}},
+		},
+	}
+	got, err := New(cli, "").Snapshot(context.Background())
+	require.NoError(t, err)
+	require.Len(t, got, 1)
+	assert.Equal(t, "gitea.shcizo.se/shcizo/slusk:1.4.2", got[0].Image)
+	assert.Equal(t, "sha256:d1", got[0].Digest)
+}
+
+func TestSnapshot_InspectErrorPropagates(t *testing.T) {
+	cli := &fakeDocker{
+		containers: []container.Summary{composeContainer("abc123", "slusk", "web", "sha256:img1", "sha256:img1", "running")},
+		inspectErr: errors.New("gone"),
+	}
+	_, err := New(cli, "").Snapshot(context.Background())
+	require.Error(t, err)
+	assert.Contains(t, err.Error(), "docker container inspect")
+}
+
 func TestSnapshot_ReplicasMergeAndRunningWins(t *testing.T) {
 	cli := &fakeDocker{containers: []container.Summary{
 		composeContainer("a1", "slusk", "worker", "slusk:1", "sha256:i", "exited"),

internal/master/hub/hub.go

@@ -78,7 +78,7 @@ func (h *Hub) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 	c.SetReadLimit(readLimit)
-	h.register(agent.ID, c)
+	h.register(agent, r.RemoteAddr, c)
 	defer h.unregister(agent.ID, c)
 
 	// Revoking commits revoked_at before calling Disconnect. Re-checking after
@@ -114,12 +114,16 @@ func (h *Hub) Disconnect(agentID int64) {
 	}
 }
 
-func (h *Hub) register(id int64, c *websocket.Conn) {
+// register makes c the agent's live connection, closing any older one. A
+// replacement is logged: it is normal after a network blip, but repeated
+// replacements mean two agents share a token.
+func (h *Hub) register(agent model.Agent, remoteAddr string, c *websocket.Conn) {
 	h.mu.Lock()
-	old := h.conns[id]
+	old := h.conns[agent.ID]
-	h.conns[id] = c
+	h.conns[agent.ID] = c
 	h.mu.Unlock()
 	if old != nil {
+		h.log.Warn("agent_replaced", "agent_id", agent.ID, "agent", agent.Name, "remote_addr", remoteAddr)
 		old.CloseNow()
 	}
 }

internal/master/hub/hub_test.go

@@ -1,6 +1,7 @@
 package hub
 
 import (
+	"bytes"
 	"context"
 	"io"
 	"log/slog"
@@ -320,9 +321,31 @@ func TestHub_ClosesConnectionThatSkipsHello(t *testing.T) {
 	assert.Equal(t, websocket.StatusPolicyViolation, websocket.CloseStatus(err))
 }
 
+// lockedBuffer is a bytes.Buffer safe for the hub's handler goroutines to
+// write while the test reads it.
+type lockedBuffer struct {
+	mu  sync.Mutex
+	buf bytes.Buffer
+}
+
+func (b *lockedBuffer) Write(p []byte) (int, error) {
+	b.mu.Lock()
+	defer b.mu.Unlock()
+	return b.buf.Write(p)
+}
+
+func (b *lockedBuffer) String() string {
+	b.mu.Lock()
+	defer b.mu.Unlock()
+	return b.buf.String()
+}
+
 func TestHub_NewConnectionReplacesOld(t *testing.T) {
 	store := newFakeStore(node2)
-	h, _, srv := newHub(t, store)
+	logs := &lockedBuffer{}
+	h := New(store, &fakePub{}, slog.New(slog.NewTextHandler(logs, nil)))
+	srv := httptest.NewServer(h)
+	t.Cleanup(srv.Close)
 	first, _, err := dial(t, srv, "pua_node-2")
 	require.NoError(t, err)
 	defer first.CloseNow()
@@ -336,6 +359,9 @@ func TestHub_NewConnectionReplacesOld(t *testing.T) {
 	_, _, err = first.Read(context.Background())
 	assert.Error(t, err, "the replaced connection must be closed")
 	assert.True(t, h.Online(node2.ID), "the agent stays online on its new connection")
+	assert.Contains(t, logs.String(), "msg=agent_replaced")
+	assert.Contains(t, logs.String(), "agent=node-2")
+	assert.NotContains(t, logs.String(), "pua_", "tokens are never logged")
 }
 
 func TestHub_DisconnectCutsLiveConnection(t *testing.T) {

internal/master/store/agents.go

@@ -24,7 +24,8 @@ func scanAgent(row pgx.Row) (model.Agent, error) {
 	return a, err
 }
 
-// CreateAgent registers an agent. Returns model.ErrConflict if the name is taken.
+// CreateAgent registers an agent. Returns model.ErrConflict if an active
+// (non-revoked) agent already has the name.
 func (s *Store) CreateAgent(ctx context.Context, name, tokenHash string) (model.Agent, error) {
 	a, err := scanAgent(s.pool.QueryRow(ctx,
 		`INSERT INTO agents (name, token_hash) VALUES ($1, $2) RETURNING `+agentColumns, name, tokenHash))
@@ -62,13 +63,29 @@ func (s *Store) AgentByTokenHash(ctx context.Context, hash string) (model.Agent,
 	return scanAgent(s.pool.QueryRow(ctx, `SELECT `+agentColumns+` FROM agents WHERE token_hash = $1`, hash))
 }
 
-// RevokeAgent marks the agent revoked. Its token stops working immediately;
+// RevokeAgent marks the agent revoked and its Compose and plain-container
-// closing a live connection is the hub's job.
+// services removed, in one transaction. Swarm rows are shared with other
+// managers and left alone. Its token stops working immediately; closing a
+// live connection is the hub's job.
 func (s *Store) RevokeAgent(ctx context.Context, id int64, at time.Time) error {
-	if _, err := s.pool.Exec(ctx,
+	tx, err := s.pool.Begin(ctx)
+	if err != nil {
+		return fmt.Errorf("revoke agent: begin: %w", err)
+	}
+	defer func() { _ = tx.Rollback(ctx) }()
+
+	if _, err := tx.Exec(ctx,
 		`UPDATE agents SET revoked_at = $2 WHERE id = $1 AND revoked_at IS NULL`, id, at); err != nil {
 		return fmt.Errorf("revoke agent: %w", err)
 	}
+	if _, err := tx.Exec(ctx,
+		`UPDATE services SET removed_at = $2
+		 WHERE agent_id = $1 AND kind IN ('compose', 'container') AND removed_at IS NULL`, id, at); err != nil {
+		return fmt.Errorf("revoke agent: remove services: %w", err)
+	}
+	if err := tx.Commit(ctx); err != nil {
+		return fmt.Errorf("revoke agent: commit: %w", err)
+	}
 	return nil
 }
 

internal/master/store/agents_integration_test.go

@@ -82,3 +82,31 @@ func TestAgents_RecordHelloAndMarkSeen(t *testing.T) {
 	require.NotNil(t, agents[0].LastSeenAt)
 	assert.True(t, agents[0].LastSeenAt.Equal(at.Add(time.Minute)))
 }
+
+func TestAgents_RevokedNameCanBeReused(t *testing.T) {
+	s := openTestStore(t)
+	ctx := context.Background()
+	old, err := s.CreateAgent(ctx, "node-2", "h1")
+	require.NoError(t, err)
+	require.NoError(t, s.RevokeAgent(ctx, old.ID, time.Now()))
+
+	fresh, err := s.CreateAgent(ctx, "node-2", "h2")
+	require.NoError(t, err)
+	assert.NotEqual(t, old.ID, fresh.ID)
+
+	_, err = s.CreateAgent(ctx, "node-2", "h3")
+	assert.ErrorIs(t, err, model.ErrConflict, "a second active agent with the name is still a conflict")
+}
+
+func TestAgents_RevokeRemovesOwnComposeAndContainerServices(t *testing.T) {
+	s := openTestStore(t)
+	ctx := context.Background()
+	id := newAgent(t, s, "swarm-mgr")
+	require.NoError(t, s.ApplySnapshot(ctx, id, manager, []protocol.Service{slusk, adhoc, gitea}, t0))
+
+	require.NoError(t, s.RevokeAgent(ctx, id, t0.Add(time.Hour)))
+
+	got, err := s.ListServices(ctx)
+	require.NoError(t, err)
+	assert.Equal(t, []string{"gitea_gitea"}, names(got), "swarm rows are shared and must survive a revoke")
+}

internal/master/store/migrations/0001_init.sql

@@ -6,7 +6,7 @@ CREATE TABLE master_identity (
 
 CREATE TABLE agents (
     id            bigserial PRIMARY KEY,
-    name          text NOT NULL UNIQUE,
+    name          text NOT NULL,
     token_hash    text NOT NULL UNIQUE,
     hostname      text NOT NULL DEFAULT '',
     agent_version text NOT NULL DEFAULT '',
@@ -16,6 +16,9 @@ CREATE TABLE agents (
     created_at    timestamptz NOT NULL DEFAULT now()
 );
 
+-- A revoked agent's name can be reused by a new agent.
+CREATE UNIQUE INDEX agents_name_active ON agents (name) WHERE revoked_at IS NULL;
+
 CREATE TABLE services (
     id               bigserial PRIMARY KEY,
     kind             text NOT NULL CHECK (kind IN ('compose', 'swarm', 'container')),

internal/master/web/server.go

@@ -39,9 +39,11 @@ type Hub interface {
 	Disconnect(agentID int64)
 }
 
-// Subscriber delivers change signals. *events.Broker satisfies it.
+// Broker delivers change signals to open streams and accepts signals from
-type Subscriber interface {
+// UI actions. *events.Broker satisfies it.
+type Broker interface {
 	Subscribe() (<-chan events.Topic, func())
+	Publish(t events.Topic)
 }
 
 // Config holds values shown in the "add agent" snippet.
@@ -53,18 +55,18 @@ type Config struct {
 
 // Server renders the UI.
 type Server struct {
-	store Store
+	store  Store
-	hub   Hub
+	hub    Hub
-	sub   Subscriber
+	broker Broker
-	cfg   Config
+	cfg    Config
-	log   *slog.Logger
+	log    *slog.Logger
-	now   func() time.Time
+	now    func() time.Time
-	pages map[string]*template.Template
+	pages  map[string]*template.Template
 }
 
 // New parses the templates. Each page is its own template set so every page
 // can define its own "content" block.
-func New(store Store, hub Hub, sub Subscriber, cfg Config, log *slog.Logger) (*Server, error) {
+func New(store Store, hub Hub, broker Broker, cfg Config, log *slog.Logger) (*Server, error) {
 	funcs := template.FuncMap{"shortRef": shortRef}
 	pages := map[string]*template.Template{}
 	for _, page := range []string{"overview.html", "agents.html", "agent_created.html"} {
@@ -74,7 +76,7 @@ func New(store Store, hub Hub, sub Subscriber, cfg Config, log *slog.Logger) (*S
 		}
 		pages[page] = t
 	}
-	return &Server{store: store, hub: hub, sub: sub, cfg: cfg, log: log, now: time.Now, pages: pages}, nil
+	return &Server{store: store, hub: hub, broker: broker, cfg: cfg, log: log, now: time.Now, pages: pages}, nil
 }
 
 // Register adds the UI routes to mux. protect wraps every route except
@@ -195,6 +197,7 @@ func (s *Server) createAgent(w http.ResponseWriter, r *http.Request) {
 	data := s.base(r, "Agent created")
 	data.Created = agent
 	data.Snippet = agentSnippet(s.cfg, token)
+	w.Header().Set("Cache-Control", "no-store") // the page shows the one-time token
 	s.render(w, http.StatusOK, "agent_created.html", "layout", data)
 }
 
@@ -219,6 +222,7 @@ func (s *Server) revokeAgent(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 	s.hub.Disconnect(id)
+	s.broker.Publish(events.Services) // the agent's services were marked removed
 	http.Redirect(w, r, "/agents", http.StatusSeeOther)
 }
 
@@ -231,7 +235,7 @@ func (s *Server) events(w http.ResponseWriter, r *http.Request) {
 	w.Header().Set("Cache-Control", "no-cache")
 	w.Header().Set("X-Accel-Buffering", "no") // stop nginx from buffering the stream
 
-	ch, cancel := s.sub.Subscribe()
+	ch, cancel := s.broker.Subscribe()
 	defer cancel()
 
 	if _, err := fmt.Fprint(w, ": connected\n\n"); err != nil {

internal/master/web/server_test.go

@@ -175,6 +175,7 @@ func TestCreateAgent_ShowsSnippetOnceAndStoresHash(t *testing.T) {
 	body, _ := io.ReadAll(resp.Body)
 
 	require.Equal(t, http.StatusOK, resp.StatusCode)
+	assert.Equal(t, "no-store", resp.Header.Get("Cache-Control"), "the one-time token must not be cached")
 	token := regexp.MustCompile(`pua_[A-Za-z0-9_-]{43}`).FindString(string(body))
 	require.NotEmpty(t, token)
 	assert.Equal(t, agenttoken.Hash(token), store.createdHash("node-4"), "only the hash is stored")
@@ -207,7 +208,9 @@ func TestCreateAgent_Validation(t *testing.T) {
 }
 
 func TestRevokeAgent_RevokesAndDisconnects(t *testing.T) {
-	srv, store, hub, _ := newTestServer(t)
+	srv, store, hub, broker := newTestServer(t)
+	ch, cancel := broker.Subscribe()
+	defer cancel()
 	resp, err := noRedirectClient().Post(srv.URL+"/agents/2/revoke", "application/x-www-form-urlencoded", nil)
 	require.NoError(t, err)
 	_ = resp.Body.Close()
@@ -216,6 +219,12 @@ func TestRevokeAgent_RevokesAndDisconnects(t *testing.T) {
 	assert.Equal(t, "/agents", resp.Header.Get("Location"))
 	assert.Equal(t, int64(2), store.revoked())
 	assert.Equal(t, int64(2), hub.lastDisconnected(), "a revoked agent must be cut off immediately")
+	select {
+	case topic := <-ch:
+		assert.Equal(t, events.Services, topic, "the overview must drop the revoked agent's services")
+	default:
+		t.Fatal("revoke published no services event")
+	}
 }
 
 func TestRevokeAgent_BadID(t *testing.T) {

internal/master/webauth/auth.go

@@ -71,7 +71,7 @@ func (a *Auth) Callback(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 	state, verifier, ok := strings.Cut(flow.Value, ".")
-	if !ok || subtle.ConstantTimeCompare([]byte(r.URL.Query().Get("state")), []byte(state)) != 1 {
+	if !ok || state == "" || subtle.ConstantTimeCompare([]byte(r.URL.Query().Get("state")), []byte(state)) != 1 {
 		http.Error(w, "invalid login state", http.StatusBadRequest)
 		return
 	}
@@ -150,7 +150,7 @@ func (a *Auth) Require(next http.Handler) http.Handler {
 		if r.Method != http.MethodGet && r.Method != http.MethodHead {
 			token := r.Header.Get("X-CSRF-Token")
 			if token == "" {
-				token = r.FormValue("csrf_token")
+				token = r.PostFormValue("csrf_token") // body only: a token in the URL leaks via logs and Referer
 			}
 			if subtle.ConstantTimeCompare([]byte(token), []byte(sess.CSRFToken)) != 1 {
 				http.Error(w, "invalid CSRF token", http.StatusForbidden)

internal/master/webauth/auth_test.go

@@ -84,7 +84,9 @@ func startLogin(t *testing.T, a *Auth) (*http.Cookie, string) {
 
 func callback(a *Auth, flow *http.Cookie, state string) *httptest.ResponseRecorder {
 	req := httptest.NewRequest(http.MethodGet, "/auth/callback?code=c&state="+url.QueryEscape(state), nil)
-	req.AddCookie(flow)
+	if flow != nil {
+		req.AddCookie(flow)
+	}
 	rec := httptest.NewRecorder()
 	a.Callback(rec, req)
 	return rec
@@ -98,11 +100,25 @@ func TestLogin_SetsFlowCookieAndRedirects(t *testing.T) {
 	assert.NotEmpty(t, state)
 }
 
-func TestCallback_StateMismatchRejected(t *testing.T) {
+func TestCallback_InvalidStateRejected(t *testing.T) {
-	a := newAuth(&fakeProvider{}, newFakeStore())
+	a := newAuth(&fakeProvider{claims: Claims{Subject: "u", Groups: []string{"updater-admins"}}}, newFakeStore())
 	flow, _ := startLogin(t, a)
-	rec := callback(a, flow, "forged")
+	emptyStateFlow := &http.Cookie{Name: flowCookieName, Value: ".verifier"}
-	assert.Equal(t, http.StatusBadRequest, rec.Code)
+	cases := []struct {
+		name  string
+		flow  *http.Cookie
+		state string
+	}{
+		{"state mismatch", flow, "forged"},
+		{"empty state in cookie and query", emptyStateFlow, ""},
+		{"missing flow cookie", nil, ""},
+	}
+	for _, c := range cases {
+		t.Run(c.name, func(t *testing.T) {
+			rec := callback(a, c.flow, c.state)
+			assert.Equal(t, http.StatusBadRequest, rec.Code)
+		})
+	}
 }
 
 func TestCallback_UserOutsideGroupForbidden(t *testing.T) {
@@ -197,6 +213,12 @@ func TestRequire(t *testing.T) {
 			r.AddCookie(&http.Cookie{Name: sessionCookieName, Value: "live"})
 			return r
 		}(), http.StatusNoContent, [2]string{}},
+		{"POST with CSRF only in query string", func() *http.Request {
+			r := httptest.NewRequest(http.MethodPost, "/agents?csrf_token=tok", strings.NewReader("name=x"))
+			r.Header.Set("Content-Type", "application/x-www-form-urlencoded")
+			r.AddCookie(&http.Cookie{Name: sessionCookieName, Value: "live"})
+			return r
+		}(), http.StatusForbidden, [2]string{}},
 		{"GET with session", protectedRequest(http.MethodGet, "/", "live"), http.StatusNoContent, [2]string{}},
 	}
 	for _, c := range cases {