From 191d84f4f8782570440cd3652bc3311cf70a364c Mon Sep 17 00:00:00 2001
From: "NGOCDG280\\ngocdg" <pnd280@gmail.com>
Date: Sun, 29 Sep 2024 19:16:36 +0700
Subject: [PATCH] feat: handle multiple native Websocket instances

---
 src/content-script/main-world/ws-hook.ts | 11 +++++++++--
 1 file changed, 9 insertions(+), 2 deletions(-)

diff --git a/src/content-script/main-world/ws-hook.ts b/src/content-script/main-world/ws-hook.ts
index 8c3a8aa0..1d3f1a83 100644
--- a/src/content-script/main-world/ws-hook.ts
+++ b/src/content-script/main-world/ws-hook.ts
@@ -8,6 +8,8 @@ import { parseUrl } from "@/utils/utils";
 
 class WsHook {
   private static instance: WsHook | null = null;
+  private capturedInstances: Set<Nullable<WebSocket | XMLHttpRequest>> =
+    new Set();
   private webSocketInstance: Nullable<WebSocket>;
   private longPollingInstance: Nullable<XMLHttpRequest>;
 
@@ -57,6 +59,8 @@ class WsHook {
   setWebSocketInstance(instance: WebSocket): void {
     if (!this.isValidWebSocketInstance(instance)) return;
 
+    this.capturedInstances.add(instance);
+
     this.webSocketInstance = instance;
     this.proxyWebSocketInstance(instance);
 
@@ -347,9 +351,12 @@ class WsHook {
     const self = this;
 
     WebSocket.prototype.send = function (data: any): void {
-      if (!this.url.includes("src=complexity")) {
+      if (
+        !this.url.includes("src=complexity") &&
+        !self.capturedInstances.has(this)
+      ) {
         //! important: must restore the original send method BEFORE capturing the instance
-        WebSocket.prototype.send = self.webSocketOriginalSend;
+        // WebSocket.prototype.send = self.webSocketOriginalSend;
 
         self.setWebSocketInstance(this);
         webpageMessenger.sendMessage({