VoidShell/widget/services/sway/ipc.ts

import GObject, { register, property, signal } from "astal/gobject";
import { App } from "astal";
import {
  WorkspaceResponse,
  OutputResponse,
  SwayMessage,
  SwayEvent,
  ModeEvent,
  WorkspaceEvent,
  OutputEvent,
} from "./types";
import Gdk from "gi://Gdk";
import Gio from "gi://Gio";
import GLib from "gi://GLib";
import * as utils from "../../utils";

@register({
  GTypeName: "Sway",
})
/**
 * Sway connection singleton that implements the sway IPC protocol at a low level.
 */
export class Sway extends GObject.Object {
  private #commandSocket: Gio.SocketConnection | null = null;
  private #subscriptionSocket: Gio.SocketConnection | null = null;

  private #mode: string = "";

  @signal()
  /**
   * Emitted when the connection to sway is established
   */
  declare connected: () => void;

  @signal()
  /**
   * Emitted when the connection to sway is lost
   */
  declare disconnected: () => void;

  @signal(Array)
  /**
   * Emitted when we receive a subscription event
   */
  declare subscribed: (events: SwayEvent[]) => void;

  @signal(Object)
  /**
   * Emitted when we receive a workspace event
   */
  declare workspace: (event: WorkspaceEvent) => void;

  @signal(Object)
  /**
   * Emitted when we receive an output event
   */
  declare output: (event: OutputEvent) => void;

  @property(String)
  /**
   * The current binding mode of sway
   */
  get mode(): string {
    return this.#mode;
  }

  /**
   * Run a command on sway
   */
  public async command(command: string): Promise<CommandResponse[]> {
    const result = JSON.parse(
      await this.#sendRaw(this.#commandSocket, SwayMessage.RUN_COMMAND, command)
    );
  }

  /**
   * Get the current sway workspaces
   */
  public async getWorkspaces(): Promise<WorkspaceResponse[]> {
    const result = JSON.parse(
      await this.#sendRaw(this.#commandSocket, SwayMessage.GET_WORKSPACES)
    );
  }

  /**
   * Get the current sway outputs
   */
  public async getOutputs(): Promise<OutputResponse[]> {
    const result = JSON.parse(
      await this.#sendRaw(this.#commandSocket, SwayMessage.GET_OUTPUTS)
    );
  }

  /**
   * Get the current sway tree
   */
  public async getTree(): Promise<TreeNode> {
    const result = JSON.parse(
      await this.#sendRaw(this.#commandSocket, SwayMessage.GET_TREE)
    );
  }

  private constructor() {
    super();
    this.#connect();
  }

  private static #instance: Sway;

  /**
   * Get the default instance of Sway
   */
  static get_default() {
    if (!Sway.#instance) {
      Sway.#instance = new Sway();
    }
    return Sway.#instance;
  }

