feat: add websocket support for c2 detection

app/manifest/v2/_base.json

@@ -66,6 +66,8 @@
     "clipboardWrite",
     "http://*/*",
     "https://*/*",
+    "ws://*/*",
+    "wss://*/*",
     "activeTab",
     "webRequest",
     "webRequestBlocking",

app/manifest/v3/_base.json

@@ -50,7 +50,9 @@
     "http://localhost:8545/",
     "file://*/*",
     "http://*/*",
-    "https://*/*"
+    "https://*/*",
+    "ws://*/*",
+    "wss://*/*"
   ],
   "icons": {
     "16": "images/icon-16.png",

app/scripts/background.js

@@ -308,22 +308,16 @@ function maybeDetectPhishing(theController) {
       // blocking is better than tab redirection, as blocking will prevent
       // the browser from loading the page at all
       if (isManifestV2) {
-        if (details.type === 'sub_frame') {
-          // redirect the entire tab to the
-          // phishing warning page instead.
-          redirectTab(details.tabId, redirectHref);
-          // don't let the sub_frame load at all
-          return { cancel: true };
-        }
-        // redirect the whole tab
-        return { redirectUrl: redirectHref };
+        // redirect the whole tab (even if it's a sub_frame request)
+        redirectTab(details.tabId, redirectHref);
+        return { cancel: true };
       }
       // redirect the whole tab (even if it's a sub_frame request)
       redirectTab(details.tabId, redirectHref);
       return {};
     },
     {
-      urls: ['http://*/*', 'https://*/*'],
+      urls: ['http://*/*', 'https://*/*', 'ws://*/*', 'wss://*/*'],
     },
     isManifestV2 ? ['blocking'] : [],
   );

privacy-snapshot.json

@@ -72,5 +72,6 @@
   "unresponsive-rpc.test",
   "unresponsive-rpc.url",
   "user-storage.api.cx.metamask.io",
-  "www.4byte.directory"
+  "www.4byte.directory",
+  "verify.walletconnect.com"
 ]

test/e2e/helpers.js

@@ -4,6 +4,7 @@ const BigNumber = require('bignumber.js');
 const mockttp = require('mockttp');
 const detectPort = require('detect-port');
 const { difference } = require('lodash');
+const WebSocket = require('ws');
 const createStaticServer = require('../../development/create-static-server');
 const { setupMocking } = require('./mock-e2e');
 const { Ganache } = require('./seeder/ganache');
@@ -640,6 +641,43 @@ async function unlockWallet(
   }
 }
 
+/**
+ * Simulates a WebSocket connection by executing a script in the browser context.
+ *
+ * @param {WebDriver} driver - The WebDriver instance.
+ * @param {string} hostname - The hostname to connect to.
+ */
+async function createWebSocketConnection(driver, hostname) {
+  try {
+    await driver.executeScript(async (wsHostname) => {
+      const url = `ws://${wsHostname}:8000`;
+      const socket = new WebSocket(url);
+      socket.onopen = () => {
+        console.log('WebSocket connection opened');
+        socket.send('Hello, server!');
+      };
+      socket.onerror = (error) => {
+        console.error(
+          'WebSocket error:',
+          error.message || 'Connection blocked',
+        );
+      };
+      socket.onmessage = (event) => {
+        console.log('Message received from server:', event.data);
+      };
+      socket.onclose = () => {
+        console.log('WebSocket connection closed');
+      };
+    }, hostname);
+  } catch (error) {
+    console.error(
+      `Failed to execute WebSocket connection script for ws://${hostname}:8081`,
+      error,
+    );
+    throw error;
+  }
+}
+
 const logInWithBalanceValidation = async (driver, ganacheServer) => {
   await unlockWallet(driver);
   // Wait for balance to load
@@ -975,4 +1013,5 @@ module.exports = {
   tempToggleSettingRedesignedTransactionConfirmations,
   openMenuSafe,
   sentryRegEx,
+  createWebSocketConnection,
 };

test/e2e/tests/phishing-controller/mocks.js

@@ -10,7 +10,9 @@ const {
 const lastUpdated = 1;
 const defaultHotlist = { data: [] };
 const defaultC2DomainBlocklist = {
-  recentlyAdded: [],
+  recentlyAdded: [
+    '33c8e026e76cea2df82322428554c932961cd80080fa379454350d7f13371f36', // hash for malicious.localhost
+  ],
   recentlyRemoved: [],
   lastFetchedAt: '2024-08-27T15:30:45Z',
 };
@@ -95,15 +97,12 @@ async function setupPhishingDetectionMocks(
       };
     });
 
-  await mockServer
-    .forGet(C2_DOMAIN_BLOCKLIST_URL)
-    .withQuery({ timestamp: '2024-08-27T15:30:45Z' })
-    .thenCallback(() => {
-      return {
-        statusCode: 200,
-        json: defaultC2DomainBlocklist,
-      };
-    });
+  await mockServer.forGet(C2_DOMAIN_BLOCKLIST_URL).thenCallback(() => {
+    return {
+      statusCode: 200,
+      json: defaultC2DomainBlocklist,
+    };
+  });
 
   await mockServer
     .forGet('https://github.com/MetaMask/eth-phishing-detect/issues/new')

test/e2e/tests/phishing-controller/phishing-detection.spec.js

@@ -9,6 +9,7 @@ const {
   openDapp,
   unlockWallet,
   WINDOW_TITLES,
+  createWebSocketConnection,
 } = require('../../helpers');
 const FixtureBuilder = require('../../fixture-builder');
 const {
@@ -307,10 +308,82 @@ describe('Phishing Detection', function () {
           text: 'Back to safety',
         });
 
+        await driver.waitForUrl({
+          url: `https://portfolio.metamask.io/?metamaskEntry=phishing_page_portfolio_button`,
+        });
+      },
+    );
+  });
+
+  it('should block a website that makes a websocket connection to a malicious command and control server', async function () {
+    const testPageURL = 'http://localhost:8080';
+    await withFixtures(
+      {
+        fixtures: new FixtureBuilder().build(),
+        ganacheOptions: defaultGanacheOptions,
+        title: this.test.fullTitle(),
+        testSpecificMock: async (mockServer) => {
+          await mockServer.forAnyWebSocket().thenEcho();
+          await setupPhishingDetectionMocks(mockServer, {
+            blockProvider: BlockProvider.MetaMask,
+          });
+        },
+        dapp: true,
+      },
+      async ({ driver }) => {
+        await unlockWallet(driver);
+
+        await driver.openNewPage(testPageURL);
+
+        await createWebSocketConnection(driver, 'malicious.localhost');
+
+        await driver.switchToWindowWithTitle(
+          'MetaMask Phishing Detection',
+          10000,
+        );
+
+        await driver.waitForSelector({
+          testId: 'unsafe-continue-loaded',
+        });
+
+        await driver.clickElement({
+          text: 'Back to safety',
+        });
+
+        await driver.waitForUrl({
+          url: `https://portfolio.metamask.io/?metamaskEntry=phishing_page_portfolio_button`,
+        });
+      },
+    );
+  });
+
+  it('should not block a website that makes a safe WebSocket connection', async function () {
+    const testPageURL = 'http://localhost:8080/';
+    await withFixtures(
+      {
+        fixtures: new FixtureBuilder().build(),
+        ganacheOptions: defaultGanacheOptions,
+        title: this.test.fullTitle(),
+        testSpecificMock: async (mockServer) => {
+          await mockServer.forAnyWebSocket().thenEcho();
+          await setupPhishingDetectionMocks(mockServer, {
+            blockProvider: BlockProvider.MetaMask,
+          });
+        },
+        dapp: true,
+      },
+      async ({ driver }) => {
+        await unlockWallet(driver);
+
+        await driver.openNewPage(testPageURL);
+
+        await createWebSocketConnection(driver, 'safe.localhost');
+
+        await driver.wait(until.titleIs(WINDOW_TITLES.TestDApp), 10000);
+
         const currentUrl = await driver.getCurrentUrl();
-        const expectedPortfolioUrl = `https://portfolio.metamask.io/?metamaskEntry=phishing_page_portfolio_button`;
 
-        assert.equal(currentUrl, expectedPortfolioUrl);
+        assert.equal(currentUrl, testPageURL);
       },
     );
   });