314 lines
9.4 KiB
Go
314 lines
9.4 KiB
Go
// Package hub accepts agent connections, authenticates them, and turns their
|
|
// messages into stored inventory and live UI signals.
|
|
package hub
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"fmt"
|
|
"log/slog"
|
|
"net/http"
|
|
"strings"
|
|
"sync"
|
|
"sync/atomic"
|
|
"time"
|
|
|
|
"github.com/coder/websocket"
|
|
"github.com/coder/websocket/wsjson"
|
|
|
|
"github.com/shcizo/package-updater/internal/agenttoken"
|
|
"github.com/shcizo/package-updater/internal/master/events"
|
|
"github.com/shcizo/package-updater/internal/master/model"
|
|
"github.com/shcizo/package-updater/internal/protocol"
|
|
)
|
|
|
|
// readLimit caps one incoming message. The library default of 32 KiB is too
|
|
// small for a full snapshot from a busy machine.
|
|
const readLimit = 1 << 20
|
|
|
|
// writeTimeout bounds one websocket write. A peer that stops reading is cut
|
|
// off after this long.
|
|
const writeTimeout = 10 * time.Second
|
|
|
|
// outQueue is how many messages may wait for a connection's writer.
|
|
const outQueue = 64
|
|
|
|
// Store is the storage the hub writes to. *store.Store satisfies it.
|
|
type Store interface {
|
|
AgentByTokenHash(ctx context.Context, hash string) (model.Agent, error)
|
|
RecordHello(ctx context.Context, agentID int64, h protocol.Hello, at time.Time) error
|
|
ApplySnapshot(ctx context.Context, agentID int64, h protocol.Hello, services []protocol.Service, at time.Time) error
|
|
UpsertService(ctx context.Context, agentID int64, clusterID string, svc protocol.Service, at time.Time) error
|
|
RemoveService(ctx context.Context, agentID int64, clusterID, agentKey string, at time.Time) error
|
|
MarkSeen(ctx context.Context, agentID int64, at time.Time) error
|
|
FailLostJobs(ctx context.Context, agentID int64, known []int64, at time.Time) ([]int64, error)
|
|
JobProgress(ctx context.Context, agentID, jobID int64, chunk string, at time.Time) (bool, error)
|
|
FinishJob(ctx context.Context, agentID, jobID int64, f model.JobFinish) (bool, error)
|
|
}
|
|
|
|
// Publisher receives "something changed" signals. *events.Broker satisfies it.
|
|
type Publisher interface {
|
|
Publish(t events.Topic)
|
|
}
|
|
|
|
// Hub is the master's endpoint for agents.
|
|
type Hub struct {
|
|
store Store
|
|
pub Publisher
|
|
log *slog.Logger
|
|
now func() time.Time
|
|
|
|
connSeq atomic.Uint64
|
|
|
|
mu sync.Mutex
|
|
conns map[int64]*agentConn
|
|
}
|
|
|
|
// agentConn is one live agent connection. Its writer goroutine is the only
|
|
// code that writes to ws; everything else queues envelopes on out.
|
|
type agentConn struct {
|
|
ws *websocket.Conn
|
|
id uint64
|
|
out chan protocol.Envelope
|
|
ctx context.Context
|
|
cancel context.CancelFunc
|
|
hello *protocol.Hello // guarded by Hub.mu; nil until the hello is reconciled
|
|
}
|
|
|
|
// send queues env for the writer. It returns false if the connection ended
|
|
// before the message could be queued.
|
|
func (c *agentConn) send(env protocol.Envelope) bool {
|
|
select {
|
|
case c.out <- env:
|
|
return true
|
|
case <-c.ctx.Done():
|
|
return false
|
|
}
|
|
}
|
|
|
|
// writeLoop writes queued envelopes until the connection's context ends. A
|
|
// failed or slow write closes the connection.
|
|
func (c *agentConn) writeLoop() {
|
|
for {
|
|
select {
|
|
case <-c.ctx.Done():
|
|
return
|
|
case env := <-c.out:
|
|
wctx, cancel := context.WithTimeout(c.ctx, writeTimeout)
|
|
err := wsjson.Write(wctx, c.ws, env)
|
|
cancel()
|
|
if err != nil {
|
|
c.ws.CloseNow()
|
|
c.cancel()
|
|
return
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// New returns a Hub. Mount it at GET /agent on the TLS listener.
|
|
func New(store Store, pub Publisher, log *slog.Logger) *Hub {
|
|
return &Hub{store: store, pub: pub, log: log, now: time.Now, conns: make(map[int64]*agentConn)}
|
|
}
|
|
|
|
// ServeHTTP authenticates the agent's token and runs its connection. The token
|
|
// is only ever hashed; it is never logged or stored.
|
|
func (h *Hub) ServeHTTP(w http.ResponseWriter, r *http.Request) {
|
|
token, ok := strings.CutPrefix(r.Header.Get("Authorization"), "Bearer ")
|
|
if !ok || token == "" {
|
|
http.Error(w, "unauthorized", http.StatusUnauthorized)
|
|
return
|
|
}
|
|
agent, err := h.store.AgentByTokenHash(r.Context(), agenttoken.Hash(token))
|
|
if err != nil && !errors.Is(err, model.ErrNotFound) {
|
|
h.log.Error("agent_lookup_failed", "error", err)
|
|
}
|
|
if err != nil || agent.RevokedAt != nil {
|
|
http.Error(w, "unauthorized", http.StatusUnauthorized)
|
|
return
|
|
}
|
|
|
|
c, err := websocket.Accept(w, r, nil)
|
|
if err != nil {
|
|
return
|
|
}
|
|
c.SetReadLimit(readLimit)
|
|
ctx, cancel := context.WithCancel(r.Context())
|
|
defer cancel()
|
|
ac := &agentConn{ws: c, id: h.connSeq.Add(1), out: make(chan protocol.Envelope, outQueue), ctx: ctx, cancel: cancel}
|
|
go ac.writeLoop()
|
|
h.register(agent, r.RemoteAddr, ac)
|
|
defer h.unregister(agent.ID, ac)
|
|
|
|
// Revoking commits revoked_at before calling Disconnect. Re-checking after
|
|
// register closes the window where a revoke lands between the first lookup
|
|
// and register, when Disconnect would find nothing to cut.
|
|
again, err := h.store.AgentByTokenHash(r.Context(), agenttoken.Hash(token))
|
|
if err != nil || again.ID != agent.ID || again.RevokedAt != nil {
|
|
_ = c.Close(websocket.StatusPolicyViolation, "revoked")
|
|
return
|
|
}
|
|
|
|
s := &session{hub: h, agent: agent, conn: ac}
|
|
err = s.run(ctx, c)
|
|
h.log.Info("agent_disconnected", "agent", agent.Name, "reason", err)
|
|
}
|
|
|
|
// Online reports whether the agent has a live connection.
|
|
func (h *Hub) Online(agentID int64) bool {
|
|
h.mu.Lock()
|
|
defer h.mu.Unlock()
|
|
_, ok := h.conns[agentID]
|
|
return ok
|
|
}
|
|
|
|
// Disconnect closes the agent's live connection, if any. Call it after
|
|
// revoking the agent so the revocation takes effect immediately.
|
|
func (h *Hub) Disconnect(agentID int64) {
|
|
h.mu.Lock()
|
|
c := h.conns[agentID]
|
|
h.mu.Unlock()
|
|
if c != nil {
|
|
c.ws.CloseNow()
|
|
}
|
|
}
|
|
|
|
// 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 *agentConn) {
|
|
h.mu.Lock()
|
|
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.ws.CloseNow()
|
|
}
|
|
}
|
|
|
|
func (h *Hub) unregister(id int64, c *agentConn) {
|
|
h.mu.Lock()
|
|
current := h.conns[id] == c
|
|
if current {
|
|
delete(h.conns, id)
|
|
}
|
|
h.mu.Unlock()
|
|
if !current {
|
|
return // replaced by a newer connection; that one owns the state
|
|
}
|
|
if err := h.store.MarkSeen(context.Background(), id, h.now()); err != nil {
|
|
h.log.Error("mark_seen_failed", "agent_id", id, "error", err)
|
|
}
|
|
h.pub.Publish(events.Agents)
|
|
}
|
|
|
|
// session is one agent connection's state.
|
|
type session struct {
|
|
hub *Hub
|
|
agent model.Agent
|
|
conn *agentConn
|
|
hello *protocol.Hello
|
|
}
|
|
|
|
// swarmManager reports whether this agent may report Swarm services. The same
|
|
// rule gates snapshots, live upserts and live removes.
|
|
func (s *session) swarmManager() bool {
|
|
return s.hello.SwarmManager && s.hello.SwarmClusterID != ""
|
|
}
|
|
|
|
func (s *session) run(ctx context.Context, c *websocket.Conn) error {
|
|
for {
|
|
var env protocol.Envelope
|
|
if err := wsjson.Read(ctx, c, &env); err != nil {
|
|
return err
|
|
}
|
|
if s.hello == nil && env.Type != protocol.TypeHello {
|
|
_ = c.Close(websocket.StatusPolicyViolation, "hello required first")
|
|
return fmt.Errorf("%q before hello", env.Type)
|
|
}
|
|
if err := s.handle(ctx, env); err != nil {
|
|
_ = c.Close(websocket.StatusPolicyViolation, "bad message")
|
|
return err
|
|
}
|
|
}
|
|
}
|
|
|
|
func (s *session) handle(ctx context.Context, env protocol.Envelope) error {
|
|
h := s.hub
|
|
switch env.Type {
|
|
case protocol.TypeHello:
|
|
if s.hello != nil {
|
|
return errors.New("duplicate hello")
|
|
}
|
|
v, err := protocol.Decode[protocol.Hello](env)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if err := h.store.RecordHello(ctx, s.agent.ID, v, h.now()); err != nil {
|
|
return err
|
|
}
|
|
// Reconcile jobs before the connection becomes routable, so a job is
|
|
// never sent to an agent whose old jobs have not been settled yet.
|
|
known := parseJobIDs(v.KnownJobs, func(id string) {
|
|
h.log.Warn("known_job_not_numeric", "agent", s.agent.Name, "job_id", id)
|
|
})
|
|
lost, err := h.store.FailLostJobs(ctx, s.agent.ID, known, h.now())
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if len(lost) > 0 {
|
|
h.pub.Publish(events.Jobs)
|
|
for _, id := range lost {
|
|
h.pub.Publish(events.JobTopic(id))
|
|
}
|
|
}
|
|
s.hello = &v
|
|
h.mu.Lock()
|
|
s.conn.hello = &v
|
|
h.mu.Unlock()
|
|
h.log.Info("agent_connected", "agent", s.agent.Name, "hostname", v.Hostname, "swarm_manager", v.SwarmManager)
|
|
h.pub.Publish(events.Agents)
|
|
case protocol.TypeSnapshot:
|
|
v, err := protocol.Decode[protocol.Snapshot](env)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if err := h.store.ApplySnapshot(ctx, s.agent.ID, *s.hello, v.Services, h.now()); err != nil {
|
|
return err
|
|
}
|
|
h.pub.Publish(events.Services)
|
|
case protocol.TypeUpsert:
|
|
v, err := protocol.Decode[protocol.Upsert](env)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if v.Service.Kind == protocol.KindSwarm && !s.swarmManager() {
|
|
return nil
|
|
}
|
|
if err := h.store.UpsertService(ctx, s.agent.ID, s.hello.SwarmClusterID, v.Service, h.now()); err != nil {
|
|
return err
|
|
}
|
|
h.pub.Publish(events.Services)
|
|
case protocol.TypeRemove:
|
|
v, err := protocol.Decode[protocol.Remove](env)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if strings.HasPrefix(v.Key, string(protocol.KindSwarm)+"|") && !s.swarmManager() {
|
|
return nil
|
|
}
|
|
if err := h.store.RemoveService(ctx, s.agent.ID, s.hello.SwarmClusterID, v.Key, h.now()); err != nil {
|
|
return err
|
|
}
|
|
h.pub.Publish(events.Services)
|
|
case protocol.TypeJobProgress:
|
|
return s.handleJobProgress(ctx, env)
|
|
case protocol.TypeJobResult:
|
|
return s.handleJobResult(ctx, env)
|
|
default:
|
|
h.log.Warn("unknown_message_type", "agent", s.agent.Name, "type", env.Type)
|
|
}
|
|
return nil
|
|
}
|