diff --git a/agent/internal/hpagent/handler.go b/agent/internal/hpagent/handler.go index 7968d3d..0b54785 100644 --- a/agent/internal/hpagent/handler.go +++ b/agent/internal/hpagent/handler.go @@ -37,7 +37,6 @@ func FollowMaster(agent *tsnet.TSAgent) { for scanner.Scan() { line := scanner.Bytes() - log.Info("Got bytes delimited by newline") var msg CborMessage decoder := cbor.NewDecoder(bytes.NewReader(line)) @@ -49,14 +48,29 @@ func FollowMaster(agent *tsnet.TSAgent) { } log.Debug("Received message from master: %s", msg) - var sshPayload sshutil.SSHConnectPayload - err = cbor.Unmarshal(msg.Payload, &sshPayload) - if err != nil { - log.Error("Unable to unmarshal SSH connect payload: %s", err) - continue - } + switch msg.Op { + case "ssh_conn": + var sshPayload sshutil.SSHConnectPayload + err = cbor.Unmarshal(msg.Payload, &sshPayload) + if err != nil { + log.Error("Unable to unmarshal SSH connect payload: %s", err) + continue + } - sshutil.StartWebSSH(agent, sshPayload) + sshutil.StartWebSSH(agent, sshPayload) + continue + + case "ssh_term": + var sshPayload sshutil.SSHClosePayload + err = cbor.Unmarshal(msg.Payload, &sshPayload) + if err != nil { + log.Error("Unable to unmarshal SSH close payload: %s", err) + continue + } + + sshutil.CloseWebSSH(agent, sshPayload) + continue + } } // var msg RecvMessage diff --git a/agent/internal/sshutil/connect.go b/agent/internal/sshutil/connect.go index 0fee7c7..1aa45b1 100644 --- a/agent/internal/sshutil/connect.go +++ b/agent/internal/sshutil/connect.go @@ -18,6 +18,10 @@ type SSHConnectPayload struct { Port int `cbor:"port"` } +type SSHClosePayload struct { + SessionId string `cbor:"sessionId"` +} + func connectToTailscaleSSH(agent *tsnet.TSAgent, params SSHConnectPayload) (*ssh.Client, error) { log := util.GetLogger() addr := strings.Join([]string{params.Hostname, ":", strconv.Itoa(params.Port)}, "") @@ -136,3 +140,27 @@ func StartWebSSH(agent *tsnet.TSAgent, params SSHConnectPayload) { RemoveSession(params.SessionId) }() } + +func CloseWebSSH(agent *tsnet.TSAgent, params SSHClosePayload) { + log := util.GetLogger() + + if agent == nil { + log.Error("tsnet.TSAgent is not initialized correctly") + return + } + + if params.SessionId == "" { + log.Error("Invalid SSH close parameters: %v", params) + return + } + + log.Debug("Closing SSH session for session ID: %s", params.SessionId) + ctx, ok := lookupSession(params.SessionId) + if !ok { + log.Info("No active SSH session found for session ID: %s", params.SessionId) + return + } + + RemoveSession(ctx.ID) + log.Info("SSH session for %s closed", params.SessionId) +} diff --git a/app/routes.ts b/app/routes.ts index b921718..9a80c50 100644 --- a/app/routes.ts +++ b/app/routes.ts @@ -25,6 +25,7 @@ export default [ route('/users', 'routes/users/overview.tsx'), route('/acls', 'routes/acls/overview.tsx'), route('/dns', 'routes/dns/overview.tsx'), + route('/ssh', 'routes/ssh/overview.tsx'), ...prefix('/settings', [ index('routes/settings/overview.tsx'), diff --git a/app/routes/ssh/overview.tsx b/app/routes/ssh/overview.tsx new file mode 100644 index 0000000..3cab350 --- /dev/null +++ b/app/routes/ssh/overview.tsx @@ -0,0 +1,145 @@ +import { decode } from 'cborg'; +import { Loader2 } from 'lucide-react'; +import { useEffect, useRef, useState } from 'react'; +import { LoaderFunctionArgs, data, useLoaderData } from 'react-router'; +import { ClientOnly } from 'remix-utils/client-only'; +import { + Command, + SSHConnectData, + SSHConnectFailedData, +} from '~/server/agent/dispatcher'; +import { useLiveData } from '~/utils/live-data'; +import toast from '~/utils/toast'; +import XTerm from './xterm.client'; + +export async function loader({ request }: LoaderFunctionArgs) { + const qp = new URL(request.url).searchParams; + const username = qp.get('username') || undefined; + const hostname = qp.get('hostname') || undefined; + const port = qp.get('port') + ? Number.parseInt(qp.get('port')!, 10) + : undefined; + + if (!username || !hostname || !port) { + throw data('Missing required parameters: username, hostname, port', 400); + } + + const baseUrl = new URL(request.url).origin; + const wsUrl = new URL('/_ssh_plexer', baseUrl); + wsUrl.protocol = wsUrl.protocol === 'https:' ? 'wss:' : 'ws:'; + wsUrl.searchParams.set('username', username); + wsUrl.searchParams.set('hostname', hostname); + wsUrl.searchParams.set('port', port.toString()); + + return { + socketUrl: wsUrl.toString(), + }; +} +type SessionStatus = 'loading' | 'connected' | 'error'; + +export default function Page() { + const { pause } = useLiveData(); + const { socketUrl } = useLoaderData(); + const [socket, setSocket] = useState(null); + const [status, setStatus] = useState('loading'); + const [sessionId, setSessionId] = useState(null); + + const queue = useRef>([]); + const validated = useRef(false); + + useEffect(() => { + // SSH connections should not use stale while revalidate logic. + pause(); + + const ws = new WebSocket(socketUrl); + ws.binaryType = 'arraybuffer'; + + ws.onopen = () => { + setSocket(ws); + setStatus('loading'); + }; + + // We need to wait for the WebSocket to open and respond with the + // connection ID. Without a session ID, we do not have a mux. + const messageHandler = (event: MessageEvent) => { + if (!(event.data instanceof ArrayBuffer)) { + toast('Invalid message received from server'); + return; + } + + const data = new Uint8Array(event.data); + const obj = decode(data) as Command; + + if (obj.op === 'ssh_conn_successful') { + const data = obj as SSHConnectData; + if (!validated.current) { + validated.current = true; + toast( + `SSH connection established with session ID: ${data.payload.sessionId}`, + ); + + setStatus('connected'); + setSessionId(data.payload.sessionId); + } + + return; + } + + if (obj.op === 'ssh_conn_failed') { + const data = obj as SSHConnectFailedData; + if (!validated.current) { + validated.current = true; + toast(`SSH connection failed: ${data.payload.reason}`); + setStatus('error'); + } + + return; + } + + if (obj.op === 'ssh_frame') { + queue.current.push(new Uint8Array(event.data)); + return; + } + }; + + ws.addEventListener('message', messageHandler); + + ws.onerror = (error) => { + setStatus('error'); + toast(`WebSocket error: ${error}`); + }; + + ws.onclose = () => { + if (status !== 'error') { + toast('SSH connection closed'); + } + + setSocket(null); + setStatus('error'); + }; + + return () => { + ws.removeEventListener('message', messageHandler); + ws.close(); + }; + }, [socketUrl]); + + if (socket === null || !sessionId || status === 'loading') { + return ( + + ); + } + + return ( +
+

Session ID: {sessionId}

+ {queue.current.length > 0 && ( +

+ {queue.current.length} frames queued +

+ )} + + +
+ ); +} diff --git a/app/routes/ssh/xterm.client.tsx b/app/routes/ssh/xterm.client.tsx new file mode 100644 index 0000000..9e2f471 --- /dev/null +++ b/app/routes/ssh/xterm.client.tsx @@ -0,0 +1,91 @@ +import * as xterm from '@xterm/xterm'; +import { useEffect, useRef } from 'react'; +import '@xterm/xterm/css/xterm.css'; +import { decode } from 'cborg'; +import type { SSHFrameData } from '~/server/agent/dispatcher'; +import { useLiveData } from '~/utils/live-data'; + +interface XTermProps { + ws: WebSocket; + sessionId: string; + queue: Array; +} + +const RED = new TextEncoder().encode('\x1b[31m'); +const RESET = new TextEncoder().encode('\x1b[0m'); + +export default function XTerm({ ws, sessionId, queue }: XTermProps) { + const { pause } = useLiveData(); + + const container = useRef(null); + const term = useRef(null); + + useEffect(() => { + pause(); + + const terminal = new xterm.Terminal({ + convertEol: true, + fontSize: 14, + theme: { + background: '#1e1e1e', + foreground: '#ffffff', + }, + }); + + terminal.open(container.current!); + terminal.focus(); + term.current = terminal; + + const handleFrame = (data: Uint8Array) => { + try { + const frame: SSHFrameData = decode(data); + if (frame.op !== 'ssh_frame') { + console.warn('Received unexpected frame type:', frame.op); + return; + } + + // If this is stderr, color it red + if (frame.payload.channel === 2) { + terminal.write( + new Uint8Array([...RED, ...frame.payload.frame, ...RESET]), + ); + } else { + terminal.write(frame.payload.frame); + } + } catch (err) { + console.error('Failed to decode CBOR frame:', err); + } + }; + + for (const buffer of queue) { + handleFrame(buffer); + } + + terminal.onData((input) => { + if (ws.readyState === WebSocket.OPEN) { + ws.send(input); + } else { + console.warn('WebSocket is not open, cannot send data'); + } + }); + + const onMessage = (event: MessageEvent) => { + if (!(event.data instanceof ArrayBuffer)) { + console.warn('Received non-binary message from WebSocket'); + return; + } + + const data = new Uint8Array(event.data); + handleFrame(data); + }; + + ws.addEventListener('message', onMessage); + + return () => { + ws.removeEventListener('message', onMessage); + term.current?.dispose(); + }; + }, [ws, queue]); + + return
; +} diff --git a/app/server/agent/dispatcher.ts b/app/server/agent/dispatcher.ts index 42290bd..e9e52cd 100644 --- a/app/server/agent/dispatcher.ts +++ b/app/server/agent/dispatcher.ts @@ -3,7 +3,7 @@ import { encode } from 'cborg'; import { WSContext } from 'hono/ws'; import { ChannelType } from './encoder'; -interface Command { +export interface Command { op: string; payload: unknown; } @@ -18,7 +18,14 @@ interface SSHConnectCommand extends Command { }; } -type AgentCommand = SSHConnectCommand; +interface SSHCloseCommand extends Command { + op: 'ssh_term'; + payload: { + sessionId: string; + }; +} + +type AgentCommand = SSHConnectCommand | SSHCloseCommand; export async function dispatchCommand( dispatcher: Writable, @@ -36,21 +43,21 @@ export async function dispatchCommand( }); } -interface SSHConnectData extends Command { +export interface SSHConnectData extends Command { op: 'ssh_conn_successful'; payload: { sessionId: string; }; } -interface SSHConnectFailedData extends Command { +export interface SSHConnectFailedData extends Command { op: 'ssh_conn_failed'; payload: { reason: string; }; } -interface SSHFrameData extends Command { +export interface SSHFrameData extends Command { op: 'ssh_frame'; payload: { channel: ChannelType; diff --git a/app/server/agent/ssh.ts b/app/server/agent/ssh.ts index 3cb1346..e6bdc08 100644 --- a/app/server/agent/ssh.ts +++ b/app/server/agent/ssh.ts @@ -143,22 +143,36 @@ export class SSHMultiplexer { this.sshInput.write(encodedFrame); }, - onClose: (_, ws) => { + onClose: async (_, ws) => { const sessionId = ws.raw; if (sessionId && this.connections.has(sessionId)) { const session = this.connections.get(sessionId); if (session) { + await dispatchCommand(this.control, { + op: 'ssh_term', + payload: { + sessionId, + }, + }); + session.connected = false; this.connections.delete(sessionId); } } }, - onError: (event, ws) => { + onError: async (event, ws) => { const sessionId = ws.raw; if (sessionId && this.connections.has(sessionId)) { const session = this.connections.get(sessionId); if (session) { + await dispatchCommand(this.control, { + op: 'ssh_term', + payload: { + sessionId, + }, + }); + session.connected = false; this.connections.delete(sessionId); } @@ -191,7 +205,7 @@ export class SSHMultiplexer { op: 'ssh_frame', payload: { channel: frame.channel, - data: frame.payload, + frame: frame.payload, }, }); }); diff --git a/package.json b/package.json index d256b0d..52ab534 100644 --- a/package.json +++ b/package.json @@ -28,6 +28,7 @@ "@uiw/codemirror-theme-github": "^4.23.12", "@uiw/codemirror-theme-xcode": "^4.23.12", "@uiw/react-codemirror": "^4.23.12", + "@xterm/xterm": "^5.5.0", "arktype": "^2.1.20", "cborg": "^4.2.11", "clsx": "^2.1.1", diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index dc8f2c0..f7a9b14 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -61,6 +61,9 @@ importers: '@uiw/react-codemirror': specifier: ^4.23.12 version: 4.23.12(@babel/runtime@7.27.3)(@codemirror/autocomplete@6.18.2(@codemirror/language@6.11.0)(@codemirror/state@6.5.2)(@codemirror/view@6.36.8)(@lezer/common@1.2.3))(@codemirror/language@6.11.0)(@codemirror/lint@6.8.2)(@codemirror/search@6.5.7)(@codemirror/state@6.5.2)(@codemirror/theme-one-dark@6.1.2)(@codemirror/view@6.36.8)(codemirror@6.0.1(@lezer/common@1.2.3))(react-dom@19.1.0(react@19.1.0))(react@19.1.0) + '@xterm/xterm': + specifier: ^5.5.0 + version: 5.5.0 arktype: specifier: ^2.1.20 version: 2.1.20 @@ -1591,6 +1594,9 @@ packages: react: '>=16.8.0' react-dom: '>=16.8.0' + '@xterm/xterm@5.5.0': + resolution: {integrity: sha512-hqJHYaQb5OptNunnyAnkHyM8aCjZ1MEIDTQu1iIbbTD/xops91NB5yq1ZK/dC2JDbVWtF23zUtl9JE2NqwT87A==} + acorn@8.14.1: resolution: {integrity: sha512-OvQ/2pUDKmgfCg++xsTX1wGxfTaszcHVcTctW4UJB4hibJx2HXxxO5UmVgyjMa+ZDsiaf5wWLXYpRWMmBI0QHg==} engines: {node: '>=0.4.0'} @@ -4713,6 +4719,8 @@ snapshots: - '@codemirror/lint' - '@codemirror/search' + '@xterm/xterm@5.5.0': {} + acorn@8.14.1: optional: true