chore: validate socket.io subscription topics (#1144)

* feat: added validation check for subscriptions

* fix: resolved export error

* refactor: added validation in middleware

* refactor: updated code according to suggestions

* fix: updated code to return multiple invalid topics upon subscription

* style: added space before comma-separated invalid topics

Author: MHassanTariq

src/api/query-helpers.ts

@@ -1,6 +1,6 @@
 import { ClarityAbi } from '@stacks/transactions';
 import { NextFunction, Request, Response } from 'express';
-import { has0xPrefix, hexToBuffer, isValidPrincipal, parseEventTypeStrings } from './../helpers';
+import { has0xPrefix, hexToBuffer, parseEventTypeStrings, isValidPrincipal } from './../helpers';
 import { InvalidRequestError, InvalidRequestErrorType } from '../errors';
 import { DbEventTypeId } from './../datastore/common';
 
@@ -286,3 +286,11 @@ export function parseEventTypeFilter(
 
   return eventTypeFilter;
 }
+export function isValidTxId(tx_id: string) {
+  try {
+    validateRequestHexInput(tx_id);
+    return true;
+  } catch {
+    return false;
+  }
+}

src/api/routes/ws/socket-io.ts

@@ -17,8 +17,42 @@ import {
   getTxFromDataStore,
   parseDbTx,
 } from '../../controllers/db-controller';
-import { isProdEnv, logError, logger } from '../../../helpers';
+import { isProdEnv, isValidPrincipal, logError, logger } from '../../../helpers';
 import { WebSocketPrometheus } from './metrics';
+import { isValidTxId } from '../../../api/query-helpers';
+
+function getInvalidSubscriptionTopics(subscriptions: Topic | Topic[]): undefined | string[] {
+  const isSubValid = (sub: Topic): undefined | string => {
+    if (sub.includes(':')) {
+      const txOrAddr = sub.split(':')[0];
+      const value = sub.split(':')[1];
+      switch (txOrAddr) {
+        case 'address-transaction':
+        case 'address-stx-balance':
+          return isValidPrincipal(value) ? undefined : sub;
+        case 'transaction':
+          return isValidTxId(value) ? undefined : sub;
+        default:
+          return sub;
+      }
+    }
+    switch (sub) {
+      case 'block':
+      case 'mempool':
+      case 'microblock':
+        return undefined;
+      default:
+        return sub;
+    }
+  };
+  if (!Array.isArray(subscriptions)) {
+    const invalidSub = isSubValid(subscriptions);
+    return invalidSub ? [invalidSub] : undefined;
+  }
+  const validatedSubs = subscriptions.map(isSubValid);
+  const invalidSubs = validatedSubs.filter(validSub => typeof validSub === 'string');
+  return invalidSubs.length === 0 ? undefined : (invalidSubs as string[]);
+}
 
 export function createSocketIORouter(db: DataStore, server: http.Server) {
   const io = new SocketIOServer<ClientToServerMessages, ServerToClientMessages>(server, {
@@ -38,7 +72,6 @@ export function createSocketIORouter(db: DataStore, server: http.Server) {
     }
     const subscriptions = socket.handshake.query['subscriptions'];
     if (subscriptions) {
-      // TODO: check if init topics are valid, reject connection with error if not
       const topics = [...[subscriptions]].flat().flatMap(r => r.split(','));
       for (const topic of topics) {
         prometheus?.subscribe(socket, topic);
@@ -51,10 +84,11 @@ export function createSocketIORouter(db: DataStore, server: http.Server) {
       prometheus?.disconnect(socket);
     });
     socket.on('subscribe', async (topic, callback) => {
-      prometheus?.subscribe(socket, topic);
-      await socket.join(topic);
-      // TODO: check if topic is valid, and return error message if not
-      callback?.(null);
+      if (!getInvalidSubscriptionTopics(topic)) {
+        prometheus?.subscribe(socket, topic);
+        await socket.join(topic);
+        callback?.(null);
+      }
     });
     socket.on('unsubscribe', async (...topics) => {
       for (const topic of topics) {
@@ -64,6 +98,23 @@ export function createSocketIORouter(db: DataStore, server: http.Server) {
     });
   });
 
+  // Middleware checks for the invalid topic subscriptions and terminates connection if found any
+  io.use((socket, next) => {
+    const subscriptions = socket.handshake.query['subscriptions'];
+    if (subscriptions) {
+      const topics = [...[subscriptions]].flat().flatMap(r => r.split(','));
+      const invalidSubs = getInvalidSubscriptionTopics(topics as Topic[]);
+      if (invalidSubs) {
+        const error = new Error(`Invalid topic: ${invalidSubs.join(', ')}`);
+        next(error);
+      } else {
+        next();
+      }
+    } else {
+      next();
+    }
+  });
+
   const adapter = io.of('/').adapter;
 
   adapter.on('create-room', room => {

src/tests/socket-io-tests.ts

@@ -224,6 +224,83 @@ describe('socket-io', () => {
     }
   });
 
+  test('socket-io > invalid topic connection', async () => {
+    const faultyAddr = 'faulty address';
+    const address = apiServer.address;
+    const socket = io(`http://${address}`, {
+      reconnection: false,
+      query: { subscriptions: `address-stx-balance:${faultyAddr}` },
+    });
+    const updateWaiter: Waiter<Error> = waiter();
+
+    socket.on(`connect_error`, err => {
+      updateWaiter.finish(err);
+    });
+
+    const result = await updateWaiter;
+    try {
+      throw result;
+    } catch (err: any) {
+      expect(err.message).toEqual(`Invalid topic: address-stx-balance:${faultyAddr}`);
+    } finally {
+      socket.close();
+    }
+  });
+
+  test('socket-io > multiple invalid topic connection', async () => {
+    const faultyAddrStx = 'address-stx-balance:faulty address';
+    const faultyTx = 'transaction:0x1';
+    const address = apiServer.address;
+    const socket = io(`http://${address}`, {
+      reconnection: false,
+      query: { subscriptions: `${faultyAddrStx},${faultyTx}` },
+    });
+    const updateWaiter: Waiter<Error> = waiter();
+
+    socket.on(`connect_error`, err => {
+      updateWaiter.finish(err);
+    });
+
+    const result = await updateWaiter;
+    try {
+      throw result;
+    } catch (err: any) {
+      expect(err.message).toEqual(`Invalid topic: ${faultyAddrStx}, ${faultyTx}`);
+    } finally {
+      socket.close();
+    }
+  });
+
+  test('socket-io > valid socket subscription', async () => {
+    const address = apiServer.address;
+    const socket = io(`http://${address}`, {
+      reconnection: false,
+      query: { subscriptions: '' },
+    });
+    const updateWaiter: Waiter<Block> = waiter();
+
+    socket.on('block', block => {
+      updateWaiter.finish(block);
+    });
+
+    socket.emit('subscribe', 'block');
+
+    const block = new TestBlockBuilder({ block_hash: '0x1234', burn_block_hash: '0x5454' })
+      .addTx({ tx_id: '0x4321' })
+      .build();
+    await db.update(block);
+
+    const result = await updateWaiter;
+    try {
+      expect(result.hash).toEqual('0x1234');
+      expect(result.burn_block_hash).toEqual('0x5454');
+      expect(result.txs[0]).toEqual('0x4321');
+    } finally {
+      socket.emit('unsubscribe', 'block');
+      socket.close();
+    }
+  });
+
   afterEach(async () => {
     await apiServer.terminate();
     dbClient.release();