package-updater/internal/agent/conn/conn.go

// 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

// 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)
	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)}}
	// 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)
	}
	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
}