SimpleRemoter/server/go/web/ws.go

package web

import (
	"encoding/json"
	"net"
	"net/http"
	"net/url"
	"strings"
	"sync"
	"time"

	"github.com/gorilla/websocket"
	"github.com/yuanyuanxiang/SimpleRemoter/server/go/devicehistory"
	"github.com/yuanyuanxiang/SimpleRemoter/server/go/hub"
	"github.com/yuanyuanxiang/SimpleRemoter/server/go/logger"
	"github.com/yuanyuanxiang/SimpleRemoter/server/go/wsauth"
)

// ----- WS framing knobs ---------------------------------------------------

const (
	wsWriteWait  = 10 * time.Second // single-frame write deadline
	wsReadLimit  = 1 << 20          // refuse incoming frames over 1 MB
	wsSendBuffer = 64               // outbound queue depth per client
)

// baseUpgrader carries the buffer-size config shared by all WS upgrades.
// CheckOrigin is set per-hub in wsHub.upgradeWS so the allowlist is
// closed-over per instance instead of being a global mutable.
var baseUpgrader = websocket.Upgrader{
	ReadBufferSize:  4096,
	WriteBufferSize: 4096,
}

// ----- per-connection client state ----------------------------------------

// wsMsg is one queued WebSocket frame. binary toggles between
// websocket.TextMessage (JSON signaling) and websocket.BinaryMessage
// (screen frames).
type wsMsg struct {
	binary bool
	data   []byte
}

type wsClient struct {
	conn   *websocket.Conn
	send   chan wsMsg
	closed chan struct{}
	once   sync.Once

	// Immutable connection metadata (set once in serve).
	id          string // short random id for audit correlation
	addr        string // resolved client IP for logs / rate limits
	peerAddr    string // actual TCP peer (differs from addr behind a proxy)
	origin      string
	userAgent   string
	connectedAt time.Time

	// Mutated under wsHub.mu (or only by the read loop owning this client).
	nonce        string // outstanding challenge — cleared after a successful login
	token        string // set once authenticated
	role         string // mirrors session role after login
	user         string // mirrors session username after login/resume
	termWatching string // device ID for an open web terminal session, "" otherwise

	// watchIDs is the set of devices this browser is streaming (the screen
	// wall watches several at once). Empty while on the device list.
	// Guarded by wsHub.mu.
	watchIDs map[string]struct{}

	// watchStartedAt records when each watch began so stop events can carry a
	// duration. Guarded by wsHub.mu.
	watchStartedAt map[string]time.Time

	// keyframeAt throttles keyframe requests to one per device per second.
	// Unix seconds, guarded by wsHub.mu.
	keyframeAt map[string]int64

	// Terminal session bookkeeping. Guarded by wsHub.mu.
	termStartedAt   time.Time
	termCloseReason string

	// Audit counters, owned by the read-loop goroutine and read in unregister
	// after that goroutine returns.
	statCommands      int
	statInvalidToken  int
	statLoginFailures int
	statAdminDenied   int
	statUnknownSalt   int
	statUnknownCmd    int

	auditInvalidTokenLogged bool
	auditUnknownSaltLogged  bool
	auditUnknownCmdLogged   bool
}

// queue writes a JSON text frame onto the send buffer. Drops silently if the
// buffer is full so a stuck reader can't back-pressure the broadcast path.
func (c *wsClient) queue(payload []byte) {
	c.enqueue(wsMsg{binary: false, data: payload})
}

// queueBinary writes a binary WS frame. Used for screen-stream packets.
func (c *wsClient) queueBinary(payload []byte) {
	c.enqueue(wsMsg{binary: true, data: payload})
}

func (c *wsClient) enqueue(m wsMsg) {
	select {
	case c.send <- m:
	case <-c.closed:
	default:
		// queue full — drop (acceptable for video; signaling clients are
		// typically not behind enough for the small text buffer to fill).
	}
}

// close signals both loops to exit. Safe to call multiple times.
func (c *wsClient) close() {
	c.once.Do(func() {
		close(c.closed)
		_ = c.conn.Close()
	})
}

