keryx 0.0.1

package/initializers/channels.ts

@@ -0,0 +1,284 @@
+import path, { join } from "path";
+import { api, logger } from "../api";
+import type { Channel } from "../classes/Channel";
+import type { Connection } from "../classes/Connection";
+import { Initializer } from "../classes/Initializer";
+import { ErrorType, TypedError } from "../classes/TypedError";
+import { config } from "../config";
+import { globLoader } from "../util/glob";
+
+const namespace = "channels";
+const PRESENCE_KEY_PREFIX = "presence:";
+const LUA_DIR = join(import.meta.dir, "..", "lua");
+
+declare module "../classes/API" {
+  export interface API {
+    [namespace]: Awaited<ReturnType<Channels["initialize"]>>;
+  }
+}
+
+export class Channels extends Initializer {
+  private addPresenceLua = "";
+  private removePresenceLua = "";
+  private refreshPresenceLua = "";
+  private heartbeatTimer: ReturnType<typeof setInterval> | null = null;
+
+  constructor() {
+    super(namespace);
+    this.loadPriority = 100;
+    this.startPriority = 600;
+    this.stopPriority = 50;
+  }
+
+  /**
+   * Find a channel definition that matches the given channel name.
+   * Returns undefined if no matching channel is found.
+   */
+  findChannel = (channelName: string): Channel | undefined => {
+    return api.channels.channels.find((c) => c.matches(channelName));
+  };
+
+  /**
+   * Authorize a connection to subscribe to a channel.
+   * Runs all middleware and the channel's authorize method.
+   * Throws TypedError if authorization fails.
+   */
+  authorizeSubscription = async (
+    channelName: string,
+    connection: Connection,
+  ): Promise<void> => {
+    const channel = this.findChannel(channelName);
+
+    if (!channel) {
+      throw new TypedError({
+        message: `Channel not found: ${channelName}`,
+        type: ErrorType.CONNECTION_CHANNEL_AUTHORIZATION,
+      });
+    }
+
+    // Run all middleware runBefore hooks
+    for (const middleware of channel.middleware) {
+      if (middleware.runBefore) {
+        await middleware.runBefore(channelName, connection);
+      }
+    }
+
+    // Run the channel's authorize method
+    await channel.authorize(channelName, connection);
+  };
+
+  /**
+   * Called when a connection unsubscribes from a channel.
+   * Runs all middleware runAfter hooks.
+   */
+  handleUnsubscription = async (
+    channelName: string,
+    connection: Connection,
+  ): Promise<void> => {
+    const channel = this.findChannel(channelName);
+
+    if (!channel) {
+      return;
+    }
+
+    // Run all middleware runAfter hooks
+    for (const middleware of channel.middleware) {
+      if (middleware.runAfter) {
+        await middleware.runAfter(channelName, connection);
+      }
+    }
+  };
+
+  /**
+   * Record a connection's presence in a channel and broadcast a join event
+   * if this is the first connection for that presence key.
+   *
+   * Redis keys:
+   *   presence:{channelName} — Set of presence keys
+   *   presence:{channelName}:{presenceKey} — Set of connection IDs
+   */
+  addPresence = async (
+    channelName: string,
+    connection: Connection,
+  ): Promise<void> => {
+    const channel = this.findChannel(channelName);
+    const key = channel ? await channel.presenceKey(connection) : connection.id;
+
+    const channelKey = `${PRESENCE_KEY_PREFIX}${channelName}`;
+    const connectionSetKey = `${PRESENCE_KEY_PREFIX}${channelName}:${key}`;
+
+    const added = await api.redis.redis.eval(
+      this.addPresenceLua,
+      2,
+      connectionSetKey,
+      channelKey,
+      connection.id,
+      key,
+      config.channels.presenceTTL,
+    );
+
+    if (added === 1) {
+      await api.pubsub.broadcast(
+        channelName,
+        JSON.stringify({ event: "join", presenceKey: key }),
+        "presence",
+      );
+    }
+  };
+
+  /**
+   * Remove a connection's presence from a channel and broadcast a leave event
+   * if this was the last connection for that presence key.
+   */
+  removePresence = async (
+    channelName: string,
+    connection: Connection,
+  ): Promise<void> => {
+    const channel = this.findChannel(channelName);
+    const key = channel ? await channel.presenceKey(connection) : connection.id;
+
+    const channelKey = `${PRESENCE_KEY_PREFIX}${channelName}`;
+    const connectionSetKey = `${PRESENCE_KEY_PREFIX}${channelName}:${key}`;
+
+    const shouldLeave = await api.redis.redis.eval(
+      this.removePresenceLua,
+      2,
+      connectionSetKey,
+      channelKey,
+      connection.id,
+      key,
+    );
+
+    if (shouldLeave === 1) {
+      await api.pubsub.broadcast(
+        channelName,
+        JSON.stringify({ event: "leave", presenceKey: key }),
+        "presence",
+      );
+    }
+  };
+
+  /**
+   * Returns the list of presence keys for members currently in the channel
+   * across all server instances.
+   */
+  members = async (channelName: string): Promise<string[]> => {
+    const channelKey = `${PRESENCE_KEY_PREFIX}${channelName}`;
+    return api.redis.redis.smembers(channelKey);
+  };
+
+  /**
+   * Refresh TTLs on all presence keys owned by local connections.
+   * Called periodically by the heartbeat timer to prevent keys from
+   * expiring while the server is still alive.
+   */
+  refreshPresence = async (): Promise<void> => {
+    const keysToRefresh = new Set<string>();
+
+    for (const connection of api.connections.connections) {
+      for (const channelName of connection.subscriptions) {
+        const channel = this.findChannel(channelName);
+        const key = channel
+          ? await channel.presenceKey(connection)
+          : connection.id;
+
+        keysToRefresh.add(`${PRESENCE_KEY_PREFIX}${channelName}`);
+        keysToRefresh.add(`${PRESENCE_KEY_PREFIX}${channelName}:${key}`);
+      }
+    }
+
+    if (keysToRefresh.size === 0) return;
+
+    const keys = [...keysToRefresh];
+    await api.redis.redis.eval(
+      this.refreshPresenceLua,
+      keys.length,
+      ...keys,
+      config.channels.presenceTTL,
+    );
+  };
+
+  /**
+   * Clear all presence data. Useful for test cleanup.
+   */
+  clearPresence = async (): Promise<void> => {
+    let cursor = "0";
+    do {
+      const [nextCursor, keys] = await api.redis.redis.scan(
+        cursor,
+        "MATCH",
+        `${PRESENCE_KEY_PREFIX}*`,
+        "COUNT",
+        100,
+      );
+      cursor = nextCursor;
+      if (keys.length > 0) {
+        await api.redis.redis.del(...keys);
+      }
+    } while (cursor !== "0");
+  };
+
+  async initialize() {
+    this.addPresenceLua = await Bun.file(
+      join(LUA_DIR, "add-presence.lua"),
+    ).text();
+    this.removePresenceLua = await Bun.file(
+      join(LUA_DIR, "remove-presence.lua"),
+    ).text();
+    this.refreshPresenceLua = await Bun.file(
+      join(LUA_DIR, "refresh-presence.lua"),
+    ).text();
+
+    let channels: Channel[] = [];
+
+    // Channels are always user-defined, load from rootDir only
+    try {
+      channels = await globLoader<Channel>(path.join(api.rootDir, "channels"));
+    } catch (e) {
+      // channels directory may not exist, which is fine
+      logger.debug(
+        `No channels directory found or error loading channels: ${e}`,
+      );
+    }
+
+    for (const c of channels) {
+      if (!c.description) c.description = `A Channel: ${c.name}`;
+    }
+
+    logger.info(`loaded ${channels.length} channels`);
+
+    return {
+      channels,
+      findChannel: this.findChannel,
+      authorizeSubscription: this.authorizeSubscription,
+      handleUnsubscription: this.handleUnsubscription,
+      addPresence: this.addPresence,
+      removePresence: this.removePresence,
+      refreshPresence: this.refreshPresence,
+      members: this.members,
+      clearPresence: this.clearPresence,
+    };
+  }
+
+  async start() {
+    const intervalMs = config.channels.presenceHeartbeatInterval * 1000;
+    this.heartbeatTimer = setInterval(async () => {
+      try {
+        await this.refreshPresence();
+      } catch (e) {
+        logger.error(`presence heartbeat error: ${e}`);
+      }
+    }, intervalMs);
+
+    logger.info(
+      `presence heartbeat started (interval=${config.channels.presenceHeartbeatInterval}s, ttl=${config.channels.presenceTTL}s)`,
+    );
+  }
+
+  async stop() {
+    if (this.heartbeatTimer) {
+      clearInterval(this.heartbeatTimer);
+      this.heartbeatTimer = null;
+    }
+  }
+}

