[pull] main from cloudflare:main

.changeset/selfish-bananas-clean.md

@@ -0,0 +1,5 @@
+---
+"partyserver": patch
+---
+
+Check for hibernated websocket connections

packages/partyserver/src/connection.ts

@@ -38,6 +38,39 @@ type ConnectionAttachments = {
   __user?: unknown;
 };
 
+function tryGetPartyServerMeta(
+  ws: WebSocket
+): ConnectionAttachments["__pk"] | null {
+  try {
+    // Avoid AttachmentCache.get() here: hibernated sockets accepted outside
+    // PartyServer can have an attachment without a __pk namespace.
+    const attachment = WebSocket.prototype.deserializeAttachment.call(
+      ws
+    ) as unknown;
+    if (!attachment || typeof attachment !== "object") {
+      return null;
+    }
+    if (!("__pk" in attachment)) {
+      return null;
+    }
+    const pk = (attachment as ConnectionAttachments).__pk as unknown;
+    if (!pk || typeof pk !== "object") {
+      return null;
+    }
+    const { id, server } = pk as { id?: unknown; server?: unknown };
+    if (typeof id !== "string" || typeof server !== "string") {
+      return null;
+    }
+    return pk as ConnectionAttachments["__pk"];
+  } catch {
+    return null;
+  }
+}
+
+export function isPartyServerWebSocket(ws: WebSocket): boolean {
+  return tryGetPartyServerMeta(ws) !== null;
+}
+
 /**
  * Cache websocket attachments to avoid having to rehydrate them on every property access.
  */
