Feature(Go): Web terminal relay with PTY mode and graceful close (Phase 6)

server/go/cmd/main.go

@@ -37,17 +37,33 @@ func (h *MyHandler) OnConnect(ctx *connection.Context) {
 
 // OnDisconnect is called when a client disconnects
 func (h *MyHandler) OnDisconnect(ctx *connection.Context) {
-	// Always clean up any screen sub-context mapping first — the connection
-	// may be a screen sub-conn (which has no ClientInfo) rather than a main
-	// login connection. UnbindScreenConn is a no-op if not tracked.
+	// Always clean up any sub-context mapping first — the connection may
+	// be a screen / terminal sub-conn rather than a main login connection.
+	// Both Unbind* calls are no-ops if not tracked. UnbindTerminalConn
+	// also fires OnTerminalClosed so the browser sees the session end on
+	// unexpected device-side drops.
 	h.hub.UnbindScreenConn(ctx)
+	h.hub.UnbindTerminalConn(ctx)
 
 	info := ctx.GetInfo()
-	if info.ClientID != "" {
+	// Only treat this disconnect as a device-going-offline event if this
+	// ctx is the device's MAIN login connection. Phase 6 added ClientID
+	// pinning to sub-conns (via ConnAuth — needed for terminal routing),
+	// so a non-empty ClientID alone no longer distinguishes main from
+	// sub. Closing a screen / terminal sub-conn must NOT remove the
+	// device from the hub.
+	if info.ClientID != "" && h.hub.MainConn(info.ClientID) == ctx {
 		h.log.ClientEvent("offline", ctx.ID, ctx.GetPeerIP(),
 			"clientID", info.ClientID,
 			"computer", info.ComputerName,
 		)
+		// Tear down any active sub-conn sessions BEFORE Unregister so the
+		// browser sees screen/terminal close events alongside the
+		// device-offline event, instead of frames/output continuing to
+		// stream from orphaned sub-conn ctxs until they time out on
+		// their own. Both calls no-op if there's no active session.
+		h.hub.CloseScreen(info.ClientID)
+		h.hub.CloseTerminalSession(info.ClientID)
 		h.hub.Unregister(info.ClientID)
 	}
 }
@@ -58,6 +74,27 @@ func (h *MyHandler) OnReceive(ctx *connection.Context, data []byte) {
 		return
 	}
 
+	// Terminal-bound sub-conns deliver RAW shell output with no leading
+	// command byte — see client/ConPTYManager.cpp:328 (Send2Server with
+	// just the buffer). We must short-circuit BEFORE the command switch
+	// or the first output byte will be misinterpreted as a token.
+	// Exception: a length-1 packet whose byte is TOKEN_TERMINAL_CLOSE
+	// is the device's "shell exited" notification, NOT data.
+	if devID := h.hub.TerminalDeviceID(ctx); devID != "" {
+		if len(data) == 1 && data[0] == protocol.TokenTerminalClose {
+			h.log.Info("terminal closed by device=%s conn=%d", devID, ctx.ID)
+			h.hub.CloseTerminalSession(devID)
+			return
+		}
+		// Wrap with the 'TRM1' magic the browser uses to demultiplex
+		// terminal output from screen frames over the shared WS.
+		packet := make([]byte, 4+len(data))
+		copy(packet[:4], protocol.TerminalBinaryMagic[:])
+		copy(packet[4:], data)
+		h.hub.PublishTerminalData(devID, packet)
+		return
+	}
+
 	cmd := data[0]
 	// Handle commands
 	switch cmd {
@@ -71,6 +108,17 @@ func (h *MyHandler) OnReceive(ctx *connection.Context, data []byte) {
 		h.handleConnAuth(ctx, data)
 	case protocol.TokenBitmapInfo:
 		h.handleBitmapInfo(ctx, data)
+	case protocol.TokenTerminalStart:
+		h.handleTerminalStart(ctx, true)
+	case protocol.TokenShellStart:
+		h.handleTerminalStart(ctx, false)
+	case protocol.TokenTerminalClose:
+		// Pre-bind close (rare — device gives up before the server
+		// finished its half of the handshake). Best-effort cleanup.
+		if devID := h.deviceIDOfSubConn(ctx); devID != "" {
+			h.log.Info("pre-bind terminal close: device=%s conn=%d", devID, ctx.ID)
+			h.hub.CloseTerminalSession(devID)
+		}
 	case protocol.TokenFirstScreen:
 		// TOKEN_FIRSTSCREEN delivers a RAW BGRA baseline frame, not an
 		// H264 unit — bytes ≈ width × height × 4. The C++ MFC dialog
@@ -107,7 +155,24 @@ func (h *MyHandler) OnReceive(ctx *connection.Context, data []byte) {
 // and the signing primitive lives in a vendored component out of scope
 // for this server, so we always reply OK and let TOKEN_BITMAPINFO carry
 // the device ID via offset 41 when the screen sub-conn proceeds.
-func (h *MyHandler) handleConnAuth(ctx *connection.Context, _ []byte) {
+func (h *MyHandler) handleConnAuth(ctx *connection.Context, data []byte) {
+	// Pin the parent device's ClientID onto the sub-conn. Without this,
+	// later 1-byte tokens (TOKEN_TERMINAL_START / TOKEN_SHELL_START) have
+	// no way to identify which device they belong to — they carry no
+	// clientID themselves. ConnAuthPacket layout has clientID at offset 1
+	// (uint64 LE); see common/commands.h::ConnAuthPacket.
+	if len(data) >= protocol.ConnAuthOffClientID+8 {
+		clientID := binary.LittleEndian.Uint64(
+			data[protocol.ConnAuthOffClientID : protocol.ConnAuthOffClientID+8])
+		if clientID != 0 {
+			// Sub-conns never go through handleLogin, so their ctx.Info
+			// is otherwise empty. We only need ClientID for routing.
+			info := ctx.GetInfo()
+			info.ClientID = strconv.FormatUint(clientID, 10)
+			ctx.SetInfo(info)
+		}
+	}
+
 	ack := make([]byte, protocol.ConnAuthAckSize)
 	ack[0] = protocol.TokenConnAuth
 	ack[protocol.ConnAuthAckOffStatus] = protocol.ConnAuthStatusOK
@@ -119,6 +184,55 @@ func (h *MyHandler) handleConnAuth(ctx *connection.Context, _ []byte) {
 	}
 }
 
+// deviceIDOfSubConn resolves the parent device of a sub-conn from the
+// ClientID pinned by handleConnAuth. Returns "" for the rare case of a
+// legacy client that skipped ConnAuth (the Go server's only target is
+// modern clients, so this is effectively a paranoia check).
+func (h *MyHandler) deviceIDOfSubConn(ctx *connection.Context) string {
+	return ctx.GetInfo().ClientID
+}
+
+// handleTerminalStart fires when the device's freshly-spawned shell
+// sub-conn announces itself. TOKEN_TERMINAL_START (232) means PTY mode
+// (Linux/macOS or Windows ConPTY); TOKEN_SHELL_START (128) means the
+// legacy Windows cmd-pipe path. Both packets are 1-byte tokens — the
+// device identity comes from ConnAuth's pinned ClientID.
+//
+// After binding we send:
+//   - For PTY only: an initial CMD_TERMINAL_RESIZE 80x24 so the shell
+//     doesn't render at the PTY default before the browser's first fit.
+//     vim/htop look broken otherwise. The browser will follow up with a
+//     real term_resize once xterm.js sizes the canvas.
+//   - Always: COMMAND_NEXT to unblock the device's read thread (the
+//     ConPTYManager ReadThread sits on m_hEventDlgOpen until then —
+//     see client/ConPTYManager.cpp:259).
+func (h *MyHandler) handleTerminalStart(ctx *connection.Context, isPTY bool) {
+	devID := h.deviceIDOfSubConn(ctx)
+	if devID == "" {
+		h.log.Warn("terminal start with no clientID: conn=%d", ctx.ID)
+		ctx.Close()
+		return
+	}
+	if !h.hub.BindTerminalConn(devID, ctx, isPTY) {
+		// No pending session — this is a stale sub-conn (e.g. browser
+		// gave up and closed term_close already). Drop it.
+		h.log.Warn("orphan terminal sub-conn: device=%s conn=%d isPTY=%v",
+			devID, ctx.ID, isPTY)
+		ctx.Close()
+		return
+	}
+
+	if isPTY {
+		if err := h.srv.Send(ctx, protocol.BuildTerminalResize(80, 24)); err != nil {
+			h.log.Error("initial resize send failed: conn=%d: %v", ctx.ID, err)
+		}
+	}
+	if err := h.srv.Send(ctx, []byte{protocol.CommandNext}); err != nil {
+		h.log.Error("COMMAND_NEXT send failed on terminal: conn=%d: %v", ctx.ID, err)
+	}
+	h.log.Info("terminal bound: device=%s conn=%d isPTY=%v", devID, ctx.ID, isPTY)
+}
+
 // handleBitmapInfo is the first packet on a freshly-arrived screen
 // sub-connection. Packet layout (after the command byte at data[0]):
 //