[ZEPPELIN-6646] validate websocket payloads

zeppelin-web-angular/projects/zeppelin-sdk/src/message-payload-guards/index.ts

@@ -0,0 +1,33 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import type { MessageReceiveDataTypeMap } from '../interfaces/message-data-type-map.interface';
+import { OP } from '../interfaces/message-operator.interface';
+
+import { isListUpdateNoteJobsPayload } from './job';
+
+export type MessagePayloadGuard = (value: unknown) => boolean;
+
+type ReceiveOP = keyof MessageReceiveDataTypeMap;
+
+/**
+ * Runtime payload guards are registered only for OPs with a demonstrated
+ * payload-shape failure. Add new guards when a concrete runtime failure
+ * shows that validation is needed.
+ */
+const MESSAGE_PAYLOAD_GUARDS: Partial<Record<ReceiveOP, MessagePayloadGuard>> = {
+  [OP.LIST_UPDATE_NOTE_JOBS]: isListUpdateNoteJobsPayload
+};
+
+export const getMessagePayloadGuard = (op: ReceiveOP): MessagePayloadGuard | undefined => {
+  return MESSAGE_PAYLOAD_GUARDS[op];
+};

zeppelin-web-angular/projects/zeppelin-sdk/src/message-payload-guards/job.ts

@@ -0,0 +1,47 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+const isRecord = (value: unknown): value is Record<string, unknown> => typeof value === 'object' && value !== null;
+
+const isJobUpdate = (value: unknown): boolean => {
+  if (!isRecord(value)) {
+    return false;
+  }
+
+  if (typeof value.noteId !== 'string') {
+    return false;
+  }
+
+  if (typeof value.isRemoved !== 'boolean') {
+    return false;
+  }
+
+  if (value.isRemoved) {
+    return true;
+  }
+
+  return typeof value.noteName === 'string';
+};
+
+export const isListUpdateNoteJobsPayload = (value: unknown): boolean => {
+  if (!isRecord(value)) {
+    return false;
+  }
+
+  const noteRunningJobs = value.noteRunningJobs;
+
+  if (!isRecord(noteRunningJobs)) {
+    return false;
+  }
+
+  return Array.isArray(noteRunningJobs.jobs) && noteRunningJobs.jobs.every(isJobUpdate);
+};

zeppelin-web-angular/projects/zeppelin-sdk/src/message.spec.ts

@@ -0,0 +1,146 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import { afterEach, describe, expect, it, vi } from 'vitest';
+
+import type { MessageReceiveDataTypeMap } from './interfaces/message-data-type-map.interface';
+import { OP } from './interfaces/message-operator.interface';
+import type { WebSocketMessage } from './interfaces/websocket-message.interface';
+import { Message } from './message';
+
+const asReceivedMessage = (message: unknown): WebSocketMessage<MessageReceiveDataTypeMap> =>
+  message as WebSocketMessage<MessageReceiveDataTypeMap>;
+
+afterEach(() => {
+  vi.restoreAllMocks();
+});
+
+describe('Message.receive', () => {
+  it('passes a non-removal job update with noteName', () => {
+    const message = new Message();
+    const listener = vi.fn();
+    const data = {
+      noteRunningJobs: {
+        jobs: [
+          {
+            noteId: 'note-1',
+            noteName: 'Test Note',
+            isRemoved: false
+          }
+        ]
+      }
+    };
+
+    message.receive(OP.LIST_UPDATE_NOTE_JOBS).subscribe(listener);
+
+    message.shortCircuit(
+      asReceivedMessage({
+        op: OP.LIST_UPDATE_NOTE_JOBS,
+        data
+      })
+    );
+
+    expect(listener).toHaveBeenCalledWith(data);
+  });
+
+  it('passes a partial removal payload without noteName', () => {
+    const message = new Message();
+    const listener = vi.fn();
+    const data = {
+      noteRunningJobs: {
+        jobs: [
+          {
+            noteId: 'note-1',
+            isRemoved: true
+          }
+        ]
+      }
+    };
+
+    message.receive(OP.LIST_UPDATE_NOTE_JOBS).subscribe(listener);
+
+    message.shortCircuit(
+      asReceivedMessage({
+        op: OP.LIST_UPDATE_NOTE_JOBS,
+        data
+      })
+    );
+
+    expect(listener).toHaveBeenCalledWith(data);
+  });
+
+  it('filters a non-removal job update without noteName and warns with the OP only', () => {
+    const message = new Message();
+    const listener = vi.fn();
+    const consoleWarn = vi.spyOn(console, 'warn').mockImplementation(() => {});
+
+    message.receive(OP.LIST_UPDATE_NOTE_JOBS).subscribe(listener);
+
+    message.shortCircuit(
+      asReceivedMessage({
+        op: OP.LIST_UPDATE_NOTE_JOBS,
+        data: {
+          noteRunningJobs: {
+            jobs: [
+              {
+                noteId: 'note-1',
+                isRemoved: false
+              }
+            ]
+          }
+        }
+      })
+    );
+
+    expect(listener).not.toHaveBeenCalled();
+    expect(consoleWarn).toHaveBeenCalledTimes(1);
+    expect(consoleWarn).toHaveBeenCalledWith(
+      `Dropped WebSocket OP ${String(OP.LIST_UPDATE_NOTE_JOBS)}: payload failed validation`
+    );
+  });
+
+  it('filters a payload without a jobs array', () => {
+    const message = new Message();
+    const listener = vi.fn();
+    vi.spyOn(console, 'warn').mockImplementation(() => {});
+
+    message.receive(OP.LIST_UPDATE_NOTE_JOBS).subscribe(listener);
+
+    message.shortCircuit(
+      asReceivedMessage({
+        op: OP.LIST_UPDATE_NOTE_JOBS,
+        data: {
+          noteRunningJobs: {}
+        }
+      })
+    );
+
+    expect(listener).not.toHaveBeenCalled();
+  });
+
+  it('keeps existing behavior for an OP without a guard', () => {
+    const message = new Message();
+    const listener = vi.fn();
+    const data = {};
+
+    message.receive(OP.NOTE).subscribe(listener);
+
+    message.shortCircuit(
+      asReceivedMessage({
+        op: OP.NOTE,
+        data
+      })
+    );
+
+    expect(listener).toHaveBeenCalledWith(data);
+  });
+});

zeppelin-web-angular/projects/zeppelin-sdk/src/message.ts

@@ -30,6 +30,8 @@ import {
 } from './interfaces/message-paragraph.interface';
 import { WebSocketMessage } from './interfaces/websocket-message.interface';
 
+import { getMessagePayloadGuard } from './message-payload-guards';
+
 export type ArgumentsType<T> = T extends (...args: infer U) => void ? U : never;
 
 export type SendArgumentsType<K extends keyof MessageSendDataTypeMap> = MessageSendDataTypeMap[K] extends undefined
@@ -173,8 +175,19 @@ export class Message {
   }
 
   receive<K extends keyof MessageReceiveDataTypeMap>(op: K): Observable<Record<K, MessageReceiveDataTypeMap[K]>[K]> {
+    const guard = getMessagePayloadGuard(op);
+
     return this.received$.pipe(
       filter(message => message.op === op),
+      filter(message => {
+        if (!guard || guard(message.data)) {
+          return true;
+        }
+
+        // The payload can be large and carries note names, so log the OP alone.
+        console.warn(`Dropped WebSocket OP ${String(op)}: payload failed validation`);
+        return false;
+      }),
       map(message => message.data)
     ) as Observable<Record<K, MessageReceiveDataTypeMap[K]>[K]>;
   }

zeppelin-web-angular/src/app/core/message-listener/message-listener.spec.ts

@@ -0,0 +1,90 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import { Subject } from 'rxjs';
+import { afterEach, describe, expect, it, vi } from 'vitest';
+
+import { Message, OP, MessageReceiveDataTypeMap } from '@zeppelin/sdk';
+
+import { MessageListener, MessageListenersManager } from './message-listener';
+
+afterEach(() => {
+  vi.restoreAllMocks();
+  vi.useRealTimers();
+});
+
+describe('MessageListener', () => {
+  it('logs handler errors with the OP, rethrows them, and keeps the subscription active', () => {
+    // RxJS rethrows an error thrown inside `next` from a timer, so the subscription itself survives.
+    vi.useFakeTimers();
+
+    const received$ = new Subject<MessageReceiveDataTypeMap[OP.NOTE]>();
+    const messageService = {
+      receive: vi.fn(() => received$.asObservable())
+    } as unknown as Message;
+
+    const error = new Error('boom');
+    const consoleError = vi.spyOn(console, 'error').mockImplementation(() => {});
+
+    class TestComponent extends MessageListenersManager {
+      calls = 0;
+
+      handleNote(_data: MessageReceiveDataTypeMap[OP.NOTE]): void {
+        this.calls++;
+
+        if (this.calls === 1) {
+          throw error;
+        }
+      }
+    }
+
+    const descriptor = Object.getOwnPropertyDescriptor(TestComponent.prototype, 'handleNote')!;
+
+    MessageListener(OP.NOTE)(TestComponent.prototype, 'handleNote', descriptor);
+
+    const component = new TestComponent(messageService);
+    const data = {} as MessageReceiveDataTypeMap[OP.NOTE];
+
+    received$.next(data);
+    received$.next(data);
+
+    expect(component.calls).toBe(2);
+    expect(consoleError).toHaveBeenCalledWith(`Failed to handle WebSocket OP ${String(OP.NOTE)}`, error);
+    expect(() => vi.runAllTimers()).toThrow(error);
+  });
+
+  it('passes received data to the handler', () => {
+    const received$ = new Subject<MessageReceiveDataTypeMap[OP.NOTE]>();
+    const messageService = {
+      receive: vi.fn(() => received$.asObservable())
+    } as unknown as Message;
+
+    class TestComponent extends MessageListenersManager {
+      receivedData?: MessageReceiveDataTypeMap[OP.NOTE];
+
+      handleNote(data: MessageReceiveDataTypeMap[OP.NOTE]): void {
+        this.receivedData = data;
+      }
+    }
+
+    const descriptor = Object.getOwnPropertyDescriptor(TestComponent.prototype, 'handleNote')!;
+
+    MessageListener(OP.NOTE)(TestComponent.prototype, 'handleNote', descriptor);
+
+    const component = new TestComponent(messageService);
+    const data = {} as MessageReceiveDataTypeMap[OP.NOTE];
+
+    received$.next(data);
+
+    expect(component.receivedData).toBe(data);
+  });
+});

zeppelin-web-angular/src/app/core/message-listener/message-listener.ts

@@ -49,8 +49,13 @@ export function MessageListener<K extends keyof MessageReceiveDataTypeMap>(op: K
 
       this.__zeppelinMessageListeners$__.add(
         this.messageService.receive(op).subscribe(data => {
-          // @ts-ignore
-          oldValue.apply(this, [data]);
+          try {
+            // @ts-ignore
+            oldValue.apply(this, [data]);
+          } catch (error) {
+            console.error(`Failed to handle WebSocket OP ${String(op)}`, error);
+            throw error;
+          }
         })
       );
     };