Commilitia-Drop/internal/hub/hub.go

package hub

import (
	"context"
	"log/slog"
	"sync"
	"time"

	"commilitia.net/cdrop/internal/db"
)

// DeviceLister is the subset of *db.Queries the hub needs.
// Declared as interface so tests can inject a fake.
type DeviceLister interface {
	ListDevicesByUser(ctx context.Context, userID string) ([]db.Device, error)
}

// Event is what the hub fan-outs to SSE clients.
type Event struct {
	Type string `json:"-"`
	Data any    `json:"-"`
}

// PresenceDevice is the per-device entry in a `presence` event payload.
// Field names align with brief §5: { name, type, online }; last_seen 是
// 设置页展示「上次活跃」需要的扩展字段。device_id 是稳定标识：name 与令牌解耦后，
// 前端按它去重 / 判本机 / 寻址，改名只换 name 字段而 device_id 不变（避免改名生重复项）。
type PresenceDevice struct {
	DeviceID string `json:"device_id"`
	Name     string `json:"name"`
	Type     string `json:"type"`
	Online   bool   `json:"online"`
	LastSeen int64  `json:"last_seen"`
}

const clientBuffer = 32

// Client is a single live SSE connection.
type Client struct {
	UserID   string
	DeviceID string
	// Name 是连接声明的设备名（X-Device-Name）。presence 的权威 name 来自 devices 表（按
	// device_id 查，改名即时反映）；本字段仅用于 (1) 无 devices 行的 live 连接（代铸前的
	// 全局 SSO 浏览器）的 presence 兜底标签，(2) 按 name 寻址时的解析（见 clientFor）。
	Name string
	// Type is the client-declared device type (browser / macos / windows / linux / ios).
	// It lets presence label a live device that has no devices-table row (a global-SSO
	// browser or a broker-authenticated desktop).
	Type string
	// key 是 users[userID] 这一级 map 的键：有 device_id 用 device_id（稳定，改名不变），
	// 无（代铸前浏览器）回落 "name:"+name。Disconnect 据此精确摘除自身。
	key string
	ch  chan Event
}

// connKey 计算一个 live 连接在 hub map 里的键：device_id 优先（稳定标识，改名不变，故改名
// 不再产生孤儿键 / 重复项 —— 这是「设备名与令牌解耦」的服务端命门）；无 device_id（代铸前
// 的全局 SSO 浏览器）回落 "name:"+name，与 device_id 命名空间不冲突。
func connKey(deviceID, name string) string {
	if deviceID != "" {
		return deviceID
	}
	return "name:" + name
}

func (c *Client) Events() <-chan Event { return c.ch }

// Hub is the in-memory presence + signaling fan-out.
//
// Concurrency: a single sync.RWMutex protects the user→device map.
// Sends to client channels are non-blocking — slow consumers drop events
// rather than stall the broadcaster.
type Hub struct {
	mu    sync.RWMutex
	users map[string]map[string]*Client

	devices DeviceLister
	grace   time.Duration

	closed bool
}

func New(devices DeviceLister) *Hub {
	return &Hub{
		users:   map[string]map[string]*Client{},
		devices: devices,
		grace:   30 * time.Second,
	}
}

// Connect registers a new SSE client and announces presence to the user's other devices.
// Keyed by the connection key (device_id when present): a reconnect of the SAME device — even
// after a rename, since device_id is stable — replaces its own entry (its channel is closed)
// rather than orphaning the old one. deviceName is the connection-declared name (presence
// fallback + name addressing); deviceType is surfaced for a live device with no devices-table row.
func (h *Hub) Connect(ctx context.Context, userID, deviceID, deviceName, deviceType string) *Client {
	key := connKey(deviceID, deviceName)
	c := &Client{
		UserID:   userID,
		DeviceID: deviceID,
		Name:     deviceName,
		Type:     deviceType,
		key:      key,
		ch:       make(chan Event, clientBuffer),
	}

	h.mu.Lock()
	if h.users[userID] == nil {
		h.users[userID] = map[string]*Client{}
	}
	if old, ok := h.users[userID][key]; ok {
		close(old.ch)
	}
	h.users[userID][key] = c
	h.mu.Unlock()

	go h.publishPresence(ctx, userID)
	return c
}

