SimpleRemoter/server/go/web/ws.go

package web

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

	"github.com/gorilla/websocket"
	"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
)

// upgrader allows any origin — this service is meant to be tunneled through
// frp, so requests can legitimately arrive from arbitrary front-end hosts.
// Adjust CheckOrigin once we have a deployment story.
var upgrader = websocket.Upgrader{
	ReadBufferSize:  4096,
	WriteBufferSize: 4096,
	CheckOrigin:     func(r *http.Request) bool { return true },
}

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

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

	// 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
	addr  string // client address for logs
}

// queue writes a payload 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) {
	select {
	case c.send <- payload:
	case <-c.closed:
	default:
		// queue full — caller is responsible for noticing if it matters.
	}
}

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

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

type wsHub struct {
	auth    *wsauth.Authenticator
	devices *hub.Hub
	log     *logger.Logger

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

	unsub func()
}

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

// 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(_ string) {
	h.broadcastAuthenticated(`{"cmd":"devices_changed"}`)
}

// 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)
	h.mu.Unlock()
	// 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()
}

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

func (h *wsHub) serve(w http.ResponseWriter, r *http.Request) {
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		h.log.Error("ws upgrade: %v", err)
		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 []byte, wsSendBuffer),
		closed: make(chan struct{}),
		nonce:  nonce,
		addr:   r.RemoteAddr,
	}
	h.register(client)
	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 + `"}`))

	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:
			_ = c.conn.SetWriteDeadline(time.Now().Add(wsWriteWait))
			if err := c.conn.WriteMessage(websocket.TextMessage, msg); 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)
	}
}