// challengeCaps advertises optional protocol features. It is appended to every
// challenge, including the rotation after a failed login, because the browser
// overwrites its capability state on each challenge. Old frontends ignore it.
const challengeCaps = `,"caps":{"multiWatch":true,"keyframeRequest":true,"offlineDevices":true,"auditLog":true}`

// maxWatchPerClient bounds the number of devices one WebSocket may stream at
// once. The wall tops out at 25 tiles; the slack covers a focus switch racing
// a teardown. This keeps the per-frame broadcast loop cheap for a misbehaving
// or buggy client.
const maxWatchPerClient = 32

// isWatching reports whether this connection subscribed to deviceID.
// Callers must hold wsHub.mu (read or write).
func (c *wsClient) isWatching(deviceID string) bool {
	_, ok := c.watchIDs[deviceID]
	return ok
}

// ----- ws hub: registry of all connected browsers -------------------------

type wsHub struct {
	auth    *wsauth.Authenticator
	devices *hub.Hub
	history *devicehistory.Store // nil → device_list stays online-only
	log     *logger.Logger

	mu      sync.RWMutex
	clients map[*wsClient]struct{}

	// screenCloseTimers delays CloseScreen by screenCloseGrace after the last
	// viewer disconnects. A reconnect within that window cancels the timer and
	// reuses the live screen sub-conn (hot path), avoiding a full cold-start.
	// Guarded by mu.
	screenCloseTimers map[string]*time.Timer

	unsub func()

	// Hardening knobs wired from server.Config. Nil/empty values mean
	// "no extra restriction" — useful for local dev where the hub is
	// exercised without server.Server wiring up the env-driven defaults.
	allowedOrigins    []string // empty → only same-origin upgrades accepted
	loginIPLimit      *wsauth.RateLimiter
	loginUserLimit    *wsauth.RateLimiter
	trustForwardedFor bool // honor X-Forwarded-For (only when behind a trusted proxy)
}

func newWSHub(auth *wsauth.Authenticator, devices *hub.Hub, history *devicehistory.Store, log *logger.Logger) *wsHub {
	h := &wsHub{
		auth:              auth,
		devices:           devices,
		history:           history,
		log:               log,
		clients:           make(map[*wsClient]struct{}),
		screenCloseTimers: make(map[string]*time.Timer),
	}
	h.unsub = devices.Subscribe(h)
	return h
}

// withOriginAllowlist returns h after installing the explicit Origin
// allowlist. Chainable. Pass empty/nil to keep "same-origin only".
func (h *wsHub) withOriginAllowlist(origins []string) *wsHub {
	h.allowedOrigins = origins
	return h
}

// withLoginRateLimiters wires per-IP and per-username throttles into
// the login flow. Either may be nil to disable that dimension.
func (h *wsHub) withLoginRateLimiters(byIP, byUser *wsauth.RateLimiter) *wsHub {
	h.loginIPLimit = byIP
	h.loginUserLimit = byUser
	return h
}

// withTrustForwardedFor opts in to using the last entry of the
// X-Forwarded-For header as the client IP. Safe only when the server is
// behind a reverse proxy that you control.
func (h *wsHub) withTrustForwardedFor(trust bool) *wsHub {
	h.trustForwardedFor = trust
	return h
}

// checkOrigin decides whether to accept a WebSocket upgrade based on
// the request's Origin header. Same-origin (Origin host == Host) is
// always accepted; explicit allowlist entries cover the
// PWA-from-different-domain or local-dev cases.
//
// An empty Origin header is rejected: a legitimate browser always sends
// it on cross-origin requests, and same-origin requests have it too in
// modern Chrome/Safari/Firefox. Non-browser clients (curl, scripts) that
// omit Origin shouldn't be talking to the WS endpoint anyway.
func (h *wsHub) checkOrigin(r *http.Request) bool {
	origin := r.Header.Get("Origin")
	if origin == "" {
		return false
	}
	u, err := url.Parse(origin)
	if err != nil || u.Host == "" {
		return false
	}
	// Same-origin (Origin host matches the Host the request came in on).
	// Strip any port mismatch: if the server is behind a proxy, Host may
	// not include a port while Origin does (or vice versa), so compare
	// the hostname components.
	originHost := u.Hostname()
	reqHost := stripPort(r.Host)
	if originHost == reqHost && originHost != "" {
		return true
	}
	// Explicit allowlist entries — match Origin in full (scheme + host
	// + port) so a customer can pin exactly one trusted PWA origin.
	for _, allowed := range h.allowedOrigins {
		if allowed == "" {
			continue
		}
		if strings.EqualFold(origin, strings.TrimSpace(allowed)) {
			return true
		}
	}
	return false
}

