@signe/room 1.4.2 → 2.0.0

package/src/shard.ts

@@ -0,0 +1,244 @@
+import type * as Party from "./types/party";
+import { response } from "./utils";
+
+interface PartyWebSocket {
+  send: (data: string | ArrayBufferLike | Blob | ArrayBufferView) => void;
+  addEventListener: (type: string, listener: (event: any) => void) => void;
+  close: () => void;
+}
+
+export interface ShardOptions {
+  worldUrl?: string;
+  worldId?: string;
+  statsInterval?: number;
+}
+
+export class Shard {
+  ws: PartyWebSocket;
+  connectionMap = new Map<string, Party.Connection>(); // Map privateId -> connection
+  mainServerStub: any;
+  worldUrl: string | null = null;
+  worldId: string = 'default';
+  lastReportedConnections: number = 0;
+  statsInterval: number = 30000; 
+  statsIntervalId: any = null;
+
+  constructor(private room: Party.Room) {}
+
+  async onStart() {
+    const roomId = this.room.id.split(':')[0];
+    const roomStub = this.room.context.parties.main.get(roomId);
+    if (!roomStub) {
+      console.warn('No room room stub found in main party context');
+      return;
+    }
+    
+    this.mainServerStub = roomStub;
+    this.ws = await roomStub.socket({
+      headers: {
+        'x-shard-id': this.room.id
+      }
+    }) as unknown as PartyWebSocket;
+    
+    // Handle messages from the main server
+    this.ws.addEventListener("message", (event) => {
+      try {
+        const message = JSON.parse(event.data);
+
+        // If the message is directed to a specific client, forward it
+        if (message.targetClientId) {
+          const clientConn = this.connectionMap.get(message.targetClientId);
+          if (clientConn) {
+            // Remove the routing information before forwarding
+            delete message.targetClientId;
+            clientConn.send(message.data);
+          }
+        } else {
+          // Broadcast to all clients if no specific target
+          this.room.broadcast(event.data);
+        }
+      } catch (error) {
+        console.error("Error processing message from main server:", error);
+      }
+    });
+
+    await this.updateWorldStats();
+    this.startPeriodicStatsUpdates();
+  }
+
+  private startPeriodicStatsUpdates() {
+    if (!this.worldUrl) {
+      return;
+    }
+
+    if (this.statsIntervalId) {
+      clearInterval(this.statsIntervalId);
+    }
+
+    this.statsIntervalId = setInterval(() => {
+      this.updateWorldStats().catch(error => {
+        console.error('Error in periodic stats update:', error);
+      });
+    }, this.statsInterval);
+  }
+
+  private stopPeriodicStatsUpdates() {
+    if (this.statsIntervalId) {
+      clearInterval(this.statsIntervalId);
+      this.statsIntervalId = null;
+    }
+  }
+
+  onConnect(conn: Party.Connection, ctx: Party.ConnectionContext) {
+    // Store connection mapping
+    this.connectionMap.set(conn.id, conn);
+    
+    // Notify the main server about the new connection with connection metadata
+    this.ws.send(JSON.stringify({
+      type: 'shard.clientConnected',
+      privateId: conn.id,
+      connectionInfo: {
+        ip: ctx.request?.headers.get('x-forwarded-for') || 'unknown',
+        userAgent: ctx.request?.headers.get('user-agent') || 'unknown',
+        // Add any other relevant connection info
+      }
+    }));
+
+    this.updateWorldStats();
+  }
+
+  onMessage(message: string | ArrayBuffer | ArrayBufferView, sender: Party.Connection) {
+    try {
+      // Parse the message if it's a string
+      const parsedMessage = typeof message === 'string' ? JSON.parse(message) : message;
+      
+      // Wrap the original message with sender information
+      const wrappedMessage = JSON.stringify({
+        type: 'shard.clientMessage',
+        privateId: sender.id,
+        publicId: (sender.state as any)?.publicId,
+        payload: parsedMessage
+      });
+      
+      // Forward to main server
+      this.ws.send(wrappedMessage);
+    } catch (error) {
+      console.error("Error forwarding message to main server:", error);
+    }
+  }
+
+  onClose(conn: Party.Connection) {
+    // Remove connection from the map
+    this.connectionMap.delete(conn.id);
+    
+    // Notify main server about disconnection
+    this.ws.send(JSON.stringify({
+      type: 'shard.clientDisconnected',
+      privateId: conn.id,
+      publicId: (conn.state as any)?.publicId
+    }));
+
+    this.updateWorldStats();
+  }
+
+  async updateWorldStats(): Promise<boolean> {
+    const currentConnections = this.connectionMap.size;
+    
+    if (currentConnections === this.lastReportedConnections) {
+      return true;
+    }
+
+    try {
+      const worldRoom = this.room.context.parties.world.get('world-default');
+      const response = await worldRoom.fetch('/update-shard', {
+        method: 'POST',
+        headers: {
+          'Content-Type': 'application/json',
+          'x-access-shard': this.room.env.SHARD_SECRET as string
+        },
+        body: JSON.stringify({
+          shardId: this.room.id,
+          connections: currentConnections
+        })
+      });
+
+      if (!response.ok) {
+        const errorData = await response.json().catch(() => ({ error: 'Unknown error' }));
+        console.error(`Failed to update World stats: ${response.status} - ${errorData.error || 'Unknown error'}`);
+        return false;
+      }
+
+      // Mettre à jour le dernier nombre rapporté
+      this.lastReportedConnections = currentConnections;
+      return true;
+    } catch (error) {
+      console.error('Error updating World stats:', error);
+      return false;
+    }
+  }
+
+  /**
+   * @method onRequest
+   * @async
+   * @param {Party.Request} req - The HTTP request to handle
+   * @description Forwards HTTP requests to the main server, preserving client context
+   * @returns {Promise<Response>} The response from the main server
+   */
+  async onRequest(req: Party.Request): Promise<Response> {
+    if (!this.mainServerStub) {
+      return response(503, { error: 'Shard not connected to main server' });
+    }
+
+    try {
+      // Extract necessary information from the request
+      const url = new URL(req.url);
+      const path = url.pathname;
+      const method = req.method;
+      let body: string | null = null;
+      
+      if (method !== 'GET' && method !== 'HEAD') {
+        body = await req.text();
+      }
+      
+      // Convert original headers to a new Headers object
+      const headers = new Headers();
+      req.headers.forEach((value, key) => {
+        headers.set(key, value);
+      });
+      
+      // Add shard identification
+      headers.set('x-shard-id', this.room.id);
+      headers.set('x-forwarded-by-shard', 'true');
+      
+      // Client IP tracking for the main server
+      const clientIp = req.headers.get('x-forwarded-for') || 'unknown';
+      if (clientIp) {
+        headers.set('x-original-client-ip', clientIp);
+      }
+      
+      // Prepare request options
+      const requestInit: RequestInit = {
+        method,
+        headers,
+        body
+      };
+
+      // Forward the request to the main server
+      const response = await this.mainServerStub.fetch(path, requestInit);
+      return response;
+    } catch (error) {
+      return response(500, { error: 'Error forwarding request' });
+    }
+  }
+  
+  /**
+   * @method onAlarm
+   * @async
+   * @description Executed periodically, used to perform maintenance tasks
+   */
+  async onAlarm() {
+    await this.updateWorldStats();
+  }
+}
+
+Shard satisfies Party.Worker;

package/src/testing.ts

@@ -1,5 +1,6 @@
 import { ServerIo } from "./mock"
 import { Server } from "./server"
+import { Shard } from "./shard"
 
 /**
  * @description Test the room with a mock server and client
@@ -27,21 +28,61 @@ import { Server } from "./server"
  * @returns The server, room, and createClient function
  */
 export async function testRoom(Room, options: {
-    hibernate?: boolean
+    hibernate?: boolean,
+    shard?: boolean,
+    env?: Record<string, string>
 } = {}) {
-    const io = new ServerIo(Room.path)
+
+    const createServer = (io: any) => {
+        const server = new Server(io)
+        server.rooms = [Room]
+        return server
+    }
+
+    const isShard = options.shard || false
+    const io = new ServerIo(Room.path, isShard ? {
+        parties: {
+            game: createServer
+        },
+        env: options.env
+    } : {
+        env: options.env
+    })
     Room.prototype.throttleSync = 0
     Room.prototype.throttleStorage = 0
     Room.prototype.options = options
-    const server = new Server(io as any)
-    server.rooms = [Room]
+    
+    let server: Server | Shard;
+    if (options.shard) {
+        const shardServer = new Shard(io as any);
+        // Add subRoom property to Shard for compatibility with Server
+        (shardServer as any).subRoom = null;
+        server = shardServer;
+    } else {
+        server = await createServer(io as any);
+    }
+    
     await server.onStart()
+    
     return {
         server,
-        room: server.subRoom,
+        room: (server as any).subRoom,
         createClient: async () => {
-            const client = await io.connection(server)
+            const client = await io.connection(server as Server)
             return client
         }
     }
 }
+
+export async function request(room: Server, path: string, options: {
+    method: 'GET' | 'POST' | 'PUT' | 'DELETE',
+    body?: any,
+    headers?: Record<string, string>
+} = {
+    method: 'GET',
+}) {
+    const url = new URL('http://localhost' + path)
+    const request = new Request(url.toString(), options)
+    const response = await room.onRequest(request as any)
+    return response
+}

package/src/utils.ts

@@ -217,4 +217,11 @@ export function buildObject(valuesMap: Map<string, any>, allMemory: Record<strin
     }
   }
   return memoryObj;
+}
+
+export function response(status: number, body: any): Response {
+  return new Response(JSON.stringify(body), { 
+    status,
+    headers: { 'Content-Type': 'application/json' }
+  });
 }

package/src/world.guard.ts

@@ -0,0 +1,28 @@
+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> => {
+    const tokenShard = req.headers.get("x-access-shard");
+    if (tokenShard) {
+        if (tokenShard !== room.env.SHARD_SECRET) {
+            return false
+        }
+        return true
+    }
+    const url = new URL(req.url);
+    const token = req.headers.get("Authorization") ?? url.searchParams.get("world-auth-token");
+    if (!token) {
+        return false;
+    }
+    const jwt = new JWTAuth(room.env.AUTH_JWT_SECRET as string);
+    try {
+        const payload = await jwt.verify(token);
+        if (!payload) {
+            return false;
+        }
+    } catch (error) {
+        return false;
+    }
+    return true;
+}