feat(agent): connect to master over pinned TLS and stream inventory

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
+}