Fix startup WebSocket readiness

firefox-extension/__tests__/client.test.ts

@@ -0,0 +1,103 @@
+import { WebsocketClient } from "../client";
+
+class MockWebSocket extends EventTarget {
+  static CONNECTING = 0;
+  static OPEN = 1;
+  static CLOSING = 2;
+  static CLOSED = 3;
+  static instances: MockWebSocket[] = [];
+
+  readyState = MockWebSocket.CONNECTING;
+  sentMessages: string[] = [];
+
+  constructor(readonly url: string) {
+    super();
+    MockWebSocket.instances.push(this);
+  }
+
+  send(message: string) {
+    this.sentMessages.push(message);
+  }
+
+  close() {
+    if (this.readyState === MockWebSocket.CLOSED) {
+      return;
+    }
+    this.readyState = MockWebSocket.CLOSED;
+    this.dispatchEvent(new Event("close"));
+  }
+}
+
+describe("WebsocketClient", () => {
+  const originalWebSocket = global.WebSocket;
+
+  beforeEach(() => {
+    jest.useFakeTimers();
+    MockWebSocket.instances = [];
+    Object.defineProperty(global, "WebSocket", {
+      value: MockWebSocket,
+      writable: true,
+      configurable: true,
+    });
+  });
+
+  afterEach(() => {
+    jest.useRealTimers();
+    Object.defineProperty(global, "WebSocket", {
+      value: originalWebSocket,
+      writable: true,
+      configurable: true,
+    });
+  });
+
+  it("reconnects after the websocket closes", () => {
+    const client = new WebsocketClient(8089, "test-secret");
+
+    client.connect();
+    expect(MockWebSocket.instances).toHaveLength(1);
+
+    MockWebSocket.instances[0].close();
+    jest.advanceTimersByTime(1999);
+    expect(MockWebSocket.instances).toHaveLength(1);
+
+    jest.advanceTimersByTime(1);
+    expect(MockWebSocket.instances).toHaveLength(2);
+    expect(MockWebSocket.instances[1].url).toBe("ws://localhost:8089");
+  });
+
+  it("retries when a websocket stays stuck connecting", () => {
+    const client = new WebsocketClient(8089, "test-secret");
+
+    client.connect();
+    expect(MockWebSocket.instances).toHaveLength(1);
+
+    jest.advanceTimersByTime(1999);
+    expect(MockWebSocket.instances).toHaveLength(1);
+    expect(MockWebSocket.instances[0].readyState).toBe(MockWebSocket.CONNECTING);
+
+    jest.advanceTimersByTime(1);
+    expect(MockWebSocket.instances[0].readyState).toBe(MockWebSocket.CLOSED);
+
+    jest.advanceTimersByTime(2000);
+    expect(MockWebSocket.instances).toHaveLength(2);
+  });
+
+  it("does not create duplicate sockets while already connecting", () => {
+    const client = new WebsocketClient(8089, "test-secret");
+
+    client.connect();
+    client.connect();
+
+    expect(MockWebSocket.instances).toHaveLength(1);
+  });
+
+  it("does not reconnect after an explicit disconnect", () => {
+    const client = new WebsocketClient(8089, "test-secret");
+
+    client.connect();
+    client.disconnect();
+    jest.advanceTimersByTime(2000);
+
+    expect(MockWebSocket.instances).toHaveLength(1);
+  });
+});

firefox-extension/client.ts

