@signe/room 2.9.4 → 3.0.0

package/src/session.guard.ts

@@ -1,5 +1,9 @@
 import type * as Party from "./types/party";
 
+function getPrivateId(sender: Party.Connection) {
+    return sender.sessionId || sender.id;
+}
+
 /**
  * @description Factory function that creates a session guard with access to room storage
  * @param {Party.Storage} storage - The room storage instance
@@ -29,7 +33,7 @@ export function createRequireSessionGuard(storage: Party.Storage) {
 
         try {
             // Check if session exists in storage
-            const session = await storage.get(`session:${sender.id}`);
+            const session = await storage.get(`session:${getPrivateId(sender)}`);
             
             // Return false if no session found
             if (!session) {
@@ -88,7 +92,7 @@ export const requireSession = async (sender: Party.Connection, value: any, room:
 
     try {
         // Check if session exists in storage
-        const session = await room.storage.get(`session:${sender.id}`);
+        const session = await room.storage.get(`session:${getPrivateId(sender)}`);
         
         // Return false if no session found
         if (!session) {

package/src/shard.ts

@@ -15,18 +15,44 @@ export interface ShardOptions {
 
 export class Shard {
   ws: PartyWebSocket;
-  connectionMap = new Map<string, Party.Connection>(); // Map privateId -> connection
+  connectionMap = new Map<string, Set<Party.Connection>>(); // Map privateId -> active connections
   mainServerStub: any;
   worldUrl: string | null = null;
-  worldId: string = 'default';
+  worldId: string;
   lastReportedConnections: number = 0;
   statsInterval: number = 30000; 
   statsIntervalId: any = null;
 
-  constructor(private room: Party.Room) {}
+  constructor(private room: Party.Room, options: ShardOptions = {}) {
+    this.worldUrl = options.worldUrl ?? null;
+    this.worldId = options.worldId
+      ?? this.getWorldIdFromShardId(room.id)
+      ?? this.getEnvString('WORLD_ID')
+      ?? this.getEnvString('SIGNE_WORLD_ID')
+      ?? 'world-default';
+    this.statsInterval = options.statsInterval ?? this.statsInterval;
+  }
+
+  private getPrivateId(conn: Party.Connection) {
+    return conn.sessionId || conn.id;
+  }
+
+  private getEnvString(key: string) {
+    const value = this.room.env?.[key];
+    return typeof value === 'string' && value.length > 0 ? value : undefined;
+  }
+
+  private getRoomIdFromShardId(shardId: string) {
+    return shardId.split(':')[0];
+  }
+
+  private getWorldIdFromShardId(shardId: string) {
+    const parts = shardId.split(':');
+    return parts.length >= 3 ? parts[1] : undefined;
+  }
 
   async onStart() {
-    const roomId = this.room.id.split(':')[0];
+    const roomId = this.getRoomIdFromShardId(this.room.id);
     const roomStub = this.room.context.parties.main.get(roomId);
     if (!roomStub) {
       console.warn('No room room stub found in main party context');
@@ -36,7 +62,9 @@ export class Shard {
     this.mainServerStub = roomStub;
     this.ws = await roomStub.socket({
       headers: {
-        'x-shard-id': this.room.id
+        'x-shard-id': this.room.id,
+        'x-shard-world-id': this.worldId,
+        'x-access-shard': this.room.env.SHARD_SECRET as string
       }
     }) as unknown as PartyWebSocket;
     
@@ -45,13 +73,25 @@ export class Shard {
       try {
         const message = JSON.parse(event.data);
 
+        if (message.type === 'shard.closeClient' && message.privateId) {
+          const clientConnections = this.connectionMap.get(message.privateId);
+          if (clientConnections?.size) {
+            for (const clientConn of [...clientConnections]) {
+              clientConn.close();
+            }
+          }
+          return;
+        }
+
         // If the message is directed to a specific client, forward it
         if (message.targetClientId) {
-          const clientConn = this.connectionMap.get(message.targetClientId);
-          if (clientConn) {
+          const clientConnections = this.connectionMap.get(message.targetClientId);
+          if (clientConnections?.size) {
             // Remove the routing information before forwarding
             delete message.targetClientId;
-            clientConn.send(message.data);
+            for (const clientConn of clientConnections) {
+              clientConn.send(message.data);
+            }
           }
         } else {
           // Broadcast to all clients if no specific target
@@ -62,12 +102,12 @@ export class Shard {
       }
     });
 
-    await this.updateWorldStats();
+    await this.updateWorldStats(true);
     this.startPeriodicStatsUpdates();
   }
 
   private startPeriodicStatsUpdates() {
-    if (!this.worldUrl) {
+    if (this.statsInterval <= 0 || !this.room.context.parties.world) {
       return;
     }
 
@@ -76,10 +116,11 @@ export class Shard {
     }
 
     this.statsIntervalId = setInterval(() => {
-      this.updateWorldStats().catch(error => {
+      this.updateWorldStats(true).catch(error => {
         console.error('Error in periodic stats update:', error);
       });
     }, this.statsInterval);
+    this.statsIntervalId?.unref?.();
   }
 
   private stopPeriodicStatsUpdates() {
@@ -90,8 +131,12 @@ export class Shard {
   }
 
   onConnect(conn: Party.Connection, ctx: Party.ConnectionContext) {
+    const privateId = this.getPrivateId(conn);
+
     // Store connection mapping
-    this.connectionMap.set(conn.id, conn);
+    const connections = this.connectionMap.get(privateId) ?? new Set<Party.Connection>();
+    connections.add(conn);
+    this.connectionMap.set(privateId, connections);
     
     // Capture all headers and request information
     const headers: Record<string, string> = {};
@@ -111,7 +156,7 @@ export class Shard {
     // Notify the main server about the new connection with complete connection metadata
     this.ws.send(JSON.stringify({
       type: 'shard.clientConnected',
-      privateId: conn.id,
+      privateId,
       requestInfo
     }));
 
@@ -126,7 +171,7 @@ export class Shard {
       // Wrap the original message with sender information
       const wrappedMessage = JSON.stringify({
         type: 'shard.clientMessage',
-        privateId: sender.id,
+        privateId: this.getPrivateId(sender),
         publicId: (sender.state as any)?.publicId,
         payload: parsedMessage
       });
@@ -139,28 +184,48 @@ export class Shard {
   }
 
   onClose(conn: Party.Connection) {
+    const privateId = this.getPrivateId(conn);
+
     // Remove connection from the map
-    this.connectionMap.delete(conn.id);
+    const connections = this.connectionMap.get(privateId);
+    connections?.delete(conn);
+
+    if (connections?.size) {
+      this.updateWorldStats();
+      return;
+    }
+
+    this.connectionMap.delete(privateId);
     
     // Notify main server about disconnection
     this.ws.send(JSON.stringify({
       type: 'shard.clientDisconnected',
-      privateId: conn.id,
+      privateId,
       publicId: (conn.state as any)?.publicId
     }));
 
     this.updateWorldStats();
   }
 
-  async updateWorldStats(): Promise<boolean> {
-    const currentConnections = this.connectionMap.size;
+  async updateWorldStats(force = false): Promise<boolean> {
+    const currentConnections = Array.from(this.connectionMap.values())
+      .reduce((total, connections) => total + connections.size, 0);
     
-    if (currentConnections === this.lastReportedConnections) {
+    if (!force && currentConnections === this.lastReportedConnections) {
       return true;
     }
 
     try {
-      const worldRoom = this.room.context.parties.world.get('world-default');
+      const worldParty = this.room.context.parties.world;
+      if (!worldParty) {
+        return false;
+      }
+
+      const worldRoom = worldParty.get(this.worldId);
+      if (!worldRoom?.fetch) {
+        return false;
+      }
+
       const response = await worldRoom.fetch('/update-shard', {
         method: 'POST',
         headers: {
@@ -169,6 +234,7 @@ export class Shard {
         },
         body: JSON.stringify({
           shardId: this.room.id,
+          worldId: this.worldId,
           connections: currentConnections
         })
       });
@@ -219,7 +285,9 @@ export class Shard {
       
       // Add shard identification
       headers.set('x-shard-id', this.room.id);
+      headers.set('x-shard-world-id', this.worldId);
       headers.set('x-forwarded-by-shard', 'true');
+      headers.set('x-access-shard', this.room.env.SHARD_SECRET as string);
       
       // Client IP tracking for the main server
       const clientIp = req.headers.get('x-forwarded-for') || 'unknown';
@@ -234,7 +302,7 @@ export class Shard {
         body
       };
       // Forward the request to the main server
-      const response = await this.mainServerStub.fetch(path, requestInit);
+      const response = await this.mainServerStub.fetch(path + url.search, requestInit);
       return response;
     } catch (error) {
       return response(500, { error: 'Error forwarding request' });
@@ -247,8 +315,8 @@ export class Shard {
    * @description Executed periodically, used to perform maintenance tasks
    */
   async onAlarm() {
-    await this.updateWorldStats();
+    await this.updateWorldStats(true);
   }
 }
 
-Shard satisfies Party.Worker;
+Shard satisfies Party.Worker;

package/src/storage.ts

@@ -1,15 +1,39 @@
 export class Storage {
   private memory = new Map();
-  async put(key, value) {
-    this.memory.set(key, value);
+  async put(key, value?) {
+    if (typeof key === "string") {
+      this.memory.set(key, value);
+      return;
+    }
+
+    for (const [entryKey, entryValue] of Object.entries(key)) {
+      this.memory.set(entryKey, entryValue);
+    }
   }
   async get(key) {
     return this.memory.get(key);
   }
   async delete(key) {
-    this.memory.delete(key);
+    if (Array.isArray(key)) {
+      let deleted = 0;
+      for (const entryKey of key) {
+        if (this.memory.delete(entryKey)) {
+          deleted += 1;
+        }
+      }
+      return deleted;
+    }
+    return this.memory.delete(key);
   }
-  async list() {
-    return this.memory;
+  async list(options?: { prefix?: string }) {
+    if (!options?.prefix) {
+      return this.memory;
+    }
+
+    return new Map(
+      Array.from(this.memory.entries()).filter(([key]) =>
+        String(key).startsWith(options.prefix!)
+      )
+    );
   }
 }

package/src/testing.ts

@@ -30,6 +30,7 @@ import { Shard } from "./shard"
 export async function testRoom(Room, options: {
     hibernate?: boolean,
     shard?: boolean,
+    id?: string,
     env?: Record<string, string>,
     parties?: Record<string, (io: any) => any>,
     partyFn?: (io: any) => any
@@ -42,7 +43,7 @@ export async function testRoom(Room, options: {
     }
 
     const isShard = options.shard || false
-    const io = new ServerIo(Room.path, isShard ? {
+    const io = new ServerIo(options.id ?? Room.path, isShard ? {
         parties: {
             game: createServer,
             ...(options.parties || {})
@@ -92,7 +93,7 @@ export async function testRoom(Room, options: {
             return client
         },
         getServerUser: async (client: any, prop = 'users') => {
-            const privateId = client.conn.id;
+            const privateId = client.conn.sessionId || client.conn.id;
             const session = await (server as Server).getSession(privateId);
             return (server as any).subRoom[prop]()[session?.publicId];
         }
@@ -114,4 +115,4 @@ export async function request(room: Server | Shard, path: string, options: {
 
 export function tick(ms: number = 0) {
     return new Promise(resolve => setTimeout(resolve, ms))
-}
+}

package/src/types/party.ts

@@ -150,6 +150,9 @@ import type {
   export type Connection<TState = unknown> = WebSocket & {
     /** Connection identifier */
     id: string;
+
+    /** Stable private session identifier shared by reconnects or multiple tabs. */
+    sessionId?: string;
   
     /** @deprecated You can access the socket properties directly on the connection*/
     socket: WebSocket;
@@ -548,4 +551,4 @@ import type {
   export type PartyKitConnection = Connection;
   
   /** @deprecated Use `Party.ServerOptions` instead */
-  export type PartyServerOptions = ServerOptions;
+  export type PartyServerOptions = ServerOptions;

package/src/world.guard.ts

@@ -1,8 +1,7 @@
 import * as Party from "./types/party";
 import { JWTAuth } from "./jwt";
-import { response } from "./utils";
 
-export const guardManageWorld = async (_, req: Party.Request, room: Party.Room): Promise<boolean> => {
+export const guardManageWorld = async (_: unknown, req: Party.Request, room: Party.Room): Promise<boolean> => {
     const tokenShard = req.headers.get("x-access-shard");
     if (tokenShard) {
         if (tokenShard !== room.env.SHARD_SECRET) {
@@ -11,7 +10,7 @@ export const guardManageWorld = async (_, req: Party.Request, room: Party.Room):
         return true
     }
     const url = new URL(req.url);
-    const token = req.headers.get("Authorization") ?? url.searchParams.get("world-auth-token");
+    const token = getAuthToken(req, url);
     if (!token) {
         return false;
     }
@@ -21,8 +20,28 @@ export const guardManageWorld = async (_, req: Party.Request, room: Party.Room):
         if (!payload) {
             return false;
         }
+        if (!canAccessWorld(payload, room.id)) {
+            return false;
+        }
     } catch (error) {
         return false;
     }
     return true;
-}
+}
+
+function getAuthToken(req: Party.Request, url: URL) {
+    const authorization = req.headers.get("Authorization");
+    if (authorization?.startsWith("Bearer ")) {
+        return authorization.slice("Bearer ".length).trim();
+    }
+    return authorization ?? url.searchParams.get("world-auth-token");
+}
+
+function canAccessWorld(payload: Record<string, unknown>, worldId: string) {
+    const worlds = payload.worlds;
+    if (!Array.isArray(worlds)) {
+        return false;
+    }
+
+    return worlds.some((world) => world === "*" || world === worldId);
+}