VoidShell/services/sway/ipc.ts

import GObject, { register, property, signal } from "astal/gobject";
import {
  WorkspaceResponse,
  OutputResponse,
  SwayMessage,
  SwayEvent,
  WorkspaceEvent,
  OutputEvent,
  TreeNode,
  CommandResponse,
} from "./types";
import Gio from "gi://Gio";
import GLib from "gi://GLib";
import * as utils from "@/utils";
import { Mutex } from "@/synchronization";

@register({
  GTypeName: "Sway",
})
export class Sway extends GObject.Object {
  #commandSocket: Gio.SocketConnection | null = null;
  #subscriptionSocket: Gio.SocketConnection | null = null;

  #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;

  // TODO: figure out how to give these signals a proper type. For now, Object is /fine/.

  @signal(Object)
  /**
   * 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;
  }

  @property(Boolean)
  get isConnected(): boolean {
    return this.#connected;
  }

  #connected = false;

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

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

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

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

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

  static _instance: Sway;

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

  async #createConnection(address: string): Promise<Gio.SocketConnection> {
    console.log(`Connecting to sway socket ${address}`);
    const client = new Gio.SocketClient();
    return new Promise<Gio.SocketConnection>((resolve, reject) => {
      const socket_addr = new Gio.UnixSocketAddress({ path: address });
      client.connect_async(socket_addr, null, (client, result) => {
        try {
          // Type assertion is safe, if it fails we reject the promise in the catch block
          const connection = client!.connect_finish(result);
          resolve(connection);
        } catch (e) {
          console.error(`Failed to connect to socket.`, e);
          reject(e);
        }
      });
    });
  }

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

  async #connect() {
    console.log("Connecting to sway");
    const address = GLib.getenv("SWAYSOCK");
    if (!address) {
      console.error("SWAYSOCK not set");
      throw new Error("SWAYSOCK not set");
    }

    this.#commandSocket = await this.#createConnection(address);
    this.#subscriptionSocket = await this.#createConnection(address);
    console.log("Connected to sway");

    // Start listening for subscriptions
    this.#startSubscriberListen();

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

    this.connected();
    this.#connected = true;
    this.notify("is-connected");
  }

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

  async #startSubscriberListen() {
    if (!this.#subscriptionSocket) {
      console.error("Not connected");
      throw new Error("Not connected");
    }
    await this.#subscribe(
      SwayEvent.OUTPUT,
      SwayEvent.WORKSPACE,
      SwayEvent.MODE,
      SwayEvent.SHUTDOWN,
    );
    const read = async () => {
      while (true) {
        try {
          const [event, payload] = await this.#readRaw(
            this.#subscriptionSocket!,
          );
          switch (event) {
            case SwayEvent.OUTPUT:
              this.output(JSON.parse(payload) as OutputEvent);
              break;

            case SwayEvent.WORKSPACE:
              this.workspace(JSON.parse(payload) as WorkspaceEvent);
              break;
            case SwayEvent.MODE:
              this.#mode = JSON.parse(payload).change;
              this.notify("mode");
              break;
            case SwayEvent.SHUTDOWN:
              this.disconnected();
              break;
            default:
              console.warn("Unhandled event", event);
              break;
          }
        } catch (e) {
          console.error("Failed to read event", e);
        }
      }
    };
    read();
  }

  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;
    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 utils.readFromStream(stream, 8)).buffer);
    const length = header.getUint32(0, true);
    const type = header.getUint32(4, true);
    return { length, type };
  }

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

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

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

  private messageMutex = new Mutex();

  async #sendRaw(
    socket: Gio.SocketConnection | null,
    type: SwayMessage,
    payload: string = "",
    waitForResponse: boolean = true,
  ): Promise<string | null> {
    return this.messageMutex.with(async () => {
      if (!socket || !socket.is_connected()) {
        console.error("Not connected");
        throw new Error("Not connected");
      }
      if (socket === this.#commandSocket) {
        if (type === SwayMessage.SUBSCRIBE) {
          console.error("Cannot subscribe on command socket");
          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 = socket.output_stream;

      // Send the message
      try {
        await utils.writeToStream(outputStream, buffer);
      } catch (e) {
        console.error("Failed to write to stream...", e);
        throw e;
      }
      if (!waitForResponse) {
        return null;
      }

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

  async #subscribe(...events: SwayEvent[]) {
    if (!this.#subscriptionSocket) {
      console.error("Not connected");
      throw new Error("Not connected");
    }
    const eventNames = events.map((e) => SwayEvent[e].toLowerCase());
    const payload = JSON.stringify(eventNames);
    console.log("Subscribing to events:", payload);
    try {
      await this.#sendRaw(
        this.#subscriptionSocket,
        SwayMessage.SUBSCRIBE,
        payload,
      );

      console.log(`Subscribed to events: ${payload}`);
      this.subscribed(events);
    } catch (e) {
      console.error(`Failed to subscribe:`, e);
      throw e;
    }
  }
}