From f76130df529e65e22d4925a7ac3060d82916debc Mon Sep 17 00:00:00 2001
From: Lucas Willems <lcswillems@gmail.com>
Date: Sun, 24 Nov 2024 20:22:17 +0100
Subject: [PATCH] Expose node URLs

---
 xsuite/src/world/fsworld.test.ts | 14 ++++++++++-
 xsuite/src/world/fsworld.ts      | 41 +++++++++++++++++++++++++-------
 xsuite/src/world/lsworld.ts      |  7 ++++--
 3 files changed, 50 insertions(+), 12 deletions(-)

diff --git a/xsuite/src/world/fsworld.test.ts b/xsuite/src/world/fsworld.test.ts
index f5b50f56..ee31eead 100644
--- a/xsuite/src/world/fsworld.test.ts
+++ b/xsuite/src/world/fsworld.test.ts
@@ -44,9 +44,21 @@ test.concurrent("FSWorld.proxy.blockNonce", async () => {
   });
 });
 
+test.concurrent("FSWorld.start - URLs", async () => {
+  const localhostRegex = /^http:\/\/localhost:\d+$/;
+  using world = await FSWorld.start();
+  expect(world.proxyUrl).toMatch(localhostRegex);
+  expect(world.nodeUrls).toEqual({
+    "0": expect.stringMatching(localhostRegex),
+    "1": expect.stringMatching(localhostRegex),
+    "2": expect.stringMatching(localhostRegex),
+    "4294967295": expect.stringMatching(localhostRegex),
+  });
+});
+
 test.concurrent("FSWorld.start - port 3000", async () => {
   using world = await FSWorld.start({ binaryPort: 3000 });
-  expect(world.proxy.proxyUrl).toEqual("http://localhost:3000");
+  expect(world.proxyUrl).toEqual("http://localhost:3000");
 });
 
 test.concurrent("FSWorld.start - epoch, round, nonce", async () => {
diff --git a/xsuite/src/world/fsworld.ts b/xsuite/src/world/fsworld.ts
index 19cf5656..516cdffb 100644
--- a/xsuite/src/world/fsworld.ts
+++ b/xsuite/src/world/fsworld.ts
@@ -27,22 +27,28 @@ import {
 
 export class FSWorld extends World {
   proxy: FSProxy;
+  proxyUrl: string;
+  nodeUrls: Record<string, string>;
   server?: ChildProcess;
   sysAcc: FSContract;
 
   constructor({
     proxy,
     gasPrice,
+    nodeUrls = {},
     explorerUrl,
     server,
   }: {
     proxy: FSProxy;
     gasPrice: number;
+    nodeUrls?: Record<string, string>;
     explorerUrl?: string;
     server?: ChildProcess;
   }) {
     super({ chainId: "chain", proxy, gasPrice, explorerUrl });
     this.proxy = proxy;
+    this.proxyUrl = this.proxy.proxyUrl;
+    this.nodeUrls = nodeUrls;
     this.server = server;
     this.sysAcc = this.newContract(fullU8AAddress);
   }
@@ -51,10 +57,11 @@ export class FSWorld extends World {
     if (params.chainId !== undefined) {
       throw new Error("chainId is not undefined.");
     }
-    const { proxyUrl, gasPrice, explorerUrl, server } = params;
+    const { proxyUrl, gasPrice, nodeUrls, explorerUrl, server } = params;
     return new FSWorld({
       proxy: new FSProxy({ proxyUrl, explorerUrl }),
       gasPrice: gasPrice ?? 1_000_000_000,
+      nodeUrls,
       explorerUrl,
       server,
     });
@@ -80,8 +87,8 @@ export class FSWorld extends World {
     gasPrice?: number;
     explorerUrl?: string;
   } & ProxyParams = {}): Promise<FSWorld> {
-    const { server, proxyUrl } = await startProxy(proxyParams);
-    return this.new({ proxyUrl, gasPrice, explorerUrl, server });
+    const { proxyUrl, nodeUrls, server } = await startProxy(proxyParams);
+    return this.new({ proxyUrl, nodeUrls, gasPrice, explorerUrl, server });
   }
 
   async restartProxy(proxyParams: ProxyParams = {}) {
@@ -368,18 +375,33 @@ const startProxy = async ({
     throw error;
   });
 
+  const nodeUrls: Record<string, string> = {};
+
   const proxyUrl = await new Promise<string>((resolve) => {
-    server.stdout.on("data", (data: Buffer) => {
-      const addressRegex =
+    const onData = (data: Buffer) => {
+      const dataStr = data.toString();
+
+      const nodeUrlRegex =
+        /node status +.{4,7}address.{4,7} = (http:\/\/[^\s]+) .{4,7}shard.{4,7} = (\d+)/g;
+      for (const match of [...dataStr.matchAll(nodeUrlRegex)]) {
+        const [, url, shard] = match;
+        nodeUrls[shard] = url;
+      }
+
+      const proxyHostRegex =
         /chain simulator's is accessible through the URL ([\w\d.:]+)/;
-      const match = data.toString().match(addressRegex);
+      const match = dataStr.match(proxyHostRegex);
       if (match) {
-        resolve(`http://${match[1]}`);
+        const [, host] = match;
+        resolve(`http://${host}`);
+        server.stdout.off("data", onData);
       }
-    });
+    };
+
+    server.stdout.on("data", onData);
   });
 
-  return { proxyUrl, server };
+  return { proxyUrl, nodeUrls, server };
 };
 
 type FSWorldNewParams =
@@ -387,6 +409,7 @@ type FSWorldNewParams =
       chainId?: undefined;
       proxyUrl: string;
       gasPrice?: number;
+      nodeUrls?: Record<string, string>;
       explorerUrl?: string;
       server?: ChildProcess;
     }
diff --git a/xsuite/src/world/lsworld.ts b/xsuite/src/world/lsworld.ts
index 88914501..18902a0d 100644
--- a/xsuite/src/world/lsworld.ts
+++ b/xsuite/src/world/lsworld.ts
@@ -98,13 +98,16 @@ export class LSWorld extends World {
     });
 
     const proxyUrl = await new Promise<string>((resolve) => {
-      server.stdout.on("data", (data: Buffer) => {
+      const onData = (data: Buffer) => {
         const addressRegex = /Server running on (http:\/\/[\w\d.:]+)/;
         const match = data.toString().match(addressRegex);
         if (match) {
           resolve(match[1]);
+          server.stdout.off("data", onData);
         }
-      });
+      };
+
+      server.stdout.on("data", onData);
     });
 
     return this.new({ proxyUrl, gasPrice, explorerUrl, server });