func stripPort(hostport string) string {
	if h, _, err := net.SplitHostPort(hostport); err == nil {
		return h
	}
	return hostport
}

// clientIP returns the source IP of an HTTP request. By default uses
// r.RemoteAddr (the actual TCP peer); this is the only safe choice when
// the server is directly exposed to the internet, because a malicious
// client can put anything in X-Forwarded-For and would otherwise rotate
// it to evade per-IP rate limits.
//
// When `trustForwardedFor` is true the LAST entry of X-Forwarded-For is
// returned instead — appropriate only when running behind a reverse
// proxy that you control and that overwrites/appends the header (caddy,
// nginx with proper config, etc). Toggled via the YAMA_WEB_TRUST_PROXY
// env var at startup.
func clientIP(r *http.Request, trustForwardedFor bool) string {
	if trustForwardedFor {
		if xff := r.Header.Get("X-Forwarded-For"); xff != "" {
			// Take the LAST entry — the closest hop, i.e. our own
			// trusted proxy's view of the peer. Trusting the first
			// entry would let a malicious client at the head of the
			// chain set an arbitrary value.
			parts := strings.Split(xff, ",")
			ip := strings.TrimSpace(parts[len(parts)-1])
			if ip != "" {
				return ip
			}
		}
	}
	return stripPort(r.RemoteAddr)
}

// stop unsubscribes from the device hub. Existing connections keep running
// until they close on their own; we only block new event delivery.
func (h *wsHub) stop() {
	if h.unsub != nil {
		h.unsub()
		h.unsub = nil
	}
}

// hub.EventHandler — invoked from hub.Register / hub.Unregister.
func (h *wsHub) OnDeviceOnline(_ hub.DeviceInfo) {
	h.broadcastAuthenticated(`{"cmd":"devices_changed"}`)
}

func (h *wsHub) OnDeviceOffline(id string) {
	// Tell everyone authenticated to refresh their device list (covers
	// users sitting on the devices-page).
	h.broadcastAuthenticated(`{"cmd":"devices_changed"}`)

	// Also tell any browser actively viewing this device's screen — the
	// devices_changed handler only refreshes the list page; viewers on the
	// screen page would otherwise see a frozen frame with "Connected" status
	// indefinitely. The browser handles this by showing "Device offline" and
	// bouncing back to the device list after 2 s.
	if id == "" {
		return
	}
	msg := mustJSON(map[string]any{
		"cmd": "device_offline",
		"id":  id,
	})
	// Drop the id from every watch set as well: the device is gone, so this
	// connection no longer holds a subscription for it. The browser does the
	// same with its local refcount. Without this, a stale entry would keep
	// countWatchers non-zero and later block the screen teardown.
	type viewer struct {
		c     *wsClient
		start time.Time
	}
	var stopped []viewer
	h.mu.Lock()
	for c := range h.clients {
		if !c.isWatching(id) {
			continue
		}
		delete(c.watchIDs, id)
		delete(c.keyframeAt, id)
		if c.token != "" {
			c.queue(msg)
		}
		if !c.watchStartedAt[id].IsZero() {
			stopped = append(stopped, viewer{c, c.watchStartedAt[id]})
		}
		delete(c.watchStartedAt, id)
	}
	h.mu.Unlock()
	// Audit outside h.mu: h.audit takes h.mu.RLock.
	for _, v := range stopped {
		h.audit(v.c, false, auditScreenStop, map[string]any{
			"device_id":   id,
			"duration_ms": durationMs(v.start),
			"reason":      "device_offline",
		})
	}
}

