kvm/ui/src/hooks/useKeyboard.ts

import { useCallback, useMemo } from "react";

import {
  hidErrorRollOver,
  hidKeyBufferSize,
  KeysDownState,
  useHidStore,
  useRTCStore,
} from "@/hooks/stores";
import { JsonRpcResponse, useJsonRpc } from "@/hooks/useJsonRpc";
import { useHidRpc } from "@/hooks/useHidRpc";
import { KeyboardLedStateMessage, KeyboardMacro, KeysDownStateMessage } from "@/hooks/hidRpc";
import { hidKeyToModifierMask, keys, modifiers } from "@/keyboardMappings";

export default function useKeyboard() {
  const { send } = useJsonRpc();
  const { rpcDataChannel } = useRTCStore();
  const { keysDownState, setKeysDownState, setKeyboardLedState } = useHidStore();

  // INTRODUCTION: The earlier version of the JetKVM device shipped with all keyboard state
  // being tracked on the browser/client-side. When adding the keyPressReport API to the
  // device-side code, we have to still support the situation where the browser/client-side code
  // is running on the cloud against a device that has not been updated yet and thus does not
  // support the keyPressReport API. In that case, we need to handle the key presses locally
  // and send the full state to the device, so it can behave like a real USB HID keyboard.
  // This flag indicates whether the keyPressReport API is available on the device which is
  // dynamically set when the device responds to the first key press event or reports its
  // keysDownState when queried since the keyPressReport was introduced together with the
  // getKeysDownState API.

  // HidRPC is a binary format for exchanging keyboard and mouse events
  const {
    reportKeyboardEvent: sendKeyboardEventHidRpc,
    reportKeypressEvent: sendKeypressEventHidRpc,
    reportKeyboardMacroEvent: sendKeyboardMacroEventHidRpc,
    rpcHidReady,
  } = useHidRpc(message => {
    switch (message.constructor) {
      case KeysDownStateMessage:
        setKeysDownState((message as KeysDownStateMessage).keysDownState);
        break;
      case KeyboardLedStateMessage:
        setKeyboardLedState((message as KeyboardLedStateMessage).keyboardLedState);
        break;
      default:
        break;
    }
  });

  // sendKeyboardEvent is used to send the full keyboard state to the device for macro handling
  // and resetting keyboard state. It sends the keys currently pressed and the modifier state.
  // The device will respond with the keysDownState if it supports the keyPressReport API
  // or just accept the state if it does not support (returning no result)
  const sendKeyboardEvent = useCallback(
    async (state: KeysDownState) => {
      if (rpcDataChannel?.readyState !== "open" && !rpcHidReady) return;

      console.debug(
        `Send keyboardReport keys: ${state.keys}, modifier: ${state.modifier}`,
      );

      if (rpcHidReady) {
        console.debug("Sending keyboard report via HidRPC");
        sendKeyboardEventHidRpc(state.keys, state.modifier);
        return;
      }

      send(
        "keyboardReport",
        { keys: state.keys, modifier: state.modifier },
        (resp: JsonRpcResponse) => {
          if ("error" in resp) {
            console.error(`Failed to send keyboard report ${state}`, resp.error);
          }
        },
      );
    },
    [rpcDataChannel?.readyState, rpcHidReady, send, sendKeyboardEventHidRpc],
  );

  const MACRO_RESET_KEYBOARD_STATE = useMemo(() => ({
    keys: new Array(hidKeyBufferSize).fill(0),
    modifier: 0,
    delay: 0,
  }), []);

  // resetKeyboardState is used to reset the keyboard state to no keys pressed and no modifiers.
  // This is useful for macros and when the browser loses focus to ensure that the keyboard state
  // is clean.
  const resetKeyboardState = useCallback(async () => {
    // Reset the keys buffer to zeros and the modifier state to zero
    sendKeyboardEvent(MACRO_RESET_KEYBOARD_STATE);
  }, [sendKeyboardEvent, MACRO_RESET_KEYBOARD_STATE]);

  // executeMacro is used to execute a macro consisting of multiple steps.
  // Each step can have multiple keys, multiple modifiers and a delay.
  // The keys and modifiers are pressed together and held for the delay duration.
  // After the delay, the keys and modifiers are released and the next step is executed.
  // If a step has no keys or modifiers, it is treated as a delay-only step.
  // A small pause is added between steps to ensure that the device can process the events.
  const executeMacro = async (
    steps: { keys: string[] | null; modifiers: string[] | null; delay: number }[],
  ) => {
    const macro: KeyboardMacro[] = [];

    for (const [_, step] of steps.entries()) {
      const keyValues = (step.keys || []).map(key => keys[key]).filter(Boolean);
      const modifierMask: number = (step.modifiers || [])
        .map(mod => modifiers[mod])
        .reduce((acc, val) => acc + val, 0);

      // If the step has keys and/or modifiers, press them and hold for the delay
      if (keyValues.length > 0 || modifierMask > 0) {
        macro.push({ keys: keyValues, modifier: modifierMask, delay: 50 });
        macro.push({ ...MACRO_RESET_KEYBOARD_STATE, delay: 200 });
      }
    }

    sendKeyboardMacroEventHidRpc(macro);
  };

  const cancelExecuteMacro = useCallback(async () => {
    send("cancelKeyboardReportMulti", {}, (resp: JsonRpcResponse) => {
      if ("error" in resp) {
        console.error(`Failed to cancel keyboard report multi`, resp.error);
      }
    });
  }, [send]);

  // handleKeyPress is used to handle a key press or release event.
  // This function handle both key press and key release events.
  // It checks if the keyPressReport API is available and sends the key press event.
  // If the keyPressReport API is not available, it simulates the device-side key
  // handling for legacy devices and updates the keysDownState accordingly.
  // It then sends the full keyboard state to the device.
  const handleKeyPress = useCallback(
    async (key: number, press: boolean) => {
      if (rpcDataChannel?.readyState !== "open" && !rpcHidReady) return;
      if ((key || 0) === 0) return; // ignore zero key presses (they are bad mappings)

      if (rpcHidReady) {
        // if the keyPress api is available, we can just send the key press event
        // sendKeypressEvent is used to send a single key press/release event to the device.
        // It sends the key and whether it is pressed or released.
        // Older device version doesn't support this API, so we will switch to local key handling
        // In that case we will switch to local key handling and update the keysDownState
        // in client/browser-side code using simulateDeviceSideKeyHandlingForLegacyDevices.
        sendKeypressEventHidRpc(key, press);
      } else {
        // if the keyPress api is not available, we need to handle the key locally
        const downState = simulateDeviceSideKeyHandlingForLegacyDevices(
          keysDownState,
          key,
          press,
        );
        sendKeyboardEvent(downState); // then we send the full state

        // if we just sent ErrorRollOver, reset to empty state
        if (downState.keys[0] === hidErrorRollOver) {
          resetKeyboardState();
        }
      }
    },
    [
      rpcHidReady,
      keysDownState,
      resetKeyboardState,
      rpcDataChannel?.readyState,
      sendKeyboardEvent,
      sendKeypressEventHidRpc,
    ],
  );

  // IMPORTANT: See the keyPressReportApiAvailable comment above for the reason this exists
  function simulateDeviceSideKeyHandlingForLegacyDevices(
    state: KeysDownState,
    key: number,
    press: boolean,
  ): KeysDownState {
    // IMPORTANT: This code parallels the logic in the kernel's hid-gadget driver
    // for handling key presses and releases. It ensures that the USB gadget
    // behaves similarly to a real USB HID keyboard. This logic is paralleled
    // in the device-side code in hid_keyboard.go so make sure to keep them in sync.
    let modifiers = state.modifier;
    const keys = state.keys;
    const modifierMask = hidKeyToModifierMask[key] || 0;

    if (modifierMask !== 0) {
      // If the key is a modifier key, we update the keyboardModifier state
      // by setting or clearing the corresponding bit in the modifier byte.
      // This allows us to track the state of dynamic modifier keys like
      // Shift, Control, Alt, and Super.
      if (press) {
        modifiers |= modifierMask;
      } else {
        modifiers &= ~modifierMask;
      }
    } else {
      // handle other keys that are not modifier keys by placing or removing them
      // from the key buffer since the buffer tracks currently pressed keys
      let overrun = true;
      for (let i = 0; i < hidKeyBufferSize; i++) {
        // If we find the key in the buffer the buffer, we either remove it (if press is false)
        // or do nothing (if down is true) because the buffer tracks currently pressed keys
        // and if we find a zero byte, we can place the key there (if press is true)
        if (keys[i] === key || keys[i] === 0) {
          if (press) {
            keys[i] = key; // overwrites the zero byte or the same key if already pressed
          } else {
            // we are releasing the key, remove it from the buffer
            if (keys[i] !== 0) {
              keys.splice(i, 1);
              keys.push(0); // add a zero at the end
            }
          }
          overrun = false; // We found a slot for the key
          break;
        }
      }

      // If we reach here it means we didn't find an empty slot or the key in the buffer
      if (overrun) {
        if (press) {
          console.warn(
            `keyboard buffer overflow current keys ${keys}, key: ${key} not added`,
          );
          // Fill all key slots with ErrorRollOver (0x01) to indicate overflow
          keys.length = hidKeyBufferSize;
          keys.fill(hidErrorRollOver);
        } else {
          // If we are releasing a key, and we didn't find it in a slot, who cares?
          console.debug(`key ${key} not found in buffer, nothing to release`);
        }
      }
    }
    return { modifier: modifiers, keys };
  }

  return { handleKeyPress, resetKeyboardState, executeMacro, cancelExecuteMacro };
}