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.
 *
 * IMPORTANT: Guacamole.Tunnel assigns connect/sendMessage/disconnect as
 * instance properties in its constructor (not prototype methods). We MUST
 * reassign them in our constructor AFTER super() to override the no-ops.
 */
class JsonWsTunnel extends Guacamole.Tunnel {
  private ws: WebSocket | null = null
  private readonly wsUrl: string
  private readonly connectMsg: object

  // 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

    // Override instance methods set by Guacamole.Tunnel constructor
    this.connect = this._connect.bind(this)
    this.disconnect = this._disconnect.bind(this)
    this.sendMessage = this._sendMessage.bind(this)
  }

  private _connect(_data?: string) {
    console.log('[RDP] Tunnel opening WebSocket:', this.wsUrl)
    this.ws = new WebSocket(this.wsUrl)

    this.ws.onopen = () => {
      console.log('[RDP] WebSocket open, sending 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':
          this.onConnected?.(msg.hostId, msg.hostName)
          break

        case 'guac': {
          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'))
    }
  }

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

  private _sendMessage(...elements: any[]) {
    // Guacamole.Client calls sendMessage(opcode, arg1, arg2, ...) with individual elements.
    // Encode into Guacamole wire format: "<len>.<value>,<len>.<value>,...;"
    if (this.ws?.readyState === WebSocket.OPEN) {
      const instruction = elements.map((e: any) => {
        const s = String(e)
        return `${s.length}.${s}`
      }).join(',') + ';'
      this.ws.send(JSON.stringify({ type: 'guac', instruction }))
    }
  }

  /**
   * Parse one or more Guacamole instructions from a raw buffer string
   * and call oninstruction for each.
   */
  private _dispatchInstructions(raw: string) {
    if (!this.oninstruction) return

    let remaining = raw
    while (remaining.length > 0) {
      const semicolonIdx = remaining.indexOf(';')
      if (semicolonIdx === -1) break

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

      if (!instruction) continue

      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
      }

      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,
    pendingSessionId: string,
    options?: { width?: number; height?: number },
  ) {
    const proto = location.protocol === 'https:' ? 'wss' : 'ws'
    const wsUrl = `${proto}://${location.host}/api/ws/rdp?token=${auth.token}`

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

    console.log(`[RDP] Connecting to ${wsUrl} for hostId=${hostId} (${width}x${height})`)

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

    const client = new Guacamole.Client(tunnel)

    const sessionId = pendingSessionId

    tunnel.onConnected = (_resolvedHostId: number, resolvedHostName: string) => {
      console.log(`[RDP] Connected to ${resolvedHostName}`)
      sessions.replaceSession(sessionId, {
        key: sessionId,
        id: sessionId,
        hostId,
        hostName: resolvedHostName || hostName,
        protocol: 'rdp',
        color,
        active: true,
      })
    }

    tunnel.onDisconnected = (reason: string) => {
      console.log(`[RDP] Disconnected: ${reason}`)
      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
    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
    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 }
}