package/initializers/connections.ts

@@ -0,0 +1,37 @@
+import { Connection, api } from "../api";
+import { Initializer } from "../classes/Initializer";
+
+const namespace = "connections";
+
+declare module "../classes/API" {
+  export interface API {
+    [namespace]: Awaited<ReturnType<Connections["initialize"]>>;
+  }
+}
+
+export class Connections extends Initializer {
+  constructor() {
+    super(namespace);
+    this.loadPriority = 1;
+  }
+
+  async initialize() {
+    function find(type: string, identifier: string, id: string) {
+      const index = api.connections.connections.findIndex(
+        (c) => c.type === type && c.id === id && c.identifier === identifier,
+      );
+
+      return { connection: api.connections.connections[index], index };
+    }
+
+    function destroy(type: string, identifier: string, id: string) {
+      const { connection, index } = find(type, identifier, id);
+      if (connection) {
+        return api.connections.connections.splice(index, 1);
+      }
+      return [];
+    }
+
+    return { connections: [] as Connection[], find, destroy };
+  }
+}

package/initializers/db.ts

@@ -0,0 +1,158 @@
+import { $ } from "bun";
+import { type Config as DrizzleMigrateConfig } from "drizzle-kit";
+import { DefaultLogger, type LogWriter, sql } from "drizzle-orm";
+import { drizzle } from "drizzle-orm/node-postgres";
+import { migrate } from "drizzle-orm/node-postgres/migrator";
+import { unlink } from "node:fs/promises";
+import path from "path";
+import { Pool } from "pg";
+import { api, logger } from "../api";
+import { Initializer } from "../classes/Initializer";
+import { ErrorType, TypedError } from "../classes/TypedError";
+import { config } from "../config";
+import { formatConnectionStringForLogging } from "../util/connectionString";
+
+const namespace = "db";
+
+declare module "../classes/API" {
+  export interface API {
+    [namespace]: Awaited<ReturnType<DB["initialize"]>>;
+  }
+}
+
+export class DB extends Initializer {
+  constructor() {
+    super(namespace);
+    this.loadPriority = 100;
+    this.startPriority = 100;
+    this.stopPriority = 910;
+  }
+
+  async initialize() {
+    const dbContainer = {} as {
+      db: ReturnType<typeof drizzle>;
+      pool: Pool;
+    };
+    return Object.assign(
+      {
+        generateMigrations: this.generateMigrations,
+        clearDatabase: this.clearDatabase,
+      },
+      dbContainer,
+    );
+  }
+
+  async start() {
+    api.db.pool = new Pool({
+      connectionString: config.database.connectionString,
+    });
+
+    class DrizzleLogger implements LogWriter {
+      write(message: string) {
+        logger.debug(message);
+      }
+    }
+
+    api.db.db = drizzle(api.db.pool, {
+      logger: new DefaultLogger({ writer: new DrizzleLogger() }),
+    });
+
+    try {
+      await api.db.db.execute(sql`SELECT NOW()`);
+    } catch (e) {
+      throw new TypedError({
+        type: ErrorType.SERVER_INITIALIZATION,
+        message: `Cannot connect to database (${formatConnectionStringForLogging(config.database.connectionString)}): ${e}`,
+      });
+    }
+
+    if (config.database.autoMigrate) {
+      try {
+        await migrate(api.db.db, {
+          migrationsFolder: path.join(api.rootDir, "drizzle"),
+        });
+        logger.info("database migrated successfully");
+      } catch (e) {
+        throw new TypedError({
+          type: ErrorType.SERVER_INITIALIZATION,
+          message: `Cannot migrate database (${formatConnectionStringForLogging(config.database.connectionString)}): ${e}`,
+        });
+      }
+    }
+
+    logger.info(
+      `database connection established (${formatConnectionStringForLogging(config.database.connectionString)})`,
+    );
+  }
+
+  async stop() {
+    if (api.db.db && api.db.pool) {
+      try {
+        await api.db.pool.end();
+        logger.info("database connection closed");
+      } catch (e) {
+        logger.error("error closing database connection", e);
+      }
+    }
+  }
+
+  /**
+   * Generate migrations for the database schema.
+   * Learn more @ https://orm.drizzle.team/kit-docs/overview
+   */
+  async generateMigrations() {
+    const migrationConfig = {
+      schema: path.join("models", "*"),
+      dbCredentials: {
+        uri: config.database.connectionString,
+      },
+      out: path.join("drizzle"),
+    } satisfies DrizzleMigrateConfig;
+
+    const fileContent = `export default ${JSON.stringify(migrationConfig, null, 2)}`;
+    const tmpfilePath = path.join(api.rootDir, "drizzle", "config.tmp.ts");
+
+    try {
+      await Bun.write(tmpfilePath, fileContent);
+      const { exitCode, stdout, stderr } =
+        await $`bun drizzle-kit generate:pg --config ${tmpfilePath}`;
+      logger.trace(stdout.toString());
+      if (exitCode !== 0) {
+        {
+          throw new TypedError({
+            message: `Failed to generate migrations: ${stderr.toString()}`,
+            type: ErrorType.SERVER_INITIALIZATION,
+          });
+        }
+      }
+    } finally {
+      const filePointer = Bun.file(tmpfilePath);
+      if (await filePointer.exists()) await unlink(tmpfilePath);
+    }
+  }
+
+  /**
+   * Erase all the tables in the active database.  Will fail on production environments.
+   */
+  async clearDatabase(restartIdentity = true, cascade = true) {
+    if (Bun.env.NODE_ENV === "production") {
+      throw new TypedError({
+        message: "clearDatabase cannot be called in production",
+        type: ErrorType.SERVER_INITIALIZATION,
+      });
+    }
+
+    const { rows } = await api.db.db.execute(
+      sql`SELECT tablename FROM pg_tables WHERE schemaname = CURRENT_SCHEMA`,
+    );
+
+    for (const row of rows) {
+      logger.debug(`truncating table ${row.tablename}`);
+      await api.db.db.execute(
+        sql.raw(
+          `TRUNCATE TABLE "${row.tablename}" ${restartIdentity ? "RESTART IDENTITY" : ""} ${cascade ? "CASCADE" : ""} `,
+        ),
+      );
+    }
+  }
+}