2026-10-12 13:50:13 +00:00
4 changed files with 375 additions and 1 deletions
@@ -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
@@ -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=
@@ -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
+}
@@ -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")
+}