account for RangeError: Array buffer allocation failed

__tests__/allocation.test.ts

@@ -0,0 +1,177 @@
+import { beforeEach, describe, test, expect, vi, assert } from 'vitest';
+import { TestSetupHelpers, transports } from '../testUtil/fixtures/transports';
+import { BinaryCodec, Codec } from '../codec';
+import {
+  advanceFakeTimersByHeartbeat,
+  createPostTestCleanups,
+} from '../testUtil/fixtures/cleanup';
+import { createServer } from '../router/server';
+import { createClient } from '../router/client';
+import { TestServiceSchema } from '../testUtil/fixtures/services';
+import { waitFor } from '../testUtil/fixtures/cleanup';
+import { numberOfConnections, closeAllConnections } from '../testUtil';
+import { cleanupTransports } from '../testUtil/fixtures/cleanup';
+import { testFinishesCleanly } from '../testUtil/fixtures/cleanup';
+import { ProtocolError } from '../transport/events';
+
+let isOom = false;
+// simulate RangeError: Array buffer allocation failed
+const OomableCodec: Codec = {
+  toBuffer(obj) {
+    if (isOom) {
+      throw new RangeError('failed allocation');
+    }
+
+    return BinaryCodec.toBuffer(obj);
+  },
+  fromBuffer: (buff: Uint8Array) => {
+    return BinaryCodec.fromBuffer(buff);
+  },
+};
+
+describe.each(transports)(
+  'failed allocation test ($name transport)',
+  async (transport) => {
+    const clientOpts = { codec: OomableCodec };
+    const serverOpts = { codec: BinaryCodec };
+
+    const { addPostTestCleanup, postTestCleanup } = createPostTestCleanups();
+    let getClientTransport: TestSetupHelpers['getClientTransport'];
+    let getServerTransport: TestSetupHelpers['getServerTransport'];
+    beforeEach(async () => {
+      // only allow client to oom, server has sane oom handling already
+      const setup = await transport.setup({
+        client: clientOpts,
+        server: serverOpts,
+      });
+      getClientTransport = setup.getClientTransport;
+      getServerTransport = setup.getServerTransport;
+      isOom = false;
+
+      return async () => {
+        await postTestCleanup();
+        await setup.cleanup();
+      };
+    });
+
+    test('oom during heartbeat kills the session, client starts new session', async () => {
+      // setup
+      const clientTransport = getClientTransport('client');
+      const serverTransport = getServerTransport();
+      const services = { test: TestServiceSchema };
+      const server = createServer(serverTransport, services);
+      const client = createClient<typeof services>(
+        clientTransport,
+        serverTransport.clientId,
+      );
+
+      const errMock = vi.fn();
+      clientTransport.addEventListener('protocolError', errMock);
+      addPostTestCleanup(async () => {
+        clientTransport.removeEventListener('protocolError', errMock);
+        await cleanupTransports([clientTransport, serverTransport]);
+      });
+
+      // establish initial connection
+      const result = await client.test.add.rpc({ n: 1 });
+      expect(result).toStrictEqual({ ok: true, payload: { result: 1 } });
+
+      await waitFor(() => expect(numberOfConnections(serverTransport)).toBe(1));
+      await waitFor(() => expect(numberOfConnections(clientTransport)).toBe(1));
+      const oldClientSession = serverTransport.sessions.get('client');
+      const oldServerSession = clientTransport.sessions.get('SERVER');
+      assert(oldClientSession);
+      assert(oldServerSession);
+
+      // simulate some OOM during heartbeat
+      for (let i = 0; i < 5; i++) {
+        isOom = i % 2 === 0;
+        await advanceFakeTimersByHeartbeat();
+      }
+
+      // verify session on client is dead
+      await waitFor(() => expect(clientTransport.sessions.size).toBe(0));
+
+      // verify we got MessageSendFailure errors
+      await waitFor(() => {
+        expect(errMock).toHaveBeenCalledWith(
+          expect.objectContaining({
+            type: ProtocolError.MessageSendFailure,
+          }),
+        );
+      });
+
+      // client should be able to reconnect and make new calls
+      isOom = false;
+      const result2 = await client.test.add.rpc({ n: 2 });
+      expect(result2).toStrictEqual({ ok: true, payload: { result: 3 } });
+
+      // verify new session IDs are different from old ones
+      const newClientSession = serverTransport.sessions.get('client');
+      const newServerSession = clientTransport.sessions.get('SERVER');
+      assert(newClientSession);
+      assert(newServerSession);
+      expect(newClientSession.id).not.toBe(oldClientSession.id);
+      expect(newServerSession.id).not.toBe(oldServerSession.id);
+
+      await testFinishesCleanly({
+        clientTransports: [clientTransport],
+        serverTransport,
+        server,
+      });
+    });
+
+    test('oom during handshake kills the session, client starts new session', async () => {
+      // setup
+      const clientTransport = getClientTransport('client');
+      const serverTransport = getServerTransport();
+      const services = { test: TestServiceSchema };
+      const server = createServer(serverTransport, services);
+      const client = createClient<typeof services>(
+        clientTransport,
+        serverTransport.clientId,
+      );
+      const errMock = vi.fn();
+      clientTransport.addEventListener('protocolError', errMock);
+      addPostTestCleanup(async () => {
+        clientTransport.removeEventListener('protocolError', errMock);
+        await cleanupTransports([clientTransport, serverTransport]);
+      });
+
+      // establish initial connection
+      await client.test.add.rpc({ n: 1 });
+      await waitFor(() => expect(numberOfConnections(serverTransport)).toBe(1));
+      await waitFor(() => expect(numberOfConnections(clientTransport)).toBe(1));
+
+      // close connection to force reconnection
+      closeAllConnections(clientTransport);
+      await waitFor(() => expect(numberOfConnections(serverTransport)).toBe(0));
+      await waitFor(() => expect(numberOfConnections(clientTransport)).toBe(0));
+
+      // simulate OOM during handshake
+      isOom = true;
+      clientTransport.connect('SERVER');
+      await waitFor(() => expect(numberOfConnections(serverTransport)).toBe(0));
+      await waitFor(() => expect(numberOfConnections(clientTransport)).toBe(0));
+
+      await waitFor(() => {
+        expect(errMock).toHaveBeenCalledWith(
+          expect.objectContaining({
+            type: ProtocolError.MessageSendFailure,
+          }),
+        );
+      });
+
+      // client should be able to reconnect and make new calls
+      isOom = false;
+      const result = await client.test.add.rpc({ n: 2 });
+      expect(result).toStrictEqual({ ok: true, payload: { result: 3 } });
+
+      await testFinishesCleanly({
+        clientTransports: [clientTransport],
+        serverTransport,
+        server,
+      });
+    });
+  },
+);