@@ -6,13 +6,14 @@ import type {
 import { getMessageSignature } from "./auth";
 
 const RECONNECT_INTERVAL = 2000; // 2 seconds
+const CONNECT_TIMEOUT = 2000; // 2 seconds
 
 export class WebsocketClient {
   private socket: WebSocket | null = null;
   private readonly port: number;
   private readonly secret: string;
   private reconnectTimer: number | null = null;
-  private connectionAttempts: number = 0;
+  private manuallyDisconnected = false;
   private messageCallback: ((data: ServerMessageRequest) => void) | null = null;
 
   constructor(port: number, secret: string) {
@@ -21,25 +22,55 @@ export class WebsocketClient {
   }
 
   public connect(): void {
-    console.log("Connecting to WebSocket server at port", this.port);
+    this.manuallyDisconnected = false;
+    if (
+      this.socket &&
+      (this.socket.readyState === WebSocket.OPEN ||
+        this.socket.readyState === WebSocket.CONNECTING)
+    ) {
+      return;
+    }
 
-    this.socket = new WebSocket(`ws://localhost:${this.port}`);
+    console.log("Connecting to WebSocket server at port", this.port);
+    this.clearReconnectTimer();
+
+    const socket = new WebSocket(`ws://localhost:${this.port}`);
+    this.socket = socket;
+    const connectTimeout = window.setTimeout(() => {
+      if (
+        this.socket === socket &&
+        socket.readyState === WebSocket.CONNECTING
+      ) {
+        console.log("WebSocket connection attempt timed out at port", this.port);
+        socket.close();
+      }
+    }, CONNECT_TIMEOUT);
 
-    this.socket.addEventListener("open", () => {
+    socket.addEventListener("open", () => {
+      window.clearTimeout(connectTimeout);
       console.log("Connected to WebSocket server at port", this.port);
-      this.connectionAttempts = 0;
     });
 
-    this.socket.addEventListener("close", () => {
+    socket.addEventListener("close", () => {
+      window.clearTimeout(connectTimeout);
       console.log("WebSocket connection closed event at port", this.port);
-      this.connectionAttempts = 0;
+      if (this.socket === socket) {
+        this.socket = null;
+      }
+      this.scheduleReconnect();
     });
 
-    this.socket.addEventListener("error", (event) => {
+    socket.addEventListener("error", (event) => {
       console.error("WebSocket error:", event);
+      if (
+        socket.readyState !== WebSocket.CLOSING &&
+        socket.readyState !== WebSocket.CLOSED
+      ) {
+        socket.close();
+      }
     });
 
-    this.socket.addEventListener("message", async (event) => {
+    socket.addEventListener("message", async (event) => {
       if (this.messageCallback === null) {
         return;
       }
@@ -62,11 +93,6 @@ export class WebsocketClient {
         console.error("Failed to parse message:", error);
       }
     });
-
-    // Start reconnection timer if not already running
-    if (this.reconnectTimer === null) {
-      this.startReconnectTimer();
-    }
   }
 
   public addMessageListener(
@@ -75,26 +101,29 @@ export class WebsocketClient {
     this.messageCallback = callback;
   }
 
-  private startReconnectTimer(): void {
-    this.reconnectTimer = window.setInterval(() => {
-      if (this.socket && this.socket.readyState === WebSocket.CONNECTING) {
-        this.connectionAttempts++;
-
-        if (this.connectionAttempts > 2) {
-          // Avoid long retry backoff periods by resetting the connection
-          this.socket.close();
-        }
-      }
+  private scheduleReconnect(): void {
+    if (this.manuallyDisconnected || this.reconnectTimer !== null) {
+      return;
+    }
 
-      if (!this.socket || this.socket.readyState === WebSocket.CLOSED) {
-        this.connect();
-      }
+    this.reconnectTimer = window.setTimeout(() => {
+      this.reconnectTimer = null;
+      this.connect();
     }, RECONNECT_INTERVAL);
   }
 
+  private clearReconnectTimer(): void {
+    if (this.reconnectTimer === null) {
+      return;
+    }
+    window.clearTimeout(this.reconnectTimer);
+    this.reconnectTimer = null;
+  }
+
   public async sendResourceToServer(resource: ExtensionMessage): Promise<void> {
     if (!this.socket || this.socket.readyState !== WebSocket.OPEN) {
       console.error("Socket is not open");
+      this.scheduleReconnect();
       return;
     }
     const signedMessage = {
@@ -113,6 +142,7 @@ export class WebsocketClient {
   ): Promise<void> {
     if (!this.socket || this.socket.readyState !== WebSocket.OPEN) {
       console.error("Socket is not open", this.socket);
+      this.scheduleReconnect();
       return;
     }
     const extensionError: ExtensionError = {
@@ -123,10 +153,8 @@ export class WebsocketClient {
   }
 
   public disconnect(): void {
-    if (this.reconnectTimer !== null) {
-      window.clearInterval(this.reconnectTimer);
-      this.reconnectTimer = null;
-    }
+    this.manuallyDisconnected = true;
+    this.clearReconnectTimer();
 
     if (this.socket) {
       this.socket.close();