fix: Improve WS request sanitization (#10231)

Author: ovr

packages/cubejs-api-gateway/package.json

@@ -49,7 +49,8 @@
     "nexus": "^1.1.0",
     "node-fetch": "^2.6.1",
     "ramda": "^0.27.0",
-    "uuid": "^8.3.2"
+    "uuid": "^8.3.2",
+    "zod": "^4.1.13"
   },
   "devDependencies": {
     "@cubejs-backend/linter": "1.5.12",

packages/cubejs-api-gateway/src/index.ts

@@ -3,5 +3,8 @@ export * from './sql-server';
 export * from './interfaces';
 export * from './cubejs-handler-error';
 export * from './user-error';
+
 export { getRequestIdFromRequest } from './request-parser';
 export { TransformDataRequest } from './types/responses';
+
+export type { SubscriptionServer } from './ws';

packages/cubejs-api-gateway/src/ws/index.ts

@@ -0,0 +1,3 @@
+export * from './local-subscription-store';
+export * from './message-schema';
+export * from './subscription-server';

packages/cubejs-api-gateway/src/ws/local-subscription-store.ts

@@ -2,8 +2,19 @@ interface LocalSubscriptionStoreOptions {
   heartBeatInterval?: number;
 }
 
+export type LocalSubscriptionStoreSubscription = {
+  message: any,
+  state: any,
+  timestamp: Date,
+};
+
+export type LocalSubscriptionStoreConnection = {
+  subscriptions: Map<string, LocalSubscriptionStoreSubscription>,
+  authContext?: any,
+};
+
 export class LocalSubscriptionStore {
-  protected connections = {};
+  protected readonly connections: Map<string, LocalSubscriptionStoreConnection> = new Map();
 
   protected readonly hearBeatInterval: number;
 
@@ -12,60 +23,68 @@ export class LocalSubscriptionStore {
   }
 
   public async getSubscription(connectionId: string, subscriptionId: string) {
-    const connection = this.getConnection(connectionId);
-    return connection.subscriptions[subscriptionId];
+    const connection = this.getConnectionOrCreate(connectionId);
+    return connection.subscriptions.get(subscriptionId);
   }
 
   public async subscribe(connectionId: string, subscriptionId: string, subscription) {
-    const connection = this.getConnection(connectionId);
-    connection.subscriptions[subscriptionId] = {
+    const connection = this.getConnectionOrCreate(connectionId);
+    connection.subscriptions.set(subscriptionId, {
       ...subscription,
       timestamp: new Date()
-    };
+    });
   }
 
   public async unsubscribe(connectionId: string, subscriptionId: string) {
-    const connection = this.getConnection(connectionId);
-    delete connection.subscriptions[subscriptionId];
+    const connection = this.getConnectionOrCreate(connectionId);
+    connection.subscriptions.delete(subscriptionId);
   }
 
-  public async getAllSubscriptions() {
-    return Object.keys(this.connections).map(connectionId => {
-      Object.keys(this.connections[connectionId].subscriptions).filter(
-        subscriptionId => new Date().getTime() -
-          this.connections[connectionId].subscriptions[subscriptionId].timestamp.getTime() >
-          this.hearBeatInterval * 4 * 1000
-      ).forEach(subscriptionId => { delete this.connections[connectionId].subscriptions[subscriptionId]; });
-
-      return Object.keys(this.connections[connectionId].subscriptions)
-        .map(subscriptionId => ({
-          connectionId,
-          ...this.connections[connectionId].subscriptions[subscriptionId]
-        }));
-    }).reduce((a, b) => a.concat(b), []);
+  public getAllSubscriptions() {
+    const now = Date.now();
+    const staleThreshold = this.hearBeatInterval * 4 * 1000;
+    const result: Array<{ connectionId: string } & LocalSubscriptionStoreSubscription> = [];
+
+    for (const [connectionId, connection] of this.connections) {
+      for (const [subscriptionId, subscription] of connection.subscriptions) {
+        if (now - subscription.timestamp.getTime() > staleThreshold) {
+          connection.subscriptions.delete(subscriptionId);
+        }
+      }
+
+      for (const [, subscription] of connection.subscriptions) {
+        result.push({ connectionId, ...subscription });
+      }
+    }
+
+    return result;
   }
 
-  public async cleanupSubscriptions(connectionId: string) {
-    delete this.connections[connectionId];
+  public async disconnect(connectionId: string) {
+    this.connections.delete(connectionId);
   }
 
   public async getAuthContext(connectionId: string) {
-    return this.getConnection(connectionId).authContext;
+    return this.getConnectionOrCreate(connectionId).authContext;
   }
 
   public async setAuthContext(connectionId: string, authContext) {
-    this.getConnection(connectionId).authContext = authContext;
+    this.getConnectionOrCreate(connectionId).authContext = authContext;
   }
 
-  protected getConnection(connectionId: string) {
-    if (!this.connections[connectionId]) {
-      this.connections[connectionId] = { subscriptions: {} };
+  protected getConnectionOrCreate(connectionId: string): LocalSubscriptionStoreConnection {
+    const connect = this.connections.get(connectionId);
+    if (connect) {
+      return connect;
     }
 
-    return this.connections[connectionId];
+    const connection = { subscriptions: new Map() };
+    this.connections.set(connectionId, connection);
+
+    return connection;
   }
 
   public clear() {
-    this.connections = {};
+    this.connections.clear();
   }
 }

packages/cubejs-api-gateway/src/ws/message-schema.ts

@@ -0,0 +1,67 @@
+import { z } from 'zod';
+
+const messageId = z.union([z.string().max(16), z.number()]);
+const requestId = z.string().max(64).optional();
+
+export const authMessageSchema = z.object({
+  authorization: z.string(),
+}).strict();
+
+export const unsubscribeMessageSchema = z.object({
+  unsubscribe: z.string().max(16),
+}).strict();
+
+const queryParams = z.object({
+  query: z.unknown(),
+  queryType: z.string().optional(),
+}).strict();
+
+const queryOnlyParams = z.object({
+  query: z.unknown(),
+}).strict();
+
+// Method-based messages using discriminatedUnion
+export const methodMessageSchema = z.discriminatedUnion('method', [
+  z.object({
+    method: z.literal('load'),
+    messageId,
+    requestId,
+    params: queryParams,
+  }).strict(),
+  z.object({
+    method: z.literal('sql'),
+    messageId,
+    requestId,
+    params: queryOnlyParams,
+  }).strict(),
+  z.object({
+    method: z.literal('dry-run'),
+    messageId,
+    requestId,
+    params: queryOnlyParams,
+  }).strict(),
+  z.object({
+    method: z.literal('meta'),
+    messageId,
+    requestId,
+    params: z.object({}).strict().optional(),
+  }).strict(),
+  z.object({
+    method: z.literal('subscribe'),
+    messageId,
+    requestId,
+    params: queryParams,
+  }).strict(),
+  z.object({
+    method: z.literal('unsubscribe'),
+    messageId,
+    requestId,
+    params: z.object({}).strict().optional(),
+  }).strict(),
+]);
+
+// Export types
+export type AuthMessage = z.infer<typeof authMessageSchema>;
+export type UnsubscribeMessage = z.infer<typeof unsubscribeMessageSchema>;
+export type MethodMessage = z.infer<typeof methodMessageSchema>;
+export type WsMessage = AuthMessage | UnsubscribeMessage | MethodMessage;