From ad6ff1e066bc1593ac6a7af3fd18f45e60c076c8 Mon Sep 17 00:00:00 2001 From: Samuel Enocsson Date: Mon, 12 Oct 2026 14:38:14 +0200 Subject: [PATCH] feat(agent): connect to master over pinned TLS and stream inventory --- go.mod | 3 +- go.sum | 2 + internal/agent/conn/conn.go | 171 ++++++++++++++++++++++++++ internal/agent/conn/conn_test.go | 200 +++++++++++++++++++++++++++++++ 4 files changed, 375 insertions(+), 1 deletion(-) create mode 100644 internal/agent/conn/conn.go create mode 100644 internal/agent/conn/conn_test.go diff --git a/go.mod b/go.mod index e916ed0..99445dc 100644 --- a/go.mod +++ b/go.mod @@ -3,8 +3,10 @@ module github.com/shcizo/package-updater go 1.26.3 require ( + github.com/coder/websocket v1.8.15 github.com/docker/docker v28.5.2+incompatible github.com/prometheus/client_golang v1.23.2 + github.com/prometheus/client_model v0.6.2 github.com/stretchr/testify v1.11.1 ) @@ -31,7 +33,6 @@ require ( github.com/opencontainers/image-spec v1.1.1 // indirect github.com/pkg/errors v0.9.1 // indirect github.com/pmezard/go-difflib v1.0.0 // indirect - github.com/prometheus/client_model v0.6.2 // indirect github.com/prometheus/common v0.66.1 // indirect github.com/prometheus/procfs v0.16.1 // indirect go.opentelemetry.io/auto/sdk v1.2.1 // indirect diff --git a/go.sum b/go.sum index db66166..e8146cc 100644 --- a/go.sum +++ b/go.sum @@ -8,6 +8,8 @@ github.com/cenkalti/backoff/v5 v5.0.3 h1:ZN+IMa753KfX5hd8vVaMixjnqRZ3y8CuJKRKj1x github.com/cenkalti/backoff/v5 v5.0.3/go.mod h1:rkhZdG3JZukswDf7f0cwqPNk4K0sa+F97BxZthm/crw= github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs= github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs= +github.com/coder/websocket v1.8.15 h1:6B2JPeOGlpff2Uz6vOEH1Vzpi0iUz20A+lPVhPHtNUA= +github.com/coder/websocket v1.8.15/go.mod h1:NX3SzP+inril6yawo5CQXx8+fk145lPDC6pumgx0mVg= github.com/containerd/errdefs v1.0.0 h1:tg5yIfIlQIrxYtu9ajqY42W3lpS19XqdxRQeEwYG8PI= github.com/containerd/errdefs v1.0.0/go.mod h1:+YBYIdtsnF4Iw6nWZhJcqGSg/dwvV7tyJ/kCkyJ2k+M= github.com/containerd/errdefs/pkg v0.3.0 h1:9IKJ06FvyNlexW690DXuQNx2KA2cUJXx151Xdx3ZPPE= diff --git a/internal/agent/conn/conn.go b/internal/agent/conn/conn.go new file mode 100644 index 0000000..ec9d92b --- /dev/null +++ b/internal/agent/conn/conn.go @@ -0,0 +1,171 @@ +// Package conn keeps an agent connected to the master: it dials over pinned +// TLS, announces itself, sends its inventory, and streams changes until the +// connection drops, then reconnects with backoff. +package conn + +import ( + "context" + "fmt" + "log/slog" + "net/http" + "time" + + "github.com/coder/websocket" + "github.com/coder/websocket/wsjson" + + "github.com/shcizo/package-updater/internal/agent/inventory" + "github.com/shcizo/package-updater/internal/pin" + "github.com/shcizo/package-updater/internal/protocol" +) + +// readLimit caps a single incoming message. The library default (32 KiB) +// is too small for the master's future job messages; matches the master. +const readLimit = 1 << 20 + +// Source supplies what the agent reports. *inventory.Collector satisfies it. +type Source interface { + Hello(ctx context.Context) (protocol.Hello, error) + Snapshot(ctx context.Context) ([]protocol.Service, error) +} + +// Config holds connection settings. Token is sent as a bearer header and must +// never be logged. +type Config struct { + URL string // wss://host:port/agent + Token string // AGENT_TOKEN + Fingerprint string // MASTER_KEY + Version string + MinBackoff time.Duration // default 1s + MaxBackoff time.Duration // default 30s +} + +// Agent is the agent's side of the master connection. +type Agent struct { + cfg Config + src Source + log *slog.Logger + changes chan struct{} +} + +// New returns an Agent. Call Run to connect. +func New(cfg Config, src Source, log *slog.Logger) *Agent { + if cfg.MinBackoff == 0 { + cfg.MinBackoff = time.Second + } + if cfg.MaxBackoff == 0 { + cfg.MaxBackoff = 30 * time.Second + } + return &Agent{cfg: cfg, src: src, log: log, changes: make(chan struct{}, 1)} +} + +// Notify tells the agent its inventory may have changed. It never blocks; +// several calls before the agent gets to them collapse into one refresh. +func (a *Agent) Notify() { + select { + case a.changes <- struct{}{}: + default: + } +} + +// Run keeps the agent connected until ctx ends, then returns nil. +func (a *Agent) Run(ctx context.Context) error { + backoff := a.cfg.MinBackoff + for { + connected, err := a.session(ctx) + if ctx.Err() != nil { + return nil + } + if connected { + backoff = a.cfg.MinBackoff + } + a.log.Warn("master_connection_lost", "error", err, "retry_in", backoff.String()) + select { + case <-ctx.Done(): + return nil + case <-time.After(backoff): + } + backoff = min(backoff*2, a.cfg.MaxBackoff) + } +} + +// session runs one connection until it fails. connected reports whether the +// 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{ + HTTPClient: httpClient, + HTTPHeader: http.Header{"Authorization": []string{"Bearer " + a.cfg.Token}}, + }) + if err != nil { + return false, fmt.Errorf("dial master: %w", err) + } + defer c.CloseNow() + c.SetReadLimit(readLimit) + + ctx, cancel := context.WithCancel(parent) + defer cancel() + // The master sends nothing yet in this phase, but reading is how a + // dropped connection is noticed, so a reader runs for the whole session. + go func() { + defer cancel() + for { + if _, _, err := c.Read(ctx); err != nil { + return + } + } + }() + + hello, err := a.src.Hello(ctx) + if err != nil { + return false, err + } + hello.AgentVersion = a.cfg.Version + if err := send(ctx, c, protocol.TypeHello, hello); err != nil { + return false, err + } + + last, err := a.src.Snapshot(ctx) + if err != nil { + return true, err + } + if err := send(ctx, c, protocol.TypeSnapshot, protocol.Snapshot{Services: last}); err != nil { + return true, err + } + a.log.Info("master_connected", "services", len(last)) + + for { + select { + case <-ctx.Done(): + _ = c.Close(websocket.StatusNormalClosure, "") + return true, ctx.Err() + case <-a.changes: + next, err := a.src.Snapshot(ctx) + if err != nil { + return true, err + } + upserts, removes := inventory.Diff(last, next) + for _, s := range upserts { + if err := send(ctx, c, protocol.TypeUpsert, protocol.Upsert{Service: s}); err != nil { + return true, err + } + } + for _, k := range removes { + if err := send(ctx, c, protocol.TypeRemove, protocol.Remove{Key: k}); err != nil { + return true, err + } + } + last = next + } + } +} + +func send(ctx context.Context, c *websocket.Conn, typ string, payload any) error { + env, err := protocol.Encode(typ, payload) + if err != nil { + return err + } + if err := wsjson.Write(ctx, c, env); err != nil { + return fmt.Errorf("send %s: %w", typ, err) + } + return nil +} diff --git a/internal/agent/conn/conn_test.go b/internal/agent/conn/conn_test.go new file mode 100644 index 0000000..ac72afb --- /dev/null +++ b/internal/agent/conn/conn_test.go @@ -0,0 +1,200 @@ +package conn + +import ( + "context" + "crypto/tls" + "io" + "log/slog" + "net" + "net/http" + "net/http/httptest" + "strings" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/coder/websocket" + "github.com/coder/websocket/wsjson" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/shcizo/package-updater/internal/pin" + "github.com/shcizo/package-updater/internal/protocol" +) + +type fakeMaster struct { + srv *httptest.Server + fingerprint string + envs chan protocol.Envelope + auth chan string + conns chan *websocket.Conn + attempts atomic.Int64 +} + +func newFakeMaster(t *testing.T) *fakeMaster { + t.Helper() + key, err := pin.GenerateKey() + require.NoError(t, err) + cert, err := pin.Certificate(key) + require.NoError(t, err) + fp, err := pin.Fingerprint(key.Public()) + require.NoError(t, err) + + fm := &fakeMaster{ + fingerprint: fp, + envs: make(chan protocol.Envelope, 100), + auth: make(chan string, 10), + conns: make(chan *websocket.Conn, 10), + } + fm.srv = httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + fm.auth <- r.Header.Get("Authorization") + c, err := websocket.Accept(w, r, nil) + if err != nil { + return + } + c.SetReadLimit(1 << 20) + fm.conns <- c + for { + var env protocol.Envelope + if err := wsjson.Read(r.Context(), c, &env); err != nil { + return + } + fm.envs <- env + } + })) + fm.srv.Config.ConnState = func(_ net.Conn, state http.ConnState) { + if state == http.StateNew { + fm.attempts.Add(1) + } + } + fm.srv.TLS = &tls.Config{Certificates: []tls.Certificate{cert}} + fm.srv.StartTLS() + t.Cleanup(fm.srv.Close) + return fm +} + +func (fm *fakeMaster) url() string { return "wss" + strings.TrimPrefix(fm.srv.URL, "https") } + +func (fm *fakeMaster) next(t *testing.T) protocol.Envelope { + t.Helper() + select { + case e := <-fm.envs: + return e + case <-time.After(2 * time.Second): + t.Fatal("timed out waiting for a message from the agent") + return protocol.Envelope{} + } +} + +type fakeSource struct { + mu sync.Mutex + services []protocol.Service +} + +func (f *fakeSource) set(s ...protocol.Service) { + f.mu.Lock() + defer f.mu.Unlock() + f.services = s +} + +func (f *fakeSource) Hello(context.Context) (protocol.Hello, error) { + return protocol.Hello{Hostname: "node-2"}, nil +} + +func (f *fakeSource) Snapshot(context.Context) ([]protocol.Service, error) { + f.mu.Lock() + defer f.mu.Unlock() + return append([]protocol.Service(nil), f.services...), nil +} + +var ( + svcA = protocol.Service{Kind: protocol.KindContainer, ContainerName: "a", Image: "a:1", Status: "running"} + svcB = protocol.Service{Kind: protocol.KindContainer, ContainerName: "b", Image: "b:1", Status: "running"} +) + +func startAgent(t *testing.T, fm *fakeMaster, fingerprint string, src Source) *Agent { + t.Helper() + a := New(Config{ + URL: fm.url(), Token: "pua_test", Fingerprint: fingerprint, Version: "v-test", + MinBackoff: 10 * time.Millisecond, MaxBackoff: 50 * time.Millisecond, + }, src, slog.New(slog.NewTextHandler(io.Discard, nil))) + ctx, cancel := context.WithCancel(context.Background()) + done := make(chan struct{}) + go func() { _ = a.Run(ctx); close(done) }() + t.Cleanup(func() { cancel(); <-done }) + return a +} + +func TestAgent_SendsHelloThenSnapshotWithToken(t *testing.T) { + fm := newFakeMaster(t) + src := &fakeSource{} + src.set(svcA) + startAgent(t, fm, fm.fingerprint, src) + + assert.Equal(t, "Bearer pua_test", <-fm.auth) + hello := fm.next(t) + require.Equal(t, protocol.TypeHello, hello.Type) + h, err := protocol.Decode[protocol.Hello](hello) + require.NoError(t, err) + assert.Equal(t, protocol.Hello{Hostname: "node-2", AgentVersion: "v-test"}, h) + + snap := fm.next(t) + require.Equal(t, protocol.TypeSnapshot, snap.Type) + s, err := protocol.Decode[protocol.Snapshot](snap) + require.NoError(t, err) + assert.Equal(t, []protocol.Service{svcA}, s.Services) +} + +func TestAgent_NotifySendsDiff(t *testing.T) { + fm := newFakeMaster(t) + src := &fakeSource{} + src.set(svcA) + a := startAgent(t, fm, fm.fingerprint, src) + fm.next(t) // hello + fm.next(t) // snapshot + + src.set(svcB) + a.Notify() + + up := fm.next(t) + require.Equal(t, protocol.TypeUpsert, up.Type) + u, err := protocol.Decode[protocol.Upsert](up) + require.NoError(t, err) + assert.Equal(t, svcB, u.Service) + + rm := fm.next(t) + require.Equal(t, protocol.TypeRemove, rm.Type) + r, err := protocol.Decode[protocol.Remove](rm) + require.NoError(t, err) + assert.Equal(t, "container|a", r.Key) +} + +func TestAgent_ReconnectsAndResendsSnapshot(t *testing.T) { + fm := newFakeMaster(t) + src := &fakeSource{} + src.set(svcA) + startAgent(t, fm, fm.fingerprint, src) + fm.next(t) + fm.next(t) + + first := <-fm.conns + first.CloseNow() + + assert.Equal(t, protocol.TypeHello, fm.next(t).Type) + assert.Equal(t, protocol.TypeSnapshot, fm.next(t).Type) +} + +func TestAgent_RefusesMasterWithWrongKey(t *testing.T) { + fm := newFakeMaster(t) + other, err := pin.GenerateKey() + require.NoError(t, err) + wrong, err := pin.Fingerprint(other.Public()) + require.NoError(t, err) + + startAgent(t, fm, wrong, &fakeSource{}) + require.Eventually(t, func() bool { return fm.attempts.Load() >= 1 }, + 2*time.Second, 10*time.Millisecond, "the agent must actually try to connect") + time.Sleep(200 * time.Millisecond) + assert.Empty(t, fm.auth, "the token must never be sent to a server with the wrong key") +}