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