// OnCursorChange relays the remote cursor index to every viewer of this
// device. The browser maps the index to a CSS cursor (desktop) or overlay
// SVG variant (touch). Hub already de-duplicates so we always have a real
// transition here.
func (h *wsHub) OnCursorChange(deviceID string, index byte) {
	msg := mustJSON(map[string]any{
		"cmd":   "cursor",
		"id":    deviceID,
		"index": index,
	})
	h.mu.RLock()
	defer h.mu.RUnlock()
	for c := range h.clients {
		if c.isWatching(deviceID) && c.token != "" {
			c.queue(msg)
		}
	}
}

// OnResolutionChange notifies viewers so the browser-side WebCodecs decoder
// can be (re)initialized with the right frame size. Without this, incoming
// binary frames after connect_result are decoded by an uninitialized
// VideoDecoder and the page stays on "Waiting for video...".
func (h *wsHub) OnResolutionChange(deviceID string, width, height int) {
	msg := mustJSON(map[string]any{
		"cmd":    "resolution_changed",
		"id":     deviceID,
		"width":  width,
		"height": height,
	})
	h.mu.RLock()
	defer h.mu.RUnlock()
	for c := range h.clients {
		if c.isWatching(deviceID) && c.token != "" {
			c.queue(msg)
		}
	}
}

// OnClipboardText relays a device's clipboard text back to every browser
// watching it (in response to a "clipboard_get" request).
func (h *wsHub) OnClipboardText(deviceID string, text string) {
	msg := mustJSON(map[string]any{
		"cmd":  "clipboard_text",
		"id":   deviceID,
		"text": text,
	})
	h.mu.RLock()
	defer h.mu.RUnlock()
	for c := range h.clients {
		if c.isWatching(deviceID) && c.token != "" {
			c.queue(msg)
		}
	}
}

// OnScreenFrame ships a screen packet to every browser currently watching
// this device. We hold the read lock for the whole iteration, but each
// queueBinary is non-blocking (drops on backpressure) so a slow viewer
// cannot stall the fast ones.
func (h *wsHub) OnScreenFrame(deviceID string, packet []byte, _ bool) {
	h.mu.RLock()
	defer h.mu.RUnlock()
	for c := range h.clients {
		if c.isWatching(deviceID) && c.token != "" {
			c.queueBinary(packet)
		}
	}
}

// OnTerminalReady notifies the requesting browser that its term_open
// handshake completed. mode is "pty" or "legacy" — xterm.js disables the
// resize callback in legacy mode (no PTY behind the cmd pipe).
func (h *wsHub) OnTerminalReady(deviceID string, isPTY bool) {
	mode := "legacy"
	if isPTY {
		mode = "pty"
	}
	msg := mustJSON(map[string]any{
		"cmd":  "term_ready",
		"id":   deviceID,
		"mode": mode,
	})
	h.mu.RLock()
	defer h.mu.RUnlock()
	for c := range h.clients {
		if c.termWatching == deviceID && c.token != "" {
			c.queue(msg)
		}
	}
}

// OnTerminalData ships one chunk of raw shell output (already wrapped in
// the "TRM1" magic header) over the binary WS frame. Single-viewer is
// enforced upstream so at most one client matches per device.
func (h *wsHub) OnTerminalData(deviceID string, packet []byte) {
	h.mu.RLock()
	defer h.mu.RUnlock()
	for c := range h.clients {
		if c.termWatching == deviceID && c.token != "" {
			c.queueBinary(packet)
		}
	}
}

