feat(master): agent hub that authenticates agents and stores their inventory
This commit is contained in:
@@ -0,0 +1,215 @@
|
||||
// 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"
|
||||
"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
|
||||
|
||||
// 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
|
||||
}
|
||||
|
||||
// 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
|
||||
|
||||
mu sync.Mutex
|
||||
conns map[int64]*websocket.Conn
|
||||
}
|
||||
|
||||
// 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]*websocket.Conn)}
|
||||
}
|
||||
|
||||
// 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)
|
||||
h.register(agent.ID, c)
|
||||
defer h.unregister(agent.ID, c)
|
||||
|
||||
s := &session{hub: h, agent: agent}
|
||||
err = s.run(r.Context(), 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.CloseNow()
|
||||
}
|
||||
}
|
||||
|
||||
func (h *Hub) register(id int64, c *websocket.Conn) {
|
||||
h.mu.Lock()
|
||||
old := h.conns[id]
|
||||
h.conns[id] = c
|
||||
h.mu.Unlock()
|
||||
if old != nil {
|
||||
old.CloseNow()
|
||||
}
|
||||
}
|
||||
|
||||
func (h *Hub) unregister(id int64, c *websocket.Conn) {
|
||||
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
|
||||
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:
|
||||
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
|
||||
}
|
||||
s.hello = &v
|
||||
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)
|
||||
default:
|
||||
return fmt.Errorf("unknown message type %q", env.Type)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
Reference in New Issue
Block a user