@enbox/dwn-server 0.0.2 → 0.0.4

package/src/admin/webhook-manager.ts

@@ -0,0 +1,245 @@
+import type { Dialect } from '@enbox/dwn-sql-store';
+import type { AdminWebhook, AdminWebhookInput } from './types.js';
+
+import { createHmac } from 'crypto';
+import { Kysely } from 'kysely';
+import log from 'loglevel';
+import { v4 as uuidv4 } from 'uuid';
+
+/**
+ * Payload delivered to webhook endpoints.
+ */
+export type WebhookPayload = {
+  /** Unique delivery ID. */
+  id : string;
+  /** ISO-8601 timestamp. */
+  timestamp : string;
+  /** Event type (e.g. `tenant.suspend`, `quota.warning`). */
+  event : string;
+  /** Optional target (e.g. tenant DID). */
+  target? : string;
+  /** Additional event data. */
+  data? : Record<string, unknown>;
+};
+
+/**
+ * Manages webhook registrations and event delivery.
+ *
+ * Webhooks are stored in a SQL table and fire asynchronously when matching
+ * events occur. Delivery includes retry logic and optional HMAC-SHA256 signatures.
+ *
+ * @see https://github.com/enboxorg/enbox/issues/395
+ */
+export class WebhookManager {
+  static readonly #tableName = 'adminWebhooks';
+  static readonly #maxRetries = 3;
+  static readonly #retryDelaysMs = [1000, 5000, 15000];
+
+  #db: Kysely<WebhookDatabase>;
+
+  private constructor(dialect: Dialect) {
+    this.#db = new Kysely<WebhookDatabase>({ dialect });
+  }
+
+  /**
+   * Creates and initializes a `WebhookManager` instance.
+   */
+  public static async create(dialect: Dialect): Promise<WebhookManager> {
+    const manager = new WebhookManager(dialect);
+    await manager.#initialize();
+    return manager;
+  }
+
+  async #initialize(): Promise<void> {
+    await this.#db.schema
+      .createTable(WebhookManager.#tableName)
+      .ifNotExists()
+      .addColumn('id', 'text', (col) => col.primaryKey())
+      .addColumn('url', 'text', (col) => col.notNull())
+      .addColumn('events', 'text', (col) => col.notNull()) // JSON array
+      .addColumn('secret', 'text')
+      .addColumn('createdAt', 'text', (col) => col.notNull())
+      .execute();
+  }
+
+  // ---------------------------------------------------------------------------
+  // CRUD
+  // ---------------------------------------------------------------------------
+
+  /**
+   * Registers a new webhook. Returns the created webhook (without secret echoed).
+   */
+  public async register(input: AdminWebhookInput): Promise<AdminWebhook> {
+    const webhook: AdminWebhook = {
+      id        : uuidv4(),
+      url       : input.url,
+      events    : input.events,
+      createdAt : new Date().toISOString(),
+    };
+
+    await this.#db
+      .insertInto(WebhookManager.#tableName)
+      .values({
+        id        : webhook.id,
+        url       : webhook.url,
+        events    : JSON.stringify(webhook.events),
+        secret    : input.secret ?? null,
+        createdAt : webhook.createdAt,
+      })
+      .execute();
+
+    return webhook;
+  }
+
+  /**
+   * Lists all registered webhooks. Secrets are redacted.
+   */
+  public async list(): Promise<AdminWebhook[]> {
+    const rows = await this.#db
+      .selectFrom(WebhookManager.#tableName)
+      .select(['id', 'url', 'events', 'secret', 'createdAt'])
+      .orderBy('createdAt', 'asc')
+      .execute();
+
+    return rows.map((row): AdminWebhook => ({
+      id        : row.id,
+      url       : row.url,
+      events    : JSON.parse(row.events),
+      secret    : row.secret ? '***' : undefined,
+      createdAt : row.createdAt,
+    }));
+  }
+
+  /**
+   * Deletes a webhook by ID. Returns `true` if found and deleted.
+   */
+  public async delete(id: string): Promise<boolean> {
+    const result = await this.#db
+      .deleteFrom(WebhookManager.#tableName)
+      .where('id', '=', id)
+      .executeTakeFirstOrThrow();
+
+    return Number(result.numDeletedRows) > 0;
+  }
+
+  // ---------------------------------------------------------------------------
+  // Event dispatch
+  // ---------------------------------------------------------------------------
+
+  /**
+   * Fires matching webhooks for a given event. Non-blocking — delivery happens
+   * asynchronously and errors are logged but never propagated.
+   */
+  public fire(event: string, target?: string, data?: Record<string, unknown>): void {
+    // Fire-and-forget.
+    this.#dispatchAll(event, target, data).catch((err): void => {
+      log.error('Webhook dispatch error:', err);
+    });
+  }
+
+  async #dispatchAll(event: string, target?: string, data?: Record<string, unknown>): Promise<void> {
+    const rows = await this.#db
+      .selectFrom(WebhookManager.#tableName)
+      .select(['id', 'url', 'events', 'secret'])
+      .execute();
+
+    const payload: WebhookPayload = {
+      id        : uuidv4(),
+      timestamp : new Date().toISOString(),
+      event,
+      target,
+      data,
+    };
+
+    const body = JSON.stringify(payload);
+
+    for (const row of rows) {
+      const patterns: string[] = JSON.parse(row.events);
+      if (this.#matchesEvent(event, patterns)) {
+        this.#deliver(row.url, body, row.secret).catch((err): void => {
+          log.warn(`Webhook delivery failed for ${row.url}:`, err);
+        });
+      }
+    }
+  }
+
+  /**
+   * Checks if an event matches any of the subscription patterns.
+   * Supports wildcard patterns like `tenant.*`.
+   */
+  #matchesEvent(event: string, patterns: string[]): boolean {
+    for (const pattern of patterns) {
+      if (pattern === '*') {
+        return true;
+      }
+      if (pattern === event) {
+        return true;
+      }
+      if (pattern.endsWith('.*')) {
+        const prefix = pattern.slice(0, -2);
+        if (event.startsWith(prefix + '.') || event === prefix) {
+          return true;
+        }
+      }
+    }
+    return false;
+  }
+
+  /**
+   * Delivers a webhook payload with retry logic and optional HMAC signature.
+   */
+  async #deliver(url: string, body: string, secret: string | null): Promise<void> {
+    const headers: Record<string, string> = {
+      'content-type': 'application/json',
+    };
+
+    if (secret) {
+      const signature = createHmac('sha256', secret).update(body).digest('hex');
+      headers['x-webhook-signature'] = `sha256=${signature}`;
+    }
+
+    for (let attempt = 0; attempt <= WebhookManager.#maxRetries; attempt++) {
+      try {
+        const response = await fetch(url, { method: 'POST', headers, body, signal: AbortSignal.timeout(10_000) });
+        if (response.ok) {
+          return;
+        }
+        log.warn(`Webhook ${url} returned ${response.status} (attempt ${attempt + 1})`);
+      } catch (err) {
+        log.warn(`Webhook ${url} fetch error (attempt ${attempt + 1}):`, err);
+      }
+
+      // Wait before retry (skip wait on last attempt).
+      if (attempt < WebhookManager.#maxRetries) {
+        await new Promise((resolve): ReturnType<typeof setTimeout> =>
+          setTimeout(resolve, WebhookManager.#retryDelaysMs[attempt]),
+        );
+      }
+    }
+
+    log.error(`Webhook delivery to ${url} failed after ${WebhookManager.#maxRetries + 1} attempts`);
+  }
+
+  /**
+   * Closes the underlying database connection.
+   */
+  public async close(): Promise<void> {
+    await this.#db.destroy();
+  }
+}
+
+// ---------------------------------------------------------------------------
+// Kysely type definitions
+// ---------------------------------------------------------------------------
+
+interface WebhookRow {
+  id : string;
+  url : string;
+  events : string; // JSON-serialized string[]
+  secret : string | null;
+  createdAt : string;
+}
+
+interface WebhookDatabase {
+  adminWebhooks : WebhookRow;
+}

