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
This commit is contained in:
2026-10-12 15:16:29 +02:00
parent 4f092eaf79
commit 0e9e34db06
12 changed files with 211 additions and 35 deletions
+9 -1
View File
@@ -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)
}
+27 -4
View File
@@ -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,
Digest: digestFor(ct.Image, repoDigests),
Image: ref,
Digest: digestFor(ref, repoDigests),
Status: containerStatus(ct.State),
Self: c.selfID != "" && strings.HasPrefix(ct.ID, c.selfID),
}
@@ -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"),
+8 -4
View File
@@ -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]
h.conns[id] = c
old := h.conns[agent.ID]
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()
}
}
+27 -1
View File
@@ -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) {
+21 -4
View File
@@ -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;
// closing a live connection is the hub's job.
// RevokeAgent marks the agent revoked and its Compose and plain-container
// 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
}
@@ -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")
}
@@ -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')),
+16 -12
View File
@@ -39,9 +39,11 @@ type Hub interface {
Disconnect(agentID int64)
}
// Subscriber delivers change signals. *events.Broker satisfies it.
type Subscriber interface {
// Broker delivers change signals to open streams and accepts signals from
// 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
hub Hub
sub Subscriber
cfg Config
log *slog.Logger
now func() time.Time
pages map[string]*template.Template
store Store
hub Hub
broker Broker
cfg Config
log *slog.Logger
now func() time.Time
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 {
+10 -1
View File
@@ -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) {
+2 -2
View File
@@ -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)
+27 -5
View File
@@ -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) {
a := newAuth(&fakeProvider{}, newFakeStore())
func TestCallback_InvalidStateRejected(t *testing.T) {
a := newAuth(&fakeProvider{claims: Claims{Subject: "u", Groups: []string{"updater-admins"}}}, newFakeStore())
flow, _ := startLogin(t, a)
rec := callback(a, flow, "forged")
assert.Equal(t, http.StatusBadRequest, rec.Code)
emptyStateFlow := &http.Cookie{Name: flowCookieName, Value: ".verifier"}
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 {