@@ -180,6 +213,12 @@ class HibernatingConnectionIterator<T> implements IterableIterator<
     while ((socket = sockets[this.index++])) {
       // only yield open sockets to match non-hibernating behaviour
       if (socket.readyState === WebSocket.READY_STATE_OPEN) {
+        // Durable Objects hibernation APIs allow storing arbitrary sockets via
+        // `state.acceptWebSocket()`. Those sockets won't have PartyServer's
+        // `__pk` attachment namespace and must be ignored.
+        if (!isPartyServerWebSocket(socket)) {
+          continue;
+        }
         const value = createLazyConnection(socket) as Connection<T>;
         return { done: false, value };
       }
@@ -263,15 +302,25 @@ export class HibernatingConnectionManager<TState> implements ConnectionManager {
   constructor(private controller: DurableObjectState) {}
 
   getCount() {
-    return Number(this.controller.getWebSockets().length);
+    // Only count sockets managed by PartyServer. Other hibernated sockets may
+    // exist on the same Durable Object via `state.acceptWebSocket()`.
+    let count = 0;
+    for (const ws of this.controller.getWebSockets()) {
+      if (isPartyServerWebSocket(ws)) count++;
+    }
+    return count;
   }
 
   getConnection<T = TState>(id: string) {
     // TODO: Should we cache the connections?
     const sockets = this.controller.getWebSockets(id);
-    if (sockets.length === 0) return undefined;
-    if (sockets.length === 1)
-      return createLazyConnection(sockets[0]) as Connection<T>;
+    const matching = sockets.filter((ws) => {
+      return tryGetPartyServerMeta(ws)?.id === id;
+    });
+
+    if (matching.length === 0) return undefined;
+    if (matching.length === 1)
+      return createLazyConnection(matching[0]) as Connection<T>;
 
     throw new Error(
       `More than one connection found for id ${id}. Did you mean to use getConnections(tag) instead?`

packages/partyserver/src/index.ts

@@ -8,7 +8,8 @@ import { nanoid } from "nanoid";
 import {
   createLazyConnection,
   HibernatingConnectionManager,
-  InMemoryConnectionManager
+  InMemoryConnectionManager,
+  isPartyServerWebSocket
 } from "./connection";
 
 import type { ConnectionManager } from "./connection";
@@ -422,6 +423,13 @@ Did you try connecting directly to this Durable Object? Try using getServerByNam
       return;
     }
 
+    // Ignore websockets accepted outside PartyServer (e.g. via
+    // `state.acceptWebSocket()` in user code). These sockets won't have the
+    // `__pk` attachment namespace required to rehydrate a Connection.
+    if (!isPartyServerWebSocket(ws)) {
+      return;
+    }
+
     const connection = createLazyConnection(ws);
 
     // rehydrate the server name if it's woken up
@@ -449,6 +457,10 @@ Did you try connecting directly to this Durable Object? Try using getServerByNam
       return;
     }
 
+    if (!isPartyServerWebSocket(ws)) {
+      return;
+    }
+
     const connection = createLazyConnection(ws);
 
     // rehydrate the server name if it's woken up
@@ -470,6 +482,10 @@ Did you try connecting directly to this Durable Object? Try using getServerByNam
       return;
     }
 
+    if (!isPartyServerWebSocket(ws)) {
+      return;
+    }
+
     const connection = createLazyConnection(ws);
 
     // rehydrate the server name if it's woken up

packages/partyserver/src/tests/index.test.ts

@@ -149,6 +149,47 @@ describe("Server", () => {
     expect(response.headers.get("Location")).toBe("https://example3.com");
   });
 
+  it("ignores foreign hibernated websockets when broadcasting", async () => {
+    const ctx = createExecutionContext();
+
+    // Create a websocket that is accepted via the DO hibernation API directly
+    // (no PartyServer `__pk` attachment).
+    const foreignReq = new Request(
+      "http://example.com/parties/mixed/room/foreign",
+      {
+        headers: { Upgrade: "websocket" }
+      }
+    );
+    const foreignRes = await worker.fetch(foreignReq, env, ctx);
+    const foreignWs = foreignRes.webSocket!;
+    foreignWs.accept();
+
+    // Now connect via PartyServer. onConnect() will call broadcast(), which must
+    // not crash due to the foreign socket.
+    const req = new Request("http://example.com/parties/mixed/room", {
+      headers: { Upgrade: "websocket" }
+    });
+    const res = await worker.fetch(req, env, ctx);
+    const ws = res.webSocket!;
+    ws.accept();
+
+    const { promise, resolve, reject } = Promise.withResolvers<void>();
+    ws.addEventListener("message", (message) => {
+      try {
+        // We should receive at least one message from the server.
+        expect(["hello", "connected"]).toContain(message.data);
+        resolve();
+      } catch (e) {
+        reject(e);
+      } finally {
+        ws.close();
+        foreignWs.close();
+      }
+    });
+
+    return promise;
+  });
+
   // it("can be connected with a query parameter");
   // it("can be connected with a header");
 

packages/partyserver/src/tests/worker.ts

@@ -11,6 +11,7 @@ function assert(condition: unknown, message: string): asserts condition {
 export type Env = {
   Stateful: DurableObjectNamespace<Stateful>;
   OnStartServer: DurableObjectNamespace<OnStartServer>;
+  Mixed: DurableObjectNamespace<Mixed>;
 };
 
 export class Stateful extends Server {
@@ -61,6 +62,37 @@ export class OnStartServer extends Server {
   }
 }
 
+export class Mixed extends Server {
+  static options = {
+    hibernate: true
+  };
+
+  async fetch(request: Request): Promise<Response> {
+    const url = new URL(request.url);
+    if (url.pathname.endsWith("/foreign")) {
+      const room = request.headers.get("x-partykit-room");
+      if (room) {
+        await this.setName(room);
+      }
+
+      const pair = new WebSocketPair();
+      const [client, server] = Object.values(pair);
+      // Accept a hibernated websocket that PartyServer does not manage. This is
+      // equivalent to user code calling `this.ctx.acceptWebSocket()` directly.
+      this.ctx.acceptWebSocket(server, ["foreign"]);
+      return new Response(null, { status: 101, webSocket: client });
+    }
+
+    return super.fetch(request);
+  }
+
+  onConnect(connection: Connection): void {
+    // Trigger a broadcast while a foreign hibernated socket exists.
+    this.broadcast("hello");
+    connection.send("connected");
+  }
+}
+
 export default {
   async fetch(request: Request, env: Env, _ctx: ExecutionContext) {
     return (

packages/partyserver/src/tests/wrangler.jsonc

@@ -18,13 +18,17 @@
       {
         "name": "OnStartServer",
         "class_name": "OnStartServer"
+      },
+      {
+        "name": "Mixed",
+        "class_name": "Mixed"
       }
     ]
   },
   "migrations": [
     {
       "tag": "v1", // Should be unique for each entry
-      "new_classes": ["Stateful", "OnStartServer"]
+      "new_classes": ["Stateful", "OnStartServer", "Mixed"]
     }
   ]
 }

packages/partysocket/src/tests/integration.test.ts

@@ -163,7 +163,7 @@
               for (let i = 0; i < messageCount; i++) {
                 expect(receivedMessages[i]).toBe(`message-${i}`);
               }
-            } catch (e) {
+            } catch (_e) {
               // If we still have Blobs, messages aren't fully processed yet
               return;
             }
@@ -582,7 +582,7 @@
   let wss: WebSocketServer;
 
   beforeAll(() => {
     wss = new WebSocketServer({ port: PORT + 4 });
   });
 
   afterAll(() => {