feat(PLA-3932): connect sockets via WebSocket gateway

packages/client/src/api.ts

@@ -1,6 +1,10 @@
 import { Socket } from "phoenix";
 
-import { exponentialBackoffFullJitter } from "./helpers";
+import {
+  createGatewayTransport,
+  fetchSocketConnectUrl,
+} from "./gatewayTransport";
+import { cacheUntilExpired, exponentialBackoffFullJitter } from "./helpers";
 import { PageVisibilityManager } from "./pageVisibility";
 
 // Ceiling for socket reconnect backoff. This is intentionally much larger than
@@ -123,7 +127,21 @@ class ApiClient {
     });
 
     if (typeof window !== "undefined" && typeof WebSocket !== "undefined") {
+      // Connect sockets via the WebSocket gateway API endpoint, which returns a signed ws:// or
+      // wss:// connect URL. This signed URL is cached and used to reconnect until it expires.
+      const connectUrlCache = cacheUntilExpired((signal: AbortSignal) =>
+        // If this fetch rejects, the error is handled by the gateway transport, and Phoenix
+        // considers it a failed connection attempt
+        fetchSocketConnectUrl({ apiClient: this, signal }),
+      );
+      const transport = createGatewayTransport({
+        baseTransport: WebSocket,
+        fetchConnectUrl: connectUrlCache.get,
+        onConnectUrlInvalid: connectUrlCache.invalidate,
+      });
+
       this.socket = new Socket(`${this.host.replace("http", "ws")}/ws/v1`, {
+        transport: transport,
         params: {
           user_token: this.userToken,
           api_key: this.apiKey,

packages/client/src/gatewayTransport.ts

@@ -0,0 +1,272 @@
+import type ApiClient from "./api";
+import { appendQueryString, getQueryString } from "./helpers";
+
+/**
+ * The signature that Phoenix expects socket transports to conform to, based on the native browser
+ * `WebSocket`. Only fields and methods that Phoenix relies on are included here.
+ */
+export type WebSocketLike = Pick<
+  WebSocket,
+  | "readyState"
+  | "bufferedAmount"
+  | "binaryType"
+  | "onopen"
+  | "onclose"
+  | "onmessage"
+  | "onerror"
+  | "send"
+  | "close"
+>;
+
+export type TransportConstructor = new (
+  url: string,
+  protocols?: string | string[],
+) => WebSocketLike;
+
+export type GatewayConnectResult = {
+  signedUrl: string;
+  expiresAt: number;
+};
+
+export type GatewayTransportOptions = {
+  /**
+   * The underlying transport constructor. An instance of this transport will be created once a
+   * connect URL has been fetched.
+   */
+  baseTransport: TransportConstructor;
+  /**
+   * Resolves the signed WebSocket URL to connect to. Rejects on failure.
+   */
+  fetchConnectUrl: (signal: AbortSignal) => Promise<GatewayConnectResult>;
+  /**
+   * Called when a connection made with the fetched URL closed before it ever
+   * opened, which most likely means the URL was rejected (though a network blip
+   * could look the same).
+   */
+  onConnectUrlInvalid?: () => void;
+};
+
+/**
+ * All possible socket states. Copied from Phoenix, based on the `readyState` property of
+ * `WebSocket`.
+ * @see {@link https://github.com/phoenixframework/phoenix/blob/1bae3a639f84cd97ccb6b984a8ef4fcdc7aa48ae/assets/js/phoenix/constants.js#L5}
+ */
+export const SOCKET_STATES = {
+  connecting: 0,
+  open: 1,
+  closing: 2,
+  closed: 3,
+} as const;
+
+// WebSocket close event status codes.
+// See https://www.rfc-editor.org/info/rfc6455/#section-7.4.1
+const CLOSE_CODE_NORMAL = 1000;
+const CLOSE_CODE_ABNORMAL = 1006;
+
+/**
+ * Requests a signed WebSocket URL from the WebSocket gateway API endpoint. Rejects on any HTTP
+ * error or malformed body so the transport treats it like a failed connection attempt.
+ */
+export async function fetchSocketConnectUrl({
+  apiClient,
+  signal,
+}: {
+  apiClient: ApiClient;
+  signal: AbortSignal;
+}): Promise<GatewayConnectResult> {
+  const result = await apiClient.makeRequest({
+    method: "POST",
+    url: "/v1/sockets/connect",
+    signal,
+    // Set maxRetries to 0, so that the API client's own retry logic is never used. Phoenix
+    // handles retries automatically via the socket's reconnect backoff.
+    maxRetries: 0,
+  });
+
+  if (result.statusCode !== "ok") {
+    throw (
+      result.error ??
+      new Error(`Socket connect request failed with status ${result.status}`)
+    );
+  }
+
+  const body = result.body as
+    | Partial<{ url: string; expires_at: string }>
+    | undefined;
+
+  const expiresAt =
+    typeof body?.expires_at === "string" ? Date.parse(body.expires_at) : NaN;
+
+  if (typeof body?.url !== "string" || !body.url || Number.isNaN(expiresAt)) {
+    throw new Error("Invalid socket connect response");
+  }
+
+  return { signedUrl: body.url, expiresAt };
+}
+
+export function createGatewayTransport({
+  baseTransport,
+  fetchConnectUrl,
+  onConnectUrlInvalid,
+}: GatewayTransportOptions): TransportConstructor {
+  // Use a named class, since Phoenix logs using the class name
+  return class GatewayTransport implements WebSocketLike {
+    binaryType: BinaryType = "blob";
+
+    onopen: ((event: Event) => void) | null = null;
+    onclose: ((event: CloseEvent) => void) | null = null;
+    onmessage: ((event: MessageEvent) => void) | null = null;
+    onerror: ((event: unknown) => void) | null = null;
+
+    private inner: WebSocketLike | null = null;
+    private state: number = SOCKET_STATES.connecting;
+    private readonly abort = new AbortController();
+    private readonly unsignedQueryString: string;
+    private opened = false;
+    private closeRequested = false;
+    private settled = false;
+    private readonly timeoutId: ReturnType<typeof setTimeout>;
+
+    constructor(
+      url: string,
+      private readonly protocols?: string | string[],
+    ) {
+      // The base URL passed in by Phoenix contains some query params we want to preserve, including
+      // vsn (protocol version to send on connect), api_key, user_token, and branch_slug. This query
+      // string is considered "unsigned" because it lacks an auth_token query param.
+      this.unsignedQueryString = getQueryString(url);
+
+      // Wait for the caller to finish setting up callbacks and properties. This pattern was copied
+      // from Phoenix's LongPoll transport.
+      this.timeoutId = setTimeout(() => this.start(), 0);
+    }
+
+    get readyState(): number {
+      return this.inner ? this.inner.readyState : this.state;
+    }
+
+    get bufferedAmount(): number {
+      return this.inner ? this.inner.bufferedAmount : 0;
+    }
+
+    send(data: string | ArrayBufferLike | ArrayBufferView | Blob): void {
+      if (!this.inner) {
+        // Phoenix only sends when a connection is open, so prevent calls to send before a
+        // connection is established.
+        throw new Error(
+          "GatewayTransport: cannot send before the connection is open",
+        );
+      }
+      this.inner.send(data);
+    }
+
+    close(code?: number, reason?: string): void {
+      if (this.closeRequested || this.settled) {
+        return;
+      }
+
+      this.closeRequested = true;
+
+      // In case start() has not run yet, clear the timeout used to invoke it in the constructor
+      clearTimeout(this.timeoutId);
+
+      if (this.inner) {
+        this.inner.close(code, reason);
+        return;
+      }
+
+      // The socket is closing before the fetch completed. This is still considered a clean close.
+      this.abort.abort();
+      this.state = SOCKET_STATES.closed;
+      this.emitClose(CLOSE_CODE_NORMAL, "", true);
+    }
+
+    private async start(): Promise<void> {
+      // Phoenix may connect() and disconnect() in the same tick; teardown then calls close() on us
+      // before this runs. This prevents a fetch that would be discarded.
+      if (this.closeRequested) {
+        return;
+      }
+
+      let result: GatewayConnectResult;
+      try {
+        result = await fetchConnectUrl(this.abort.signal);
+      } catch (error) {
+        if (this.closeRequested) {
+          // Ignore this error; the socket was already closed cleanly.
+          return;
+        }
+        // An unexepected error occurred while fetching
+        this.fail(error);
+        return;
+      }
+
+      // The socket was closed while a fetch was in-progress
+      if (this.closeRequested) {
+        return;
+      }
+
+      const connectUrl = appendQueryString(
+        result.signedUrl,
+        this.unsignedQueryString,
+      );
+
+      try {
+        this.inner = this.initBaseTransport(connectUrl);
+      } catch (error) {
+        // The base transport's constructor threw an exception. We never even attempted a
+        // connection, so onclose won't fire. The connect URL must be invalid.
+        onConnectUrlInvalid?.();
+        this.fail(error);
+      }
+    }
+
+    private initBaseTransport(url: string): WebSocketLike {
+      const inner = new baseTransport(url, this.protocols);
+
+      inner.binaryType = this.binaryType;
+      inner.onopen = (ev) => {
+        this.opened = true;
+        this.onopen?.(ev);
+      };
+      inner.onmessage = (ev) => this.onmessage?.(ev);
+      inner.onerror = (ev) => this.onerror?.(ev);
+      inner.onclose = (ev) => {
+        if (!this.opened && !this.closeRequested) {
+          // A connection was never opened and this close event was unexpected.
+          // The signed connect URL is likely invalid.
+          onConnectUrlInvalid?.();
+        }
+        this.settled = true;
+        this.onclose?.(ev);
+      };
+
+      return inner;
+    }
+
+    private fail(error: unknown): void {
+      this.state = SOCKET_STATES.closed;
+      this.onerror?.(error);
+      this.emitClose(CLOSE_CODE_ABNORMAL, "gateway_connect_failed", false);
+    }
+
+    private emitClose(code: number, reason: string, wasClean: boolean): void {
+      // Check this.settled to prevent triggering onclose more than once. For example, this could
+      // occur if the transport is closed manually inside of an onerror callback.
+      if (this.settled) {
+        return;
+      }
+
+      this.settled = true;
+
+      // Mirrors Phoenix's LongPoll.close: CloseEvent exists in browsers.
+      // Fall back to a plain object elsewhere (e.g. React Native).
+      const opts = { code, reason, wasClean };
+      if (typeof CloseEvent !== "undefined") {
+        this.onclose?.(new CloseEvent("close", opts));
+      } else {
+        this.onclose?.(opts as CloseEvent);
+      }
+    }
+  };
+}

packages/client/src/helpers.ts

@@ -1,10 +1,41 @@
+import { Store } from "@tanstack/store";
+
 const uuidRegex =
   /^[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}$/;
 
 export function isValidUuid(uuid: string) {
   return uuidRegex.test(uuid);
 }
 
+/**
+ * Returns the query string for a given URL, minus the leading `"?"` character. When the URL has no
+ * query string, returns the empty string.
+ *
+ * This function does not use the `URL` class due to known issues in React Native's implementation
+ * of `URL`.
+ *
+ * @param url a valid string URL. For example, `"https://example.com/test?hello=world"`.
+ */
+export function getQueryString(url: string): string {
+  const idx = url.indexOf("?");
+  return idx === -1 ? "" : url.slice(idx + 1);
+}
+
+/**
+ * Appends a query string to a given URL. Note that this does not dedupe query params.
+ *
+ * This function does not use the `URL` class due to known issues in React Native's implementation
+ * of `URL`.
+ *
+ * @param url a valid string URL.
+ * @param queryString a query string, minus any leading `"?"` character.
+ */
+export function appendQueryString(url: string, queryString: string): string {
+  return queryString
+    ? `${url}${url.includes("?") ? "&" : "?"}${queryString}`
+    : url;
+}
+
 /**
  * Exponential backoff with full jitter and a minimum delay floor.
  *
@@ -43,3 +74,44 @@ export function exponentialBackoffFullJitter(
   const jitterRange = exponentialDelay - minDelayMs;
   return minDelayMs + Math.floor(Math.random() * jitterRange);
 }
+
+export type CachedFetch<Args extends unknown[], T> = {
+  /**
+   * Returns the cached value if still valid, otherwise calls `fetchFunc` with the same args and
+   * caches the result.
+   */
+  get: (...args: Args) => Promise<T>;
+  /**
+   * Drops any cached value so the next `get` call fetches again.
+   */
+  invalidate: () => void;
+};
+
+export const CACHE_EXPIRY_PADDING_MS = 5_000;
+
+/**
+ * Wraps `fetchFunc` so its result is cached until `expiresAt` (epoch ms), minus
+ * `CACHE_EXPIRY_PADDING_MS`.
+ */
+export function cacheUntilExpired<
+  Args extends unknown[],
+  T extends { expiresAt: number },
+>(fetchFunc: (...args: Args) => Promise<T>): CachedFetch<Args, T> {
+  const store = new Store<T | null>(null);
+
+  const get = async (...args: Args) => {
+    const cached = store.state;
+    if (cached && cached.expiresAt - CACHE_EXPIRY_PADDING_MS > Date.now()) {
+      return cached;
+    }
+    const value = await fetchFunc(...args);
+    store.setState(() => value);
+    return value;
+  };
+
+  const invalidate = () => {
+    store.setState(() => null);
+  };
+
+  return { get, invalidate };
+}