// OnTerminalClosed fires when the device's shell exits or the sub-conn
// drops. The browser closes its xterm panel. We also clear termWatching
// so a subsequent term_open from the same browser isn't rejected as
// "already open" by stale state.
func (h *wsHub) OnTerminalClosed(deviceID string, reason string) {
	msg := mustJSON(map[string]any{
		"cmd":    "term_closed",
		"ok":     true,
		"reason": reason,
	})
	type closed struct {
		c     *wsClient
		start time.Time
		why   string
	}
	var done []closed
	h.mu.Lock()
	for c := range h.clients {
		if c.termWatching == deviceID && c.token != "" {
			why := c.termCloseReason
			if why == "" {
				why = reason
			}
			done = append(done, closed{c, c.termStartedAt, why})
			c.termWatching = ""
			c.termStartedAt = time.Time{}
			c.termCloseReason = ""
			c.queue(msg)
		}
	}
	h.mu.Unlock()
	// Audit outside h.mu: h.audit takes h.mu.RLock.
	for _, d := range done {
		h.audit(d.c, false, auditTermStop, map[string]any{
			"device_id":   deviceID,
			"duration_ms": durationMs(d.start),
			"reason":      d.why,
		})
	}
}

// OnDeviceUpdate forwards heartbeat-derived liveness data so the device-list
// rows can refresh RTT and active-window labels without re-fetching.
func (h *wsHub) OnDeviceUpdate(id string, rtt int, activeWindow string) {
	payload := mustJSON(map[string]any{
		"cmd":          "device_update",
		"id":           id,
		"rtt":          rtt,
		"activeWindow": activeWindow,
	})
	h.mu.RLock()
	defer h.mu.RUnlock()
	for c := range h.clients {
		if c.token != "" {
			c.queue(payload)
		}
	}
}

func (h *wsHub) broadcastAuthenticated(msg string) {
	payload := []byte(msg)
	h.mu.RLock()
	defer h.mu.RUnlock()
	for c := range h.clients {
		if c.token != "" {
			c.queue(payload)
		}
	}
}

func (h *wsHub) register(c *wsClient) {
	h.mu.Lock()
	h.clients[c] = struct{}{}
	h.mu.Unlock()
}

func (h *wsHub) unregister(c *wsClient) {
	h.mu.Lock()
	delete(h.clients, c)
	// Snapshot the watch set and clear it under the same lock. The teardown
	// below must run OUTSIDE the lock: countWatchers takes h.mu.RLock and
	// deferredCloseScreen takes h.mu.Lock.
	watched := make([]string, 0, len(c.watchIDs))
	started := make(map[string]time.Time, len(c.watchIDs))
	for id := range c.watchIDs {
		watched = append(watched, id)
		started[id] = c.watchStartedAt[id]
	}
	c.watchIDs = nil
	c.watchStartedAt = nil
	c.keyframeAt = nil

	termDevice := c.termWatching
	termStart := c.termStartedAt
	c.termWatching = ""
	c.termStartedAt = time.Time{}
	c.termCloseReason = ""
	h.mu.Unlock()

	// If this client was the last viewer of a device, tear down the screen
	// session so the device stops encoding.
	for _, id := range watched {
		if !started[id].IsZero() {
			h.audit(c, false, auditScreenStop, map[string]any{
				"device_id":   id,
				"duration_ms": durationMs(started[id]),
				"reason":      "ws_close",
			})
		}
		if h.countWatchers(id) == 0 {
			h.deferredCloseScreen(id)
		}
	}
	// Terminal sessions are single-viewer by design, so any open session
	// belongs to this client. Tear it down so the next viewer doesn't
	// hit ErrTerminalBusy from an abandoned session. The client is already
	// out of h.clients, so OnTerminalClosed won't reach it — audit here.
	if termDevice != "" {
		h.devices.CloseTerminalSession(termDevice)
		h.audit(c, false, auditTermStop, map[string]any{
			"device_id":   termDevice,
			"duration_ms": durationMs(termStart),
			"reason":      "ws_close",
		})
	}
	h.auditClose(c, len(watched))
	// Do NOT revoke the token: tokens are session-scoped, not WS-scoped.
	// Frontend may close+reopen the WS at any time (visibilitychange handler,
	// brief network blip, reload) and must be able to resume with the same
	// cached token. The token expires on its own TTL.
	c.close()
}

// screenCloseGrace is how long we keep a device's screen sub-conn alive after
// the last browser viewer disconnects. A reconnect within this window (e.g.
// the user briefly switched to a TOTP app on mobile) cancels the timer and
// resumes on the hot path — no cold-start, no waiting for a new IDR.
const screenCloseGrace = 30 * time.Second