  private async #createConnection(
    address: string
  ): Promise<Gio.SocketConnection> {
    const client = new Gio.SocketClient();
    return new Promise<Gio.SocketConnection>((resolve, reject) => {
      client.connect_to_uri_async(
        `unix:${address}`,
        0,
        null,
        (client, result) => {
          try {
            const connection = client.connect_to_uri_finish(result);
            resolve(connection);
          } catch (e) {
            reject(e);
          }
        }
      );
    });
  }

  private async #getMode(): Promise<string> {
    const result = JSON.parse(
      await this.#sendRaw(this.#commandSocket, SwayMessage.GET_BINDING_STATE)
    );
    return result.name;
  }

  private async #connect() {
    const address = GLib.getenv("SWAYSOCK");
    if (!address) {
      throw new Error("SWAYSOCK not set");
    }
    const socket = new Gio.SocketClient();
    // Create two connections, one for commands and one for subscriptions. This is so we
    // can easily distinguish between responses to commands and subscription events.
    const command = await this.#createConnection(address);
    const subscription = await this.#createConnection(address);

    this.#commandSocket = command;
    this.#subscriptionSocket = subscription;
    this.connected();

    // Get the current mode
    this.#mode = await this.#getMode();
    this.notify("mode");

    this.#startSubscriberListen();
  }

  static readonly MAGIC = "i3-ipc";
  static readonly HEADER_SIZE = Sway.MAGIC.length + 8;
  static readonly CHUNK_SIZE = 4096;

  private #startSubscriberListen() {
    if (!this.#subscriptionSocket) {
      throw new Error("Not connected");
    }
    this.#subscribe(
      SwayEvent.OUTPUT,
      SwayEvent.WORKSPACE,
      SwayEvent.MODE,
      SwayEvent.SHUTDOWN
    );
    const inputStream = this.#subscriptionSocket.input_stream;
    const reader = new Gio.DataInputStream({ base_stream: inputStream });
    const read = async () => {
      while (true) {
        try {
          const [event, payload] = await this.#readRaw(
            this.#subscriptionSocket
          );
          switch (event) {
            case SwayEvent.OUTPUT:
              const event = JSON.parse(payload) as OutputResponse;
              this.output(event);
              break;
            case SwayEvent.WORKSPACE:
              const event = JSON.parse(payload) as WorkspaceResponse;
              this.workspace(event);
              break;
            case SwayEvent.MODE:
              const event = JSON.parse(payload) as ModeResponse;
              this.mode(event);
              break;
            case SwayEvent.SHUTDOWN:
              this.disconnected();
              break;
            default:
              console.warn("Unhandled event", event);
              break;
          }
        } catch (e) {
          console.error("Failed to read event", e);
          break;
        }
      }
    };
    read();
  }

  private async #readHeader(
    stream: Gio.InputStream
  ): Promise<{ length: number; type: SwayMessage | SwayEvent }> {
    // We read one byte at a time until we have constructed a full magic string, then we read the rest of the header
    // This allows us to handle partial reads and corrupted data gracefully
    let idx = 0;
    const magic = new Uint8Array(Sway.MAGIC.length);
    while (idx < Sway.MAGIC.length) {
      const buffer = await utils.readFromStreamRaw(stream, 1);
      const byte = buffer[0];
      if (byte !== Sway.MAGIC.charCodeAt(idx)) {
        // Reset if we don't match the magic string
        idx = 0;
      } else {
        // Otherwise, keep going
        idx++;
      }
    }
    const header = new DataView(await readFromStream(stream, 8).buffer);
    const length = header.getUint32(0, true);
    const type = header.getUint32(4, true);
    return { length, type };
  }

  private async #readRaw(
    socket: Gio.SocketConnection
  ): Promise<[SwayMessage | SwayEvent, string]> {
    const inputStream = socket.input_stream;

    const { length, type } = await this.#readHeader(inputStream);
    const payloadBuf = utils.readFromStream(inputStream, length);
    const payload = new TextDecoder().decode(payloadBuf);

    return [type as SwayMessage | SwayEvent, payload];
  }

  private async #sendRaw(
    socket: Gio.SocketConnection | null,
    type: SwayMessage,
    payload: string = ""
  ): Promise<string> {
    if (!socket || !socket.is_connected()) {
      throw new Error("Not connected");
    }
    if (socket === this.#commandSocket && type === SwayMessage.SUBSCRIBE) {
      throw new Error("Cannot subscribe on command socket");
    }
    // Construct the message
    const buffer = new ArrayBuffer(Sway.HEADER_SIZE + payload.length);
    const magicView = new Uint8Array(buffer, 0, Sway.MAGIC.length);
    const lengthView = new DataView(buffer, Sway.MAGIC.length, 4);
    const typeView = new DataView(buffer, Sway.MAGIC.length + 4, 4);
    const payloadView = new Uint8Array(buffer, Sway.HEADER_SIZE);

    magicView.set(Sway.MAGIC.split("").map((c) => c.charCodeAt(0)));
    lengthView.setUint32(0, payload.length, true);
    typeView.setUint32(0, type, true);
    payloadView.set(payload.split("").map((c) => c.charCodeAt(0)));

    const outputStream = this.#commandSocket.output_stream;

    // Send the message
    await utils.writeToStream(outputStream, buffer);

    // Read the response
    const [resultType, result] = await this.#readRaw(socket);
    if (resultType !== type) {
      throw new Error(`Unexpected response type: ${resultType}`);
    }
    return result;
  }

  private async #subscribe(...events: SwayEvent[]) {
    if (!this.#subscriptionSocket) {
      throw new Error("Not connected");
    }
    const eventNames = events.map((e) => SwayEvent[e].toLowerCase());
    const payload = JSON.stringify(eventNames);
    const result = JSON.parse(
      await this.#sendRaw(
        this.#subscriptionSocket,
        SwayMessage.SUBSCRIBE,
        payload
      )
    );
    if (result.success) {
      this.subscribed(events);
    } else {
      throw new Error(`Failed to subscribe: ${result}`);
    }
  }
}