package/src/config.ts

@@ -1,3 +1,5 @@
+import { readFileSync } from 'fs';
+
 import bytes from 'bytes';
 
 export type DwnServerConfig = typeof config;
@@ -5,11 +7,8 @@ export type DwnServerConfig = typeof config;
 export const config = {
   /**
    * Used to populate the `server` property returned by the `/info` endpoint.
-   *
-   * If running using `npm` the `process.env.npm_package_name` variable exists and we use that,
-   * otherwise we fall back on the use defined `DWN_SERVER_PACKAGE_NAME` or `@enbox/dwn-server`.
    */
-  serverName: process.env.npm_package_name || process.env.DWN_SERVER_PACKAGE_NAME || '@enbox/dwn-server',
+  serverName: process.env.DWN_SERVER_PACKAGE_NAME || '@enbox/dwn-server',
 
   /**
    * The base external URL of this DWN.
@@ -26,7 +25,7 @@ export const config = {
   /**
    * The URL of the TTL cache used by the DWN.
    * NOTE: Used for session/state keeping, thus requires the cache to be commonly addressable by nodes in a cloud cluster environment.
-   * 
+   *
    * Currently only supports SQL databases, e.g.
    * Postgres: 'postgres://root:dwn@localhost:5432/dwn'
    * MySQL: 'mysql://root:dwn@localhost:3306/dwn'
@@ -37,35 +36,204 @@ export const config = {
    * Used to populate the `version` and `sdkVersion` properties returned by the `/info` endpoint.
    *
    * The `version` and `sdkVersion` are pulled from `package.json` at runtime.
-   * If running using `npm` the `process.env.npm_package_json` variable exists as the filepath, so we use that.
-   * Otherwise we check to see if a specific `DWN_SERVER_PACKAGE_JSON` exists, if it does we use that.
-   * Finally if both of those options don't exist we resort to the path within the docker server image, located at `/dwn-server/package.json`
+   * If `DWN_SERVER_PACKAGE_JSON` is set, we use that path.
+   * Otherwise we resort to the path within the docker server image, located at `/dwn-server/package.json`.
+   */
+  packageJsonPath   : process.env.DWN_SERVER_PACKAGE_JSON || '/dwn-server/package.json',
+  /**
+   * Maximum size of data that can be provided with a RecordsWrite.
+   * Request bodies up to this size are buffered fully into memory, so the
+   * default should be conservative enough to prevent memory exhaustion from
+   * concurrent large uploads. Operators can raise the limit via the
+   * `MAX_RECORD_DATA_SIZE` env var (e.g. `'1gb'`).
    */
-  packageJsonPath:  process.env.npm_package_json ||  process.env.DWN_SERVER_PACKAGE_JSON || '/dwn-server/package.json',
-  // max size of data that can be provided with a RecordsWrite
-  maxRecordDataSize: bytes(process.env.MAX_RECORD_DATA_SIZE || '1gb'),
+  maxRecordDataSize : bytes(process.env.MAX_RECORD_DATA_SIZE || '100mb'),
+
+  /**
+   * Maximum number of unacknowledged subscription events the server will send
+   * per subscription before pausing delivery. Clients must send `rpc.ack` to
+   * advance the window. Configurable via `DWN_MAX_IN_FLIGHT` env var.
+   */
+  maxInFlight: parseInt(process.env.DWN_MAX_IN_FLIGHT || '32'),
 
   // whether to enable 'ws:'
   webSocketSupport: { on: true, off: false }[process.env.DS_WEBSOCKET_SERVER] ?? true,
 
   /**
-   * Path to DWN Event Stream plugin to use. Default single-node implementation will be used if left empty.
+   * Path to DWN EventLog plugin to use. Default in-memory implementation will be used if left empty.
+   * Also accepts the legacy `DWN_EVENT_STREAM_PLUGIN_PATH` env var for backward compatibility.
    */
-  eventStreamPluginPath: process.env.DWN_EVENT_STREAM_PLUGIN_PATH,
+  eventLogPluginPath: process.env.DWN_EVENT_LOG_PLUGIN_PATH || process.env.DWN_EVENT_STREAM_PLUGIN_PATH,
 
   // where to store persistent data
-  messageStore: process.env.DWN_STORAGE_MESSAGES || process.env.DWN_STORAGE || 'level://data',
-  dataStore: process.env.DWN_STORAGE_DATA || process.env.DWN_STORAGE || 'level://data',
-  eventLog: process.env.DWN_STORAGE_EVENTS || process.env.DWN_STORAGE || 'level://data',
-  resumableTaskStore: process.env.DWN_STORAGE_RESUMABLE_TASKS || process.env.DWN_STORAGE || 'level://data',
+  messageStore       : process.env.DWN_STORAGE_MESSAGES || process.env.DWN_STORAGE || 'level://data',
+  dataStore          : process.env.DWN_STORAGE_DATA || process.env.DWN_STORAGE || 'level://data',
+  stateIndex         : process.env.DWN_STORAGE_STATE_INDEX || process.env.DWN_STORAGE || 'level://data',
+  resumableTaskStore : process.env.DWN_STORAGE_RESUMABLE_TASKS || process.env.DWN_STORAGE || 'level://data',
+
+  /**
+   * PostgreSQL connection pool tuning. When multiple DWN stores share the same
+   * Postgres connection URL, a single shared pool is used instead of one pool
+   * per store (which would be 4 pools x 10 default connections = 40 connections).
+   */
+  pgPoolMin         : parseInt(process.env.DWN_PG_POOL_MIN || '5'),
+  pgPoolMax         : parseInt(process.env.DWN_PG_POOL_MAX || '30'),
+  pgPoolIdleTimeout : parseInt(process.env.DWN_PG_POOL_IDLE_TIMEOUT || '30000'),
 
   // tenant registration feature configuration
-  registrationStoreUrl: process.env.DWN_REGISTRATION_STORE_URL || process.env.DWN_STORAGE,
-  registrationProofOfWorkSeed: process.env.DWN_REGISTRATION_PROOF_OF_WORK_SEED,
-  registrationProofOfWorkEnabled: process.env.DWN_REGISTRATION_PROOF_OF_WORK_ENABLED === 'true',
-  registrationProofOfWorkInitialMaxHash: process.env.DWN_REGISTRATION_PROOF_OF_WORK_INITIAL_MAX_HASH,
-  termsOfServiceFilePath: process.env.DWN_TERMS_OF_SERVICE_FILE_PATH,
+  registrationStoreUrl                  : process.env.DWN_REGISTRATION_STORE_URL || process.env.DWN_STORAGE,
+  registrationProofOfWorkSeed           : process.env.DWN_REGISTRATION_PROOF_OF_WORK_SEED,
+  registrationProofOfWorkEnabled        : process.env.DWN_REGISTRATION_PROOF_OF_WORK_ENABLED === 'true',
+  registrationProofOfWorkInitialMaxHash : process.env.DWN_REGISTRATION_PROOF_OF_WORK_INITIAL_MAX_HASH,
+  termsOfServiceFilePath                : process.env.DWN_TERMS_OF_SERVICE_FILE_PATH,
+
+  // Provider auth configuration for paid DWN registration
+  providerAuthEnabled       : process.env.DWN_PROVIDER_AUTH_ENABLED === 'true',
+  providerAuthAuthorizeUrl  : process.env.DWN_PROVIDER_AUTH_AUTHORIZE_URL,
+  providerAuthTokenUrl      : process.env.DWN_PROVIDER_AUTH_TOKEN_URL,
+  providerAuthRefreshUrl    : process.env.DWN_PROVIDER_AUTH_REFRESH_URL,
+  providerAuthManagementUrl : process.env.DWN_PROVIDER_AUTH_MANAGEMENT_URL,
+  providerAuthPluginPath    : process.env.DWN_PROVIDER_AUTH_PLUGIN_PATH,
+  providerAuthJwtSecret     : process.env.DWN_PROVIDER_AUTH_JWT_SECRET,
+  providerAuthJwtJwksUrl    : process.env.DWN_PROVIDER_AUTH_JWT_JWKS_URL,
 
   // log level - trace/debug/info/warn/error
   logLevel: process.env.DWN_SERVER_LOG_LEVEL || 'INFO',
+
+  /**
+   * Bearer token for the admin API. If unset (or empty), the admin API is disabled entirely.
+   * Can also be read from a file path via `DWN_ADMIN_TOKEN_FILE` (useful for Docker secrets).
+   */
+  adminToken: process.env.DWN_ADMIN_TOKEN || (
+    process.env.DWN_ADMIN_TOKEN_FILE
+      ? readAdminTokenFromFile(process.env.DWN_ADMIN_TOKEN_FILE)
+      : undefined
+  ),
+
+  /**
+   * Maximum number of recent DWN activity events retained in the in-memory
+   * ring buffer for the admin `/events` endpoint. Defaults to 10,000.
+   */
+  adminActivityLogCapacity: parseInt(process.env.DWN_ADMIN_ACTIVITY_LOG_CAPACITY || '10000'),
+
+  /**
+   * Interval (in seconds) at which Prometheus gauge metrics are updated from
+   * the admin store. Defaults to 30 seconds.
+   */
+  adminMetricsUpdateIntervalSeconds: parseInt(process.env.DWN_ADMIN_METRICS_UPDATE_INTERVAL || '30'),
+
+  // ---------------------------------------------------------------------------
+  // Per-tenant storage quotas
+  // ---------------------------------------------------------------------------
+
+  /**
+   * Default maximum number of messages a tenant may store. 0 = unlimited (default).
+   * Per-tenant overrides are managed via the admin API.
+   */
+  quotaMaxMessages: parseInt(process.env.DWN_QUOTA_MAX_MESSAGES || '0'),
+
+  /**
+   * Default maximum data storage in bytes a tenant may use. 0 = unlimited (default).
+   * Per-tenant overrides are managed via the admin API.
+   */
+  quotaMaxStorageBytes: parseInt(process.env.DWN_QUOTA_MAX_STORAGE_BYTES || '0'),
+
+  // ---------------------------------------------------------------------------
+  // Audit log retention
+  // ---------------------------------------------------------------------------
+
+  /**
+   * Maximum age of audit log entries in days. Entries older than this are purged.
+   * 0 = no age limit (default: 90 days).
+   *
+   * @see https://github.com/enboxorg/enbox/issues/394
+   */
+  auditLogMaxAgeDays: parseInt(process.env.DWN_AUDIT_LOG_MAX_AGE_DAYS || '90'),
+
+  /**
+   * Maximum number of audit log rows to retain. Oldest entries are purged when exceeded.
+   * 0 = no row limit (default: 100000).
+   *
+   * @see https://github.com/enboxorg/enbox/issues/394
+   */
+  auditLogMaxRows: parseInt(process.env.DWN_AUDIT_LOG_MAX_ROWS || '100000'),
+
+  // ---------------------------------------------------------------------------
+  // Rate limiting
+  // ---------------------------------------------------------------------------
+
+  /**
+   * Maximum HTTP requests per second per IP address. Set to 0 to disable.
+   * Defaults to 30 req/s which is generous for normal usage while limiting abuse.
+   * Can be reconfigured at runtime via the admin `PATCH /config` endpoint.
+   */
+  rateLimitRequestsPerSecond: parseInt(process.env.DWN_RATE_LIMIT_REQUESTS_PER_SECOND || '30'),
+
+  /**
+   * Maximum burst size for per-IP rate limiting. Allows short spikes above the
+   * sustained rate without triggering 429s. Defaults to 50.
+   */
+  rateLimitBurst: parseInt(process.env.DWN_RATE_LIMIT_BURST || '50'),
+
+  /**
+   * Maximum DWN requests per second per tenant DID. Set to 0 to disable.
+   * Defaults to 20 req/s. Applies to both HTTP and WebSocket transports.
+   * Can be reconfigured at runtime via the admin `PATCH /config` endpoint.
+   */
+  rateLimitTenantRequestsPerSecond: parseInt(process.env.DWN_RATE_LIMIT_TENANT_REQUESTS_PER_SECOND || '20'),
+
+  /**
+   * Maximum burst size for per-tenant rate limiting. Defaults to 50.
+   */
+  rateLimitTenantBurst: parseInt(process.env.DWN_RATE_LIMIT_TENANT_BURST || '50'),
+
+  // ---------------------------------------------------------------------------
+  // Record delivery & endpoint forwarding
+  // ---------------------------------------------------------------------------
+
+  /**
+   * Enable endpoint forwarding: when a RecordsWrite/RecordsDelete is processed,
+   * forward the original signed message to the tenant's other DWN service
+   * endpoints (discovered via DID resolution). Disabled by default.
+   */
+  forwardingEnabled: process.env.DWN_FORWARDING_ENABLED === 'true',
+
+  /**
+   * Enable protocol-aware record delivery: when a RecordsWrite/RecordsDelete is
+   * processed at a protocol path with `$delivery`, proactively deliver to
+   * participants' DWN endpoints. Disabled by default.
+   */
+  deliveryEnabled: process.env.DWN_DELIVERY_ENABLED === 'true',
+
+  /**
+   * Maximum number of concurrent outbound delivery/forwarding requests.
+   * Prevents unbounded parallelism when delivering to many providers.
+   * Defaults to 10.
+   */
+  deliveryMaxConcurrency: parseInt(process.env.DWN_DELIVERY_MAX_CONCURRENCY || '10'),
+
+  /**
+   * TTL in seconds for caching DID document service endpoint resolution results.
+   * Avoids resolving the same DID document on every delivery. Defaults to 300 (5 min).
+   */
+  deliveryEndpointCacheTtlSeconds: parseInt(process.env.DWN_DELIVERY_ENDPOINT_CACHE_TTL || '300'),
+
+  /**
+   * TTL in seconds for the recently-forwarded messageCid deduplication cache.
+   * Messages with CIDs in this cache are not forwarded again, reducing redundant
+   * outbound requests between peer endpoints. Defaults to 60.
+   */
+  forwardingDeduplicationTtlSeconds: parseInt(process.env.DWN_FORWARDING_DEDUP_TTL || '60'),
 };