// deferredCloseScreen schedules CloseScreen after screenCloseGrace unless a
// viewer reconnects first. Must be called without h.mu held.
func (h *wsHub) deferredCloseScreen(deviceID string) {
	h.mu.Lock()
	if t, ok := h.screenCloseTimers[deviceID]; ok {
		t.Stop()
	}
	h.screenCloseTimers[deviceID] = time.AfterFunc(screenCloseGrace, func() {
		h.mu.Lock()
		delete(h.screenCloseTimers, deviceID)
		h.mu.Unlock()
		if h.countWatchers(deviceID) == 0 {
			h.devices.CloseScreen(deviceID)
		}
	})
	h.mu.Unlock()
}

// ----- HTTP handler -------------------------------------------------------

func (h *wsHub) serve(w http.ResponseWriter, r *http.Request) {
	// Reject disallowed origins explicitly so the rejection is auditable.
	// Upgrader.CheckOrigin below is kept as defense in depth.
	if !h.checkOrigin(r) {
		fields := requestFields(r, h.trustForwardedFor)
		fields["reason"] = "origin_not_allowed"
		h.log.AuditWarn(auditWSReject, fields)
		http.Error(w, "origin not allowed", http.StatusForbidden)
		return
	}

	// Build a per-call upgrader so CheckOrigin closes over this hub's
	// allowlist instead of a package-level mutable.
	upgrader := baseUpgrader
	upgrader.CheckOrigin = h.checkOrigin

	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		fields := requestFields(r, h.trustForwardedFor)
		fields["reason"] = "upgrade_failed"
		fields["error"] = err.Error()
		h.log.AuditWarn(auditWSReject, fields)
		return
	}
	conn.SetReadLimit(wsReadLimit)

	nonce, err := wsauth.NewNonce()
	if err != nil {
		h.log.Error("nonce gen: %v", err)
		_ = conn.Close()
		return
	}

	client := &wsClient{
		conn:           conn,
		send:           make(chan wsMsg, wsSendBuffer),
		closed:         make(chan struct{}),
		id:             newConnID(),
		addr:           clientIP(r, h.trustForwardedFor),
		peerAddr:       stripPort(r.RemoteAddr),
		origin:         r.Header.Get("Origin"),
		userAgent:      r.UserAgent(),
		connectedAt:    time.Now(),
		nonce:          nonce,
		watchIDs:       make(map[string]struct{}),
		watchStartedAt: make(map[string]time.Time),
		keyframeAt:     make(map[string]int64),
	}
	h.register(client)
	h.audit(client, false, auditWSConnect, client.wsRequestFields())
	defer h.unregister(client)

	go h.writeLoop(client)

	// Greet with a challenge nonce so the browser can compute the login response.
	client.queue([]byte(`{"cmd":"challenge","nonce":"` + nonce + `"` + challengeCaps + `}`))

	h.readLoop(client)
}

// writeLoop drains the send queue. Exits when the channel is closed or a
// write fails. Closing the underlying connection is the read loop's job.
func (h *wsHub) writeLoop(c *wsClient) {
	for {
		select {
		case msg := <-c.send:
			msgType := websocket.TextMessage
			if msg.binary {
				msgType = websocket.BinaryMessage
			}
			_ = c.conn.SetWriteDeadline(time.Now().Add(wsWriteWait))
			if err := c.conn.WriteMessage(msgType, msg.data); err != nil {
				c.close()
				return
			}
		case <-c.closed:
			return
		}
	}
}

// readLoop dispatches incoming messages. Exits on read error (peer closed,
// timeout, malformed frame, etc.), which then triggers unregister cleanup.
func (h *wsHub) readLoop(c *wsClient) {
	for {
		_, raw, err := c.conn.ReadMessage()
		if err != nil {
			return
		}
		var env struct {
			Cmd string `json:"cmd"`
		}
		if err := json.Unmarshal(raw, &env); err != nil {
			continue // ignore garbage frames
		}
		h.dispatch(c, env.Cmd, raw)
	}
}