// Disconnect removes a client and, after a grace period, re-publishes presence
// to peers if the client has not reconnected.
func (h *Hub) Disconnect(c *Client) {
	h.mu.Lock()
	if active, ok := h.users[c.UserID][c.key]; ok && active == c {
		delete(h.users[c.UserID], c.key)
		if len(h.users[c.UserID]) == 0 {
			delete(h.users, c.UserID)
		}
	}
	h.mu.Unlock()

	time.AfterFunc(h.grace, func() {
		h.mu.RLock()
		_, stillOnline := h.users[c.UserID][c.key]
		h.mu.RUnlock()
		if stillOnline {
			return
		}
		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
		defer cancel()
		h.publishPresence(ctx, c.UserID)
	})
}

// SendTo routes an event to one of the user's live clients addressed by `target`. The
// target may be a stable device_id (direct map hit) OR a device name — peer addressing
// across cdrop (signaling / messages / transfers) is by name, so this resolves a name to
// its live client via clientFor while presence/online/kick key by the stable device_id.
// Reports whether the target was online and the event was queued.
//
// The send happens while still holding the read lock so it can never race a Kick / Connect-
// replace / Close that closes the channel (those hold the write lock): a send on a closed
// channel panics even inside a select, so close-vs-send must be mutually exclusive. The send
// is non-blocking (select default), so holding the read lock across it is brief.
func (h *Hub) SendTo(userID, target string, ev Event) bool {
	h.mu.RLock()
	defer h.mu.RUnlock()
	c := h.clientFor(userID, target)
	if c == nil {
		return false
	}
	select {
	case c.ch <- ev:
		return true
	default:
		slog.Warn("client buffer full; dropping event",
			"user", userID, "device", target, "type", ev.Type)
		return false
	}
}

// clientFor resolves `target` (device_id or device name) to a live client. Caller holds at
// least the read lock. A direct map hit handles device_id (and the "name:"+name fallback key
// of a code-less browser is never addressed directly); otherwise a scan matches the current
// display name — so a renamed device is reachable by its NEW name as soon as Rename lands.
func (h *Hub) clientFor(userID, target string) *Client {
	if c, ok := h.users[userID][target]; ok {
		return c
	}
	for _, c := range h.users[userID] {
		if c.Name == target {
			return c
		}
	}
	return nil
}

// Broadcast fans an event out to every live client of a user. The non-blocking sends run
// under the read lock so they can't race a concurrent channel close (see SendTo).
func (h *Hub) Broadcast(userID string, ev Event) {
	h.mu.RLock()
	defer h.mu.RUnlock()
	for _, c := range h.users[userID] {
		select {
		case c.ch <- ev:
		default:
			slog.Warn("client buffer full; dropping event",
				"user", c.UserID, "device", c.DeviceID, "type", ev.Type)
		}
	}
}

// Online reports whether a device (by stable device_id) currently has a live SSE.
func (h *Hub) Online(userID, deviceID string) bool {
	h.mu.RLock()
	defer h.mu.RUnlock()
	_, ok := h.users[userID][deviceID]
	return ok
}

// Rename updates a live client's display name in place (keyed by stable device_id), so a
// device renamed via PATCH /api/devices/{device_id} stays the SAME hub entry — no reconnect,
// no token re-mint, no duplicate — and becomes reachable by its new name (clientFor) at once.
// No-op if the device has no live connection (its next connect carries the new name anyway).
func (h *Hub) Rename(userID, deviceID, name string) {
	if deviceID == "" {
		return
	}
	h.mu.Lock()
	defer h.mu.Unlock()
	if c, ok := h.users[userID][deviceID]; ok {
		c.Name = name
	}
}

// Kick force-removes a device's live entry from the hub and closes its event channel; the SSE
// handler exits on the next iteration. Keyed by the connection key (device_id when present,
// else "name:"+name) so it摘除 the exact entry Connect inserted. The close happens UNDER the
// write lock — the same discipline as Connect-replace and Close — so it can never race a send
// from publishPresence / Broadcast / SendTo (those hold the read lock), which would otherwise
// panic on a send to a closed channel. revokeDevice now Kicks on every logout / device delete /
// session revoke, so this path is hot, not rare.
func (h *Hub) Kick(userID, deviceID, name string) {
	key := connKey(deviceID, name)
	h.mu.Lock()
	defer h.mu.Unlock()
	if c, ok := h.users[userID][key]; ok {
		delete(h.users[userID], key)
		if len(h.users[userID]) == 0 {
			delete(h.users, userID)
		}
		close(c.ch)
	}
}

// PublishPresence broadcasts the user's current device list to all live
// clients. 用于在变更设备集合后（如 DELETE /api/devices/{name}）立刻广播，
// 而不是等待 Disconnect 的宽限期。
func (h *Hub) PublishPresence(ctx context.Context, userID string) {
	h.publishPresence(ctx, userID)
}

// Close drops every connected client. Used on graceful shutdown.
func (h *Hub) Close() {
	h.mu.Lock()
	defer h.mu.Unlock()
	if h.closed {
		return
	}
	h.closed = true
	for _, devices := range h.users {
		for _, c := range devices {
			close(c.ch)
		}
	}
	h.users = map[string]map[string]*Client{}
}

func (h *Hub) publishPresence(ctx context.Context, userID string) {
	devs, err := h.devices.ListDevicesByUser(ctx, userID)
	if err != nil {
		// Log but continue with an empty managed-device set: a transient DB error must
		// not blank out the presence of live, unmanaged devices that need no row.
		slog.Error("presence: list devices failed", "user", userID, "err", err)
	}

	now := time.Now().Unix()
	// Hold the read lock across the build AND the sends: the non-blocking sends below must
	// be mutually exclusive with any channel close (Kick / Connect-replace / Close hold the
	// write lock), or a send could hit a closed channel and panic.
	h.mu.RLock()
	defer h.mu.RUnlock()
	live := h.users[userID]
	seen := make(map[string]bool, len(devs))
	items := make([]PresenceDevice, 0, len(devs)+len(live))
	for _, d := range devs {
		// Online keyed by the stable device_id (not the name) — so a renamed device stays
		// online through the rename, and the authoritative name is the (just-updated) DB row.
		_, online := live[d.DeviceID]
		items = append(items, PresenceDevice{
			DeviceID: d.DeviceID, Name: d.Name, Type: d.Type, Online: online, LastSeen: d.LastSeen,
		})
		seen[d.DeviceID] = true
	}
	// Live connections without a devices-table row. Only a code-less connection (DeviceID=="" —
	// a global-SSO browser before 代铸, keyed by name) is surfaced here, by its declared name.
	//
	// A connection WITH a device_id but no row is deliberately NOT surfaced: its row is either
	// deleted (a device that was just revoked but reconnected on a still-valid access token —
	// the "需要移除两次" zombie) or a 代铸 row that hasn't landed yet. Hiding it stops the revoked
	// device from reappearing as a phantom online peer; a legitimate re-login mints a fresh
	// devices row and shows normally via the loop above.
	for _, c := range live {
		if c.DeviceID != "" {
			continue
		}
		items = append(items, PresenceDevice{
			Name: c.Name, Type: c.Type, Online: true, LastSeen: now,
		})
	}

	ev := Event{
		Type: "presence",
		Data: map[string]any{"devices": items},
	}
	for _, c := range live {
		select {
		case c.ch <- ev:
		default:
			slog.Warn("presence: client buffer full; dropping",
				"user", c.UserID, "device", c.DeviceID)
		}
	}
}