+
+/**
+ * Reads the admin token from a file path, trimming whitespace.
+ * Returns `undefined` if the file cannot be read.
+ */
+function readAdminTokenFromFile(filePath: string): string | undefined {
+  try {
+    return readFileSync(filePath).toString().trim() || undefined;
+  } catch {
+    return undefined;
+  }
+}

package/src/connection/connection-manager.ts

@@ -1,39 +1,89 @@
-import type { Dwn } from "@enbox/dwn-sdk-js";
+import type { Dwn } from '@enbox/dwn-sdk-js';
+import type { ServerWebSocket } from 'bun';
 
-import type { IncomingMessage } from "http";
-import type { WebSocket } from 'ws';
+import type { ActivityLog } from '../admin/activity-log.js';
+import type { AdminConnectionSnapshot } from '../admin/types.js';
+import type { AdminStore } from '../admin/admin-store.js';
+import type { DwnServerConfig } from '../config.js';
+import type { MessageProcessedHook } from '../message-processed-hook.js';
+import type { RateLimiter } from '../rate-limiter.js';
+import type { RegistrationStore } from '../registration/registration-store.js';
+import type { WsData } from '../http-api.js';
 
-import { SocketConnection } from "./socket-connection.js";
+import { SocketConnection } from './socket-connection.js';
 
 /**
  * Interface for managing `WebSocket` connections as they arrive.
  */
 export interface ConnectionManager {
-  /** connect handler used for the `WebSockets` `'connection'` event. */
-  connect(socket: WebSocket, request?: IncomingMessage): Promise<void>;
+  /** connect handler invoked when a new WebSocket connection is established. */
+  connect(socket: ServerWebSocket<WsData>): Promise<void>;
   /** closes all of the connections */
-  closeAll(): Promise<void>
+  closeAll(): Promise<void>;
+  /** Returns the number of active connections. */
+  getConnectionCount(): number;
+  /** Returns the total number of active subscriptions across all connections. */
+  getSubscriptionCount(): number;
+  /** Returns serializable snapshots of all active connections. */
+  getConnectionSnapshots(): AdminConnectionSnapshot[];
 }
 
 /**
  * A Simple In Memory ConnectionManager implementation.
- * It uses a `Map<WebSocket, SocketConnection>` to manage connections.
+ * It uses a `Map<ServerWebSocket, SocketConnection>` to manage connections.
  */
 export class InMemoryConnectionManager implements ConnectionManager {
-  constructor(private dwn: Dwn, private connections: Map<WebSocket, SocketConnection> = new Map()) {}
+  constructor(
+    private dwn: Dwn,
+    private connections: Map<ServerWebSocket<WsData>, SocketConnection> = new Map(),
+    private maxInFlight?: number,
+    private activityLog?: ActivityLog,
+    private adminStore?: AdminStore,
+    private registrationStore?: RegistrationStore,
+    private serverConfig?: DwnServerConfig,
+    private tenantRateLimiter?: RateLimiter,
+    private messageProcessedHooks?: MessageProcessedHook[],
+  ) {}
 
-  async connect(socket: WebSocket): Promise<void> {
-    const connection = new SocketConnection(socket, this.dwn, () => {
-      // this is the onClose handler to clean up any closed connections.
-      this.connections.delete(socket);
-    });
+  async connect(socket: ServerWebSocket<WsData>): Promise<void> {
+    const connection = new SocketConnection(
+      socket, this.dwn, () => {
+        // this is the onClose handler to clean up any closed connections.
+        this.connections.delete(socket);
+      },
+      this.maxInFlight, this.activityLog,
+      this.adminStore, this.registrationStore, this.serverConfig, this.tenantRateLimiter, this.messageProcessedHooks,
+    );
+
+    // Attach the connection to the ws.data so Bun's websocket handlers can delegate to it.
+    socket.data.connection = connection;
 
     this.connections.set(socket, connection);
   }
 
   async closeAll(): Promise<void> {
-    const closePromises = [];
-    this.connections.forEach(connection => closePromises.push(connection.close()));
+    const closePromises: Promise<void>[] = [];
+    this.connections.forEach((connection) => closePromises.push(connection.close()));
     await Promise.all(closePromises);
   }
-}
+
+  getConnectionCount(): number {
+    return this.connections.size;
+  }
+
+  getSubscriptionCount(): number {
+    let count = 0;
+    this.connections.forEach((conn) => {
+      count += conn.subscriptionCount;
+    });
+    return count;
+  }
+
+  getConnectionSnapshots(): AdminConnectionSnapshot[] {
+    const snapshots: AdminConnectionSnapshot[] = [];
+    this.connections.forEach((conn) => {
+      snapshots.push(conn.toSnapshot());
+    });
+    return snapshots;
+  }
+}