package-updater/internal/master/hub/hub.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"
	"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, r.RemoteAddr, c)
	defer h.unregister(agent.ID, c)

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

// 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 *websocket.Conn) {
	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.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:
		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
		}
		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
}