From e3dec45ed808c972356e9b9de1cb161e599824fc Mon Sep 17 00:00:00 2001
From: Logan Nguyen <logan.nguyen@swirldslabs.com>
Date: Sat, 4 May 2024 12:29:10 -0400
Subject: [PATCH] fix: fixed setting WS_NEW_HEADS_ENABLED true as default logic
 at Relay level

Signed-off-by: Logan Nguyen <logan.nguyen@swirldslabs.com>
---
 packages/relay/src/lib/poller.ts              |  3 +-
 .../src/controllers/eth_subscribe.ts          |  4 +--
 .../acceptance/subscribeNewHeads.spec.ts      | 30 +++++++++++++++++++
 3 files changed, 34 insertions(+), 3 deletions(-)

diff --git a/packages/relay/src/lib/poller.ts b/packages/relay/src/lib/poller.ts
index 15fa1eb54b..1ce8c63084 100644
--- a/packages/relay/src/lib/poller.ts
+++ b/packages/relay/src/lib/poller.ts
@@ -48,7 +48,8 @@ export class Poller {
     this.logger = logger;
     this.polls = [];
     this.pollingInterval = Number(process.env.WS_POLLING_INTERVAL) || 500;
-    this.newHeadsEnabled = process.env.WS_NEW_HEADS_ENABLED ? Boolean(Number(process.env.WS_NEW_HEADS_ENABLED)) : true;
+    this.newHeadsEnabled =
+      typeof process.env.WS_NEW_HEADS_ENABLED !== 'undefined' ? process.env.WS_NEW_HEADS_ENABLED === 'true' : true;
 
     const activePollsGaugeName = 'rpc_websocket_active_polls';
     register.removeSingleMetric(activePollsGaugeName);
diff --git a/packages/ws-server/src/controllers/eth_subscribe.ts b/packages/ws-server/src/controllers/eth_subscribe.ts
index 7918e65771..24624e8483 100644
--- a/packages/ws-server/src/controllers/eth_subscribe.ts
+++ b/packages/ws-server/src/controllers/eth_subscribe.ts
@@ -77,9 +77,9 @@ const handleEthSubscribeNewHeads = (
   requestIdPrefix: string,
 ): { response: any; subscriptionId: any } => {
   const wsNewHeadsEnabled =
-    typeof process.env.WS_NEW_HEADS_ENABLED !== 'undefined' ? process.env.WS_NEW_HEADS_ENABLED : 'true';
+    typeof process.env.WS_NEW_HEADS_ENABLED !== 'undefined' ? process.env.WS_NEW_HEADS_ENABLED === 'true' : true;
 
-  if (wsNewHeadsEnabled === 'true') {
+  if (wsNewHeadsEnabled) {
     ({ response, subscriptionId } = subscribeToNewHeads(filters, response, subscriptionId, ctx, event, relay, logger));
   } else {
     logger.warn(
diff --git a/packages/ws-server/tests/acceptance/subscribeNewHeads.spec.ts b/packages/ws-server/tests/acceptance/subscribeNewHeads.spec.ts
index 8a18abc57c..7111c2589f 100644
--- a/packages/ws-server/tests/acceptance/subscribeNewHeads.spec.ts
+++ b/packages/ws-server/tests/acceptance/subscribeNewHeads.spec.ts
@@ -198,6 +198,36 @@ describe('@release @web-socket-batch-3 eth_subscribe newHeads', async function (
 
       await new Promise((resolve) => setTimeout(resolve, 500));
     });
+
+    it('should subscribe to newHeads even when WS_NEW_HEADS_ENABLED=undefined, and receive a valid JSON RPC response', async (done) => {
+      expect(process.env.WS_NEW_HEADS_ENABLED).to.eq('undefined');
+
+      const webSocket = new WebSocket(WS_RELAY_URL);
+      const subscriptionId = 1;
+      webSocket.on('open', function open() {
+        webSocket.send(
+          JSON.stringify({
+            id: subscriptionId,
+            jsonrpc: '2.0',
+            method: 'eth_subscribe',
+            params: ['newHeads', { includeTransactions: true }],
+          }),
+        );
+      });
+
+      let responseCounter = 0;
+
+      Utils.sendTransaction(ONE_TINYBAR, CHAIN_ID, accounts, rpcServer, requestId, mirrorNodeServer);
+      webSocket.on('message', function incoming(data) {
+        const response = JSON.parse(data);
+        responseCounter++;
+        verifyResponse(response, done, webSocket, true);
+        if (responseCounter > 1) {
+          webSocket.close();
+        }
+      });
+      done();
+    });
   });
 
   describe('Subscriptions for newHeads', async function () {