Files
SimpleRemoter/server/go/web/ws_handlers.go

287 lines
9.6 KiB
Go

package web
import (
"encoding/json"
"github.com/yuanyuanxiang/SimpleRemoter/server/go/hub"
"github.com/yuanyuanxiang/SimpleRemoter/server/go/protocol"
)
// dispatch routes one inbound message to its handler. The `raw` payload is
// passed through so handlers can re-parse to their own shape.
//
// Phase 3 implements: get_salt, login, get_devices, ping, disconnect.
// Phase 4/5/6 commands (connect, mouse, key, term_*, etc.) get a friendly
// "not yet implemented" reply so the browser UI doesn't hang silently.
func (h *wsHub) dispatch(c *wsClient, cmd string, raw []byte) {
switch cmd {
case "get_salt":
h.handleGetSalt(c, raw)
case "login":
h.handleLogin(c, raw)
case "get_devices":
h.handleGetDevices(c, raw)
case "ping":
// no-op heartbeat; the read itself was the keep-alive signal
case "disconnect":
h.handleDisconnect(c, raw)
case "connect":
h.handleConnect(c, raw)
case "rdp_reset":
// silently ignored — UI uses this as a fire-and-forget
case "mouse", "key":
// silently ignored — no remote screen yet
case "term_open":
h.replyNotImplemented(c, "term_closed", "Web terminal not yet implemented on Go server")
case "term_input", "term_resize", "term_close":
// silently ignored — no terminal session
// Admin operations (Phase 7).
case "create_user":
h.replyNotImplemented(c, "create_user_result", "User management not yet implemented")
case "delete_user":
h.replyNotImplemented(c, "delete_user_result", "User management not yet implemented")
case "list_users":
h.replyNotImplemented(c, "list_users_result", "User management not yet implemented")
case "get_groups":
c.queue([]byte(`{"cmd":"groups","ok":true,"groups":[]}`))
}
}
func (h *wsHub) replyNotImplemented(c *wsClient, replyCmd, msg string) {
c.queue(mustJSON(map[string]any{
"cmd": replyCmd,
"ok": false,
"msg": msg,
}))
}
// ----- handlers ------------------------------------------------------------
func (h *wsHub) handleGetSalt(c *wsClient, raw []byte) {
var in struct {
Username string `json:"username"`
}
_ = json.Unmarshal(raw, &in)
salt, ok := h.auth.GetSalt(in.Username)
// Do not leak which usernames exist: always return ok=true with a salt.
// For unknown users hand back the empty salt (matches admin convention)
// so the timing/shape of the response is uniform.
if !ok {
salt = ""
}
c.queue(mustJSON(map[string]any{
"cmd": "salt",
"ok": true,
"salt": salt,
}))
}
func (h *wsHub) handleLogin(c *wsClient, raw []byte) {
var in struct {
Username string `json:"username"`
Response string `json:"response"`
Nonce string `json:"nonce"`
}
if err := json.Unmarshal(raw, &in); err != nil {
c.queue(mustJSON(map[string]any{"cmd": "login_result", "ok": false, "msg": "Invalid request"}))
return
}
// Bind the response to the challenge we issued at connect time so that
// replays from a different connection can't reuse a captured response.
if in.Nonce == "" || in.Nonce != c.nonce {
c.queue(mustJSON(map[string]any{"cmd": "login_result", "ok": false, "msg": "Invalid challenge"}))
return
}
token, role, err := h.auth.VerifyLogin(in.Username, in.Response, in.Nonce)
if err != nil {
// Burn the challenge on failure too — forces a new round on retry.
c.nonce = ""
c.queue(mustJSON(map[string]any{"cmd": "login_result", "ok": false, "msg": "Invalid credentials"}))
return
}
c.nonce = ""
c.token = token
c.role = role
h.log.Info("ws login: user=%s role=%s addr=%s", in.Username, role, c.addr)
c.queue(mustJSON(map[string]any{
"cmd": "login_result",
"ok": true,
"token": token,
"role": role,
}))
}
// handleConnect kicks off a screen-sharing session for the browser. We send
// COMMAND_SCREEN_SPY to the device's main TCP connection; the device then
// opens a new sub-connection (TOKEN_BITMAPINFO) which the TCP side binds to
// the device via hub.BindScreenConn. Frame relay to the browser is handled
// in Phase 4.2 once frames actually arrive.
//
// Reply semantics: returning connect_result.ok=true (without width/height)
// triggers the browser's "Waiting for video..." spinner. We can't deliver
// width/height here because we don't yet know them — they show up in the
// first TOKEN_BITMAPINFO from the device.
// handleDisconnect detaches this client from any device it was watching and
// — if no other authenticated client is still watching — closes the device's
// screen sub-connection. Closing the TCP sub-conn is the signal the C++
// device firmware uses to stop screen capture, so this is how we ask the
// device to free its encoder.
func (h *wsHub) handleDisconnect(c *wsClient, _ []byte) {
// Mirror handleConnect: take h.mu so event-handler readers
// (OnResolutionChange/OnScreenFrame) get a consistent view of c.watching.
h.mu.Lock()
prev := c.watching
c.watching = ""
h.mu.Unlock()
c.queue([]byte(`{"cmd":"disconnect_result","ok":true}`))
if prev != "" && h.countWatchers(prev) == 0 {
h.devices.CloseScreen(prev)
}
}
// countWatchers returns how many authenticated clients still have their
// `watching` field pointing at deviceID. Called from disconnect paths.
func (h *wsHub) countWatchers(deviceID string) int {
h.mu.RLock()
defer h.mu.RUnlock()
n := 0
for c := range h.clients {
if c.watching == deviceID {
n++
}
}
return n
}
func (h *wsHub) handleConnect(c *wsClient, raw []byte) {
if !h.requireAuth(c, raw, "connect_result") {
return
}
var in struct {
ID string `json:"id"`
}
if err := json.Unmarshal(raw, &in); err != nil || in.ID == "" {
c.queue(mustJSON(map[string]any{"cmd": "connect_result", "ok": false, "msg": "Bad request"}))
return
}
// If a screen session is already live for this device (another browser
// is already watching), reuse it: hand the new viewer the current
// resolution and the most recent IDR keyframe so its decoder can start
// rendering immediately, without waiting for the next IDR (~15 s).
cache := h.devices.ScreenState(in.ID)
if cache.Active {
c.queue(mustJSON(map[string]any{
"cmd": "connect_result", "ok": true,
"width": cache.Width, "height": cache.Height,
}))
if len(cache.Keyframe) > 0 {
c.queueBinary(cache.Keyframe)
}
h.mu.Lock()
c.watching = in.ID
h.mu.Unlock()
return
}
// No active session — kick the device to start capturing. We send the
// same 32-byte COMMAND_SCREEN_SPY payload the C++ WebService sends:
// [0]=COMMAND_SCREEN_SPY, [1]=0 (GDI), [2]=ALGORITHM_H264, [3]=1 (multi-screen),
// [4..31]=0.
cmd := make([]byte, 32)
cmd[0] = protocol.CommandScreenSpy
cmd[2] = protocol.AlgorithmH264
cmd[3] = 1
// CRITICAL: bind c.watching BEFORE asking the device to start capturing.
// On fast reconnects the device's screen sub-conn handshake completes in
// <100 ms, so TOKEN_BITMAPINFO and even the first H264 frame can arrive
// before this handler finishes — and the resolution_changed / frame
// broadcasts in wsHub filter on c.watching. With the assignment after
// SendToDevice the new viewer silently misses the very first IDR and
// resolution_changed, leaving the page stuck on "Waiting for video".
//
// The write needs to share the lock event handlers use to read c.watching
// (they iterate h.clients under h.mu.RLock). Without that the write is a
// data race; on a fast reconnect the reader goroutine can keep observing
// the previous value ("") long enough to drop the first resolution_changed
// and the first IDR, which produces the exact "every other quick reconnect
// goes black" symptom — the C++ server avoids it because it does the same
// state mutation under std::mutex and reaps the memory-barrier as a bonus.
h.mu.Lock()
c.watching = in.ID
h.mu.Unlock()
if err := h.devices.SendToDevice(in.ID, cmd); err != nil {
// Roll back the watching flag if we never managed to kick capture.
h.mu.Lock()
c.watching = ""
h.mu.Unlock()
msg := "Device offline"
if err != hub.ErrDeviceOffline {
msg = "Failed to start screen capture"
h.log.Error("SendToDevice(%s): %v", in.ID, err)
}
c.queue(mustJSON(map[string]any{"cmd": "connect_result", "ok": false, "msg": msg}))
return
}
h.log.Info("[timing] COMMAND_SCREEN_SPY sent to device=%s (cold start)", in.ID)
c.queue(mustJSON(map[string]any{"cmd": "connect_result", "ok": true}))
}
func (h *wsHub) handleGetDevices(c *wsClient, raw []byte) {
if !h.requireAuth(c, raw, "device_list") {
return
}
devices := h.devices.ListDevices()
c.queue(mustJSON(map[string]any{
"cmd": "device_list",
"ok": true,
"devices": devices,
}))
}
// requireAuth validates the token embedded in raw against the authenticator's
// session store (not against c.token). Tokens live independently of WS
// connections — the browser may reconnect after a visibility/network blip and
// resume with the same token, so we must not tie validity to one WS lifetime.
// On the first authenticated message we cache the token/role on the wsClient
// so broadcasts know to deliver to this connection.
func (h *wsHub) requireAuth(c *wsClient, raw []byte, replyCmd string) bool {
var in struct {
Token string `json:"token"`
}
_ = json.Unmarshal(raw, &in)
if in.Token == "" {
c.queue(mustJSON(map[string]any{"cmd": replyCmd, "ok": false}))
return false
}
sess, err := h.auth.ValidateToken(in.Token)
if err != nil {
c.queue(mustJSON(map[string]any{"cmd": replyCmd, "ok": false}))
return false
}
if c.token == "" {
c.token = in.Token
c.role = sess.Role
}
return true
}
func mustJSON(v any) []byte {
b, err := json.Marshal(v)
if err != nil {
// All callers pass simple map[string]any with primitive values;
// marshal can't realistically fail. If it does, return a safe fallback.
return []byte(`{"cmd":"error","msg":"internal encode error"}`)
}
return b
}