codec/adapter.ts

@@ -0,0 +1,53 @@
+import { Value } from '@sinclair/typebox/value';
+import {
+  OpaqueTransportMessage,
+  OpaqueTransportMessageSchema,
+} from '../transport';
+import { Codec } from './types';
+import { DeserializeResult, SerializeResult } from '../transport/results';
+import { coerceErrorString } from '../transport/stringifyError';
+
+/**
+ * Adapts a {@link Codec} to the {@link OpaqueTransportMessage} format,
+ * accounting for fallibility of toBuffer and fromBuffer and wrapping
+ * it with a Result type.
+ */
+export class CodecMessageAdapter {
+  constructor(private readonly codec: Codec) {}
+
+  toBuffer(msg: OpaqueTransportMessage): SerializeResult {
+    try {
+      return {
+        ok: true,
+        value: this.codec.toBuffer(msg),
+      };
+    } catch (e) {
+      return {
+        ok: false,
+        reason: coerceErrorString(e),
+      };
+    }
+  }
+
+  fromBuffer(buf: Uint8Array): DeserializeResult {
+    try {
+      const parsedMsg = this.codec.fromBuffer(buf);
+      if (!Value.Check(OpaqueTransportMessageSchema, parsedMsg)) {
+        return {
+          ok: false,
+          reason: 'transport message schema mismatch',
+        };
+      }
+
+      return {
+        ok: true,
+        value: parsedMsg,
+      };
+    } catch (e) {
+      return {
+        ok: false,
+        reason: coerceErrorString(e),
+      };
+    }
+  }
+}

codec/binary.ts

@@ -10,15 +10,11 @@ export const BinaryCodec: Codec = {
     return encode(obj, { ignoreUndefined: true });
   },
   fromBuffer: (buff: Uint8Array) => {
-    try {
-      const res = decode(buff);
-      if (typeof res !== 'object') {
-        return null;
-      }
-
-      return res;
-    } catch {
-      return null;
+    const res = decode(buff);
+    if (typeof res !== 'object' || res === null) {
+      throw new Error('unpacked msg is not an object');
     }
+
+    return res;
   },
 };

codec/codec.test.ts

@@ -41,10 +41,10 @@ describe.each(codecs)('codec -- $name', ({ codec }) => {
     expect(codec.fromBuffer(codec.toBuffer(msg))).toStrictEqual(msg);
   });
 
-  test('invalid json returns null', () => {
-    expect(codec.fromBuffer(Buffer.from(''))).toBeNull();
-    expect(codec.fromBuffer(Buffer.from('['))).toBeNull();
-    expect(codec.fromBuffer(Buffer.from('[{}'))).toBeNull();
-    expect(codec.fromBuffer(Buffer.from('{"a":1}[]'))).toBeNull();
+  test('invalid json throws', () => {
+    expect(() => codec.fromBuffer(Buffer.from(''))).toThrow();
+    expect(() => codec.fromBuffer(Buffer.from('['))).toThrow();
+    expect(() => codec.fromBuffer(Buffer.from('[{}'))).toThrow();
+    expect(() => codec.fromBuffer(Buffer.from('{"a":1}[]'))).toThrow();
   });
 });

codec/index.ts

@@ -1,3 +1,4 @@
 export { BinaryCodec } from './binary';
 export { NaiveJsonCodec } from './json';
 export type { Codec } from './types';
+export { CodecMessageAdapter } from './adapter';

codec/json.ts

@@ -48,23 +48,21 @@ export const NaiveJsonCodec: Codec = {
     );
   },
   fromBuffer: (buff: Uint8Array) => {
-    try {
-      const parsed = JSON.parse(
-        decoder.decode(buff),
-        function reviver(_key, val: unknown) {
-          if ((val as Base64EncodedValue | undefined)?.$t) {
-            return base64ToUint8Array((val as Base64EncodedValue).$t);
-          } else {
-            return val;
-          }
-        },
-      ) as unknown;
-
-      if (typeof parsed === 'object') return parsed;
+    const parsed = JSON.parse(
+      decoder.decode(buff),
+      function reviver(_key, val: unknown) {
+        if ((val as Base64EncodedValue | undefined)?.$t) {
+          return base64ToUint8Array((val as Base64EncodedValue).$t);
+        } else {
+          return val;
+        }
+      },
+    ) as unknown;
 
-      return null;
-    } catch {
-      return null;
+    if (typeof parsed !== 'object' || parsed === null) {
+      throw new Error('unpacked msg is not an object');
     }
+
+    return parsed;
   },
 };

codec/types.ts

@@ -14,5 +14,5 @@ export interface Codec {
    * @param buf - The Uint8 buffer to decode.
    * @returns The decoded object, or null if decoding failed.
    */
-  fromBuffer(buf: Uint8Array): object | null;
+  fromBuffer(buf: Uint8Array): object;
 }

package.json

@@ -1,7 +1,7 @@
 {
   "name": "@replit/river",
   "description": "It's like tRPC but... with JSON Schema Support, duplex streaming and support for service multiplexing. Transport agnostic!",
-  "version": "0.207.2",
+  "version": "0.207.3",
   "type": "module",
   "exports": {
     ".": {