wraith/frontend/composables/useRdp.ts

import Guacamole from 'guacamole-common-js'
import { useAuthStore } from '~/stores/auth.store'
import { useSessionStore } from '~/stores/session.store'

/**
 * Custom Guacamole Tunnel that speaks JSON over WebSocket.
 *
 * Why custom instead of Guacamole.WebSocketTunnel:
 *   - The standard WebSocketTunnel expects raw Guacamole protocol over WS.
 *   - Our /ws/rdp gateway wraps all messages in JSON envelopes (same pattern
 *     as /ws/terminal and /ws/sftp) for auth, connect handshake, and error
 *     signaling.
 *   - This tunnel bridges the two: guacamole-common-js sees raw Guacamole
 *     instructions while our gateway sees consistent JSON messages.
 *
 * Protocol:
 *   Browser → Gateway: { type: 'connect', hostId, width, height, ... }
 *   Browser → Gateway: { type: 'guac', instruction: '<raw guac instruction>' }
 *   Gateway → Browser: { type: 'connected', hostId, hostName }
 *   Gateway → Browser: { type: 'guac', instruction: '<raw guac instruction>' }
 *   Gateway → Browser: { type: 'error', message: '...' }
 *   Gateway → Browser: { type: 'disconnected', reason: '...' }
 */
class JsonWsTunnel extends Guacamole.Tunnel {
  private ws: WebSocket | null = null
  private readonly wsUrl: string
  private readonly connectMsg: object

  // Callbacks set by Guacamole.Client after we're constructed
  oninstruction: ((opcode: string, args: string[]) => void) | null = null
  onerror: ((status: Guacamole.Status) => void) | null = null

  // Expose for external callers (connected event, disconnect action)
  onConnected: ((hostId: number, hostName: string) => void) | null = null
  onDisconnected: ((reason: string) => void) | null = null
  onGatewayError: ((message: string) => void) | null = null

  constructor(wsUrl: string, connectMsg: object) {
    super()
    this.wsUrl = wsUrl
    this.connectMsg = connectMsg
  }

  connect(_data?: string) {
    this.ws = new WebSocket(this.wsUrl)

    this.ws.onopen = () => {
      // Step 1: send our JSON connect handshake
      this.ws!.send(JSON.stringify(this.connectMsg))
    }

    this.ws.onmessage = (event: MessageEvent) => {
      const msg = JSON.parse(event.data as string)

      switch (msg.type) {
        case 'connected':
          // Gateway completed the guacd SELECT→CONNECT handshake.
          // Notify RdpCanvas so it can finalize the session store entry.
          this.onConnected?.(msg.hostId, msg.hostName)
          break

        case 'guac': {
          // Raw Guacamole instruction from guacd — parse and dispatch to client.
          // Instructions arrive as: "4.size,4.1024,3.768;" or multiple batched.
          // We dispatch each complete instruction individually.
          const instruction: string = msg.instruction
          this._dispatchInstructions(instruction)
          break
        }

        case 'error':
          this.onGatewayError?.(msg.message as string)
          this.onerror?.(new Guacamole.Status(Guacamole.Status.Code.SERVER_ERROR, msg.message))
          break

        case 'disconnected':
          this.onDisconnected?.(msg.reason as string)
          break
      }
    }

    this.ws.onclose = (event: CloseEvent) => {
      if (event.code === 4001) {
        this.onerror?.(new Guacamole.Status(Guacamole.Status.Code.CLIENT_FORBIDDEN, 'Unauthorized'))
      }
      this.onDisconnected?.('WebSocket closed')
    }

    this.ws.onerror = () => {
      this.onerror?.(new Guacamole.Status(Guacamole.Status.Code.SERVER_ERROR, 'WebSocket error'))
    }
  }

  disconnect() {
    if (this.ws && this.ws.readyState === WebSocket.OPEN) {
      this.ws.close()
    }
    this.ws = null
  }

  sendMessage(message: string) {
    // Guacamole.Client calls this with raw Guacamole instruction strings.
    // Wrap in JSON envelope and forward to gateway → guacd.
    if (this.ws?.readyState === WebSocket.OPEN) {
      this.ws.send(JSON.stringify({ type: 'guac', instruction: message }))
    }
  }

  /**
   * Parse one or more Guacamole instructions from a raw buffer string
   * and call oninstruction for each.
   *
   * Guacamole instruction format: "<len>.<opcode>,<len>.<arg1>,...;"
   * Multiple instructions may arrive in one data chunk.
   */
  private _dispatchInstructions(raw: string) {
    if (!this.oninstruction) return

    let remaining = raw
    while (remaining.length > 0) {
      const semicolonIdx = remaining.indexOf(';')
      if (semicolonIdx === -1) break // incomplete instruction — guacd shouldn't send partials but be safe

      const instruction = remaining.substring(0, semicolonIdx)
      remaining = remaining.substring(semicolonIdx + 1)

      if (!instruction) continue

      // Parse length-prefixed fields
      const parts: string[] = []
      let pos = 0
      while (pos < instruction.length) {
        const dotIdx = instruction.indexOf('.', pos)
        if (dotIdx === -1) break
        const len = parseInt(instruction.substring(pos, dotIdx), 10)
        if (isNaN(len)) break
        const value = instruction.substring(dotIdx + 1, dotIdx + 1 + len)
        parts.push(value)
        pos = dotIdx + 1 + len + 1 // skip comma separator
      }

      if (parts.length > 0) {
        const opcode = parts[0]
        const args = parts.slice(1)
        this.oninstruction(opcode, args)
      }
    }
  }
}

// ─── Composable ────────────────────────────────────────────────────────────────

export function useRdp() {
  const auth = useAuthStore()
  const sessions = useSessionStore()

  function connectRdp(
    container: HTMLElement,
    hostId: number,
    hostName: string,
    color: string | null,
    options?: { width?: number; height?: number },
  ) {
    const proto = location.protocol === 'https:' ? 'wss' : 'ws'
    const wsUrl = `${proto}://${location.host}/ws/rdp?token=${auth.token}`

    const width = options?.width || container.clientWidth || 1920
    const height = options?.height || container.clientHeight || 1080

    const tunnel = new JsonWsTunnel(wsUrl, {
      type: 'connect',
      hostId,
      width,
      height,
    })

    const client = new Guacamole.Client(tunnel)

    // Session store ID — created optimistically, will be confirmed on 'connected'
    const sessionId = `rdp-${hostId}-${Date.now()}`

    // Wire tunnel callbacks
    tunnel.onConnected = (_resolvedHostId: number, resolvedHostName: string) => {
      sessions.addSession({
        key: sessionId,
        id: sessionId,
        hostId,
        hostName: resolvedHostName || hostName,
        protocol: 'rdp',
        color,
        active: true,
      })
    }

    tunnel.onDisconnected = (_reason: string) => {
      sessions.removeSession(sessionId)
      client.disconnect()
    }

    tunnel.onGatewayError = (message: string) => {
      console.error('[RDP] Gateway error:', message)
    }

    // Attach Guacamole display element to container
    const displayEl = client.getDisplay().getElement()
    container.appendChild(displayEl)

    // Mouse input
    const mouse = new Guacamole.Mouse(displayEl)
    mouse.onEach(
      ['mousedown', 'mousemove', 'mouseup'],
      (e: Guacamole.Mouse.Event) => {
        client.sendMouseState(e.state)
      },
    )

    // Keyboard input — attach to document so key events fire even when
    // focus is on other elements (toolbar buttons, etc.)
    const keyboard = new Guacamole.Keyboard(document)
    keyboard.onkeydown = (keysym: number) => client.sendKeyEvent(1, keysym)
    keyboard.onkeyup = (keysym: number) => client.sendKeyEvent(0, keysym)

    // Initiate the WebSocket connection (triggers connect handshake)
    client.connect()

    function disconnect() {
      keyboard.onkeydown = null
      keyboard.onkeyup = null
      client.disconnect()
      sessions.removeSession(sessionId)
    }

    function sendClipboardText(text: string) {
      const stream = client.createClipboardStream('text/plain')
      const writer = new Guacamole.StringWriter(stream)
      writer.sendText(text)
      writer.sendEnd()
    }

    return { client, sessionId, disconnect, sendClipboardText }
  }

  return { connectRdp }
}