feat: added reverse response message to raw duplex streams

Author: tegefaulkes

src/rpc/RPCClient.ts

@@ -1,12 +1,13 @@
 import type { WritableStream, ReadableStream } from 'stream/web';
-import type { ContextTimed, ContextTimedInput } from '@matrixai/contexts';
+import type { ContextTimedInput } from '@matrixai/contexts';
 import type {
   HandlerType,
   JSONRPCRequestMessage,
   StreamFactory,
   ClientManifest,
-  RPCStream, JSONRPCResponseResult
-} from "./types";
+  RPCStream,
+  JSONRPCResponseResult,
+} from './types';
 import type { JSONValue } from '../types';
 import type {
   JSONRPCRequest,
@@ -20,8 +21,7 @@ import { Timer } from '@matrixai/timer';
 import * as rpcUtilsMiddleware from './utils/middleware';
 import * as rpcErrors from './errors';
 import * as rpcUtils from './utils/utils';
-import { never, promise } from "../utils";
-import { parseJSONRPCResponse } from "./utils/utils";
+import { never, promise } from '../utils';
 
 const timerCleanupReasonSymbol = Symbol('timerCleanUpReasonSymbol');
 
@@ -255,15 +255,18 @@ class RPCClient<M extends ClientManifest> {
     method: string,
     ctx: Partial<ContextTimedInput> = {},
   ): Promise<RPCStream<O, I>> {
+    // Setting up abort signal and timer
     const abortController = new AbortController();
     const signal = abortController.signal;
     // A promise that will reject if there is an abort signal or timeout
     const abortRaceProm = promise<never>();
     // Prevent unhandled rejection when we're done with the promise
     abortRaceProm.p.catch(() => {});
-    signal.addEventListener('abort', () => {
+    const abortRacePromHandler = () => {
       abortRaceProm.rejectP(signal.reason);
-    }, { once: true});
+    };
+    signal.addEventListener('abort', abortRacePromHandler);
+
     let abortHandler: () => void;
     if (ctx.signal != null) {
       // Propagate signal events
@@ -284,7 +287,10 @@ class RPCClient<M extends ClientManifest> {
     const cleanUp = () => {
       // Clean up the timer and signal
       if (ctx.timer == null) timer.cancel(timerCleanupReasonSymbol);
-      signal.removeEventListener('abort', abortHandler);
+      if (ctx.signal != null) {
+        ctx.signal.removeEventListener('abort', abortHandler);
+      }
+      signal.addEventListener('abort', abortRacePromHandler);
     };
     // Setting up abort events for timeout
     const timeoutError = new rpcErrors.ErrorRPCTimedOut();
@@ -294,6 +300,7 @@ class RPCClient<M extends ClientManifest> {
       },
       () => {}, // Ignore cancellation error
     );
+
     // Hooking up agnostic stream side
     let rpcStream: RPCStream<Uint8Array, Uint8Array>;
     const streamFactoryProm = this.streamFactory({ signal, timer });
@@ -302,19 +309,10 @@ class RPCClient<M extends ClientManifest> {
     } catch (e) {
       cleanUp();
       void streamFactoryProm.then((stream) =>
-        stream.cancel('stream timed out early'),
+        stream.cancel(Error('TMP stream timed out early')),
       );
       throw e;
     }
-    const cancelStream = () => {
-      rpcStream.cancel(signal.reason);
-    };
-    if (signal.aborted) {
-      cancelStream();
-    } else {
-      signal.addEventListener('abort', cancelStream);
-    }
-    // Setting up event for stream timeout
     void timer.then(
       () => {
         rpcStream.cancel(new rpcErrors.ErrorRPCTimedOut());
@@ -356,7 +354,6 @@ class RPCClient<M extends ClientManifest> {
         .catch(() => {}),
     ]).finally(() => {
       cleanUp();
-      signal.removeEventListener('abort', cancelStream);
     });
 
     // Returning interface
@@ -387,9 +384,25 @@ class RPCClient<M extends ClientManifest> {
     method: string,
     headerParams: JSONValue,
     ctx: Partial<ContextTimedInput> = {},
-  ): Promise<RPCStream<Uint8Array, Uint8Array, Record<string, JSONValue> & {result: JSONValue, command: string}>> {
+  ): Promise<
+    RPCStream<
+      Uint8Array,
+      Uint8Array,
+      Record<string, JSONValue> & { result: JSONValue; command: string }
+    >
+  > {
+    // Setting up abort signal and timer
     const abortController = new AbortController();
     const signal = abortController.signal;
+    // A promise that will reject if there is an abort signal or timeout
+    const abortRaceProm = promise<never>();
+    // Prevent unhandled rejection when we're done with the promise
+    abortRaceProm.p.catch(() => {});
+    const abortRacePromHandler = () => {
+      abortRaceProm.rejectP(signal.reason);
+    };
+    signal.addEventListener('abort', abortRacePromHandler);
+
     let abortHandler: () => void;
     if (ctx.signal != null) {
       // Propagate signal events
@@ -410,24 +423,34 @@ class RPCClient<M extends ClientManifest> {
     const cleanUp = () => {
       // Clean up the timer and signal
       if (ctx.timer == null) timer.cancel(timerCleanupReasonSymbol);
-      signal.removeEventListener('abort', abortHandler);
+      if (ctx.signal != null) {
+        ctx.signal.removeEventListener('abort', abortHandler);
+      }
+      signal.addEventListener('abort', abortRacePromHandler);
     };
+    // Setting up abort events for timeout
     const timeoutError = new rpcErrors.ErrorRPCTimedOut();
     void timer.then(
       () => {
         abortController.abort(timeoutError);
       },
-      () => {},
+      () => {}, // Ignore cancellation error
     );
-    let streamCreation: [JSONValue, RPCStream<Uint8Array, Uint8Array>];
-    const setupStream = async (): Promise<[JSONValue, RPCStream<Uint8Array, Uint8Array>]> => {
+
+    const setupStream = async (): Promise<
+      [JSONValue, RPCStream<Uint8Array, Uint8Array>]
+    > => {
       if (signal.aborted) throw signal.reason;
       const abortProm = promise<never>();
-      // ignore error if orphaned
+      // Ignore error if orphaned
       void abortProm.p.catch(() => {});
-      signal.addEventListener('abort', () => {
-        abortProm.rejectP(signal.reason);
-      }, {once: true});
+      signal.addEventListener(
+        'abort',
+        () => {
+          abortProm.rejectP(signal.reason);
+        },
+        { once: true },
+      );
       const rpcStream = await Promise.race([
         this.streamFactory({ signal, timer }),
         abortProm.p,
@@ -441,48 +464,48 @@ class RPCClient<M extends ClientManifest> {
       };
       await tempWriter.write(Buffer.from(JSON.stringify(header)));
       tempWriter.releaseLock();
-      const headTransformStream = rpcUtilsMiddleware.binaryToJsonMessageStream(
+      const headTransformStream = rpcUtils.parseHeadStream(
         rpcUtils.parseJSONRPCResponse,
       );
       void rpcStream.readable
         // Allow us to re-use the readable after reading the first message
-        .pipeTo(headTransformStream.writable, {
-          preventClose: true,
-          preventCancel: true,
-        })
+        .pipeTo(headTransformStream.writable)
         // Ignore any errors here, we only care that it ended
         .catch(() => {});
       const tempReader = headTransformStream.readable.getReader();
       let leadingMessage: JSONRPCResponseResult;
       try {
-        const message = await Promise.race([
-          tempReader.read(),
-          abortProm.p,
-        ]);
+        const message = await Promise.race([tempReader.read(), abortProm.p]);
+        const messageValue = message.value as JSONRPCResponse;
         if (message.done) never();
-        if ('error'in message.value) {
+        if ('error' in messageValue) {
           const metadata = {
             ...(rpcStream.meta ?? {}),
             command: method,
           };
-          throw rpcUtils.toError(message.value.error.data, metadata);
+          throw rpcUtils.toError(messageValue.error.data, metadata);
         }
-        leadingMessage = message.value;
+        leadingMessage = messageValue;
       } catch (e) {
-        await tempReader.cancel();
         rpcStream.cancel(Error('TMP received error in leading response'));
         throw e;
       }
-      // Downgrade back to the raw stream
-      await tempReader.cancel();
-      return [leadingMessage.result, rpcStream];
+      tempReader.releaseLock();
+      const newRpcStream: RPCStream<Uint8Array, Uint8Array> = {
+        writable: rpcStream.writable,
+        readable: headTransformStream.readable as ReadableStream<Uint8Array>,
+        cancel: rpcStream.cancel,
+        meta: rpcStream.meta,
+      };
+      return [leadingMessage.result, newRpcStream];
     };
+    let streamCreation: [JSONValue, RPCStream<Uint8Array, Uint8Array>];
     try {
       streamCreation = await setupStream();
     } finally {
       cleanUp();
     }
-    const [result, rpcStream] = streamCreation
+    const [result, rpcStream] = streamCreation;
     const metadata = {
       ...(rpcStream.meta ?? {}),
       result,

src/rpc/RPCServer.ts

@@ -31,7 +31,6 @@ import * as rpcEvents from './events';
 import * as rpcUtils from './utils/utils';
 import * as rpcErrors from './errors';
 import * as rpcUtilsMiddleware from './utils/middleware';
-import * as utils from '../utils/utils';
 import { never } from '../utils/utils';
 import { sysexits } from '../errors';
 
@@ -586,8 +585,8 @@ class RPCServer extends EventTarget {
           error: rpcError,
           id: null,
         };
-        await headerWriter.write(Buffer.from(JSON.stringify(rpcErrorMessage)))
-        // clean up and return
+        await headerWriter.write(Buffer.from(JSON.stringify(rpcErrorMessage)));
+        // Clean up and return
         timer.cancel(cleanupReason);
         abortController.signal.removeEventListener('abort', handleAbort);
         graceTimer?.cancel(cleanupReason);
@@ -597,12 +596,12 @@ class RPCServer extends EventTarget {
       }
       const [leadingResult, outputStream] = handlerResult;
 
-      if (leadingResult !== undefined){
-        // writing leading metadata
+      if (leadingResult !== undefined) {
+        // Writing leading metadata
         const leadingMessage: JSONRPCResponseResult = {
-          jsonrpc: "2.0",
+          jsonrpc: '2.0',
           result: leadingResult,
-          id: null
+          id: null,
         };
         await headerWriter.write(Buffer.from(JSON.stringify(leadingMessage)));
       }

src/rpc/types.ts

@@ -264,7 +264,13 @@ type DuplexCallerImplementation<
 type RawCallerImplementation = (
   headerParams: JSONValue,
   ctx?: Partial<ContextTimedInput>,
-) => Promise<RPCStream<Uint8Array, Uint8Array, Record<string, JSONValue> & {result: JSONValue, command: string}>>;
+) => Promise<
+  RPCStream<
+    Uint8Array,
+    Uint8Array,
+    Record<string, JSONValue> & { result: JSONValue; command: string }
+  >
+>;
 
 type ConvertDuplexCaller<T> = T extends DuplexCaller<infer I, infer O>
   ? DuplexCallerImplementation<I, O>

src/rpc/utils/utils.ts

@@ -13,6 +13,7 @@ import type {
 import type { JSONValue } from '../../types';
 import type { Timer } from '@matrixai/timer';
 import { TransformStream } from 'stream/web';
+import { JSONParser } from '@streamparser/json';
 import { AbstractError } from '@matrixai/errors';
 import * as rpcErrors from '../errors';
 import * as utils from '../../utils';
@@ -429,6 +430,72 @@ function getHandlerTypes(
   return out;
 }
 
+/**
+ * This function is a factory to create a TransformStream that will
+ * transform a `Uint8Array` stream to a JSONRPC message stream.
+ * The parsed messages will be validated with the provided messageParser, this
+ * also infers the type of the stream output.
+ * @param messageParser - Validates the JSONRPC messages, so you can select for a
+ *  specific type of message
+ * @param bufferByteLimit - sets the number of bytes buffered before throwing an
+ *  error. This is used to avoid infinitely buffering the input.
+ */
+function parseHeadStream<T extends JSONRPCMessage>(
+  messageParser: (message: unknown) => T,
+  bufferByteLimit: number = 1024 * 1024,
+): TransformStream<Uint8Array, T | Uint8Array> {
+  const parser = new JSONParser({
+    separator: '',
+    paths: ['$'],
+  });
+  let bytesWritten: number = 0;
+  let parsing = true;
+  let ended = false;
+
+  const endP = utils.promise();
+  parser.onEnd = () => endP.resolveP();
+
+  return new TransformStream<Uint8Array, T | Uint8Array>(
+    {
+      flush: async () => {
+        if (!parser.isEnded) parser.end();
+        await endP.p;
+      },
+      start: (controller) => {
+        parser.onValue = async (value) => {
+          const jsonMessage = messageParser(value.value);
+          controller.enqueue(jsonMessage);
+          bytesWritten = 0;
+          parsing = false;
+        };
+      },
+      transform: async (chunk, controller) => {
+        if (parsing) {
+          try {
+            bytesWritten += chunk.byteLength;
+            parser.write(chunk);
+          } catch (e) {
+            throw new rpcErrors.ErrorRPCParse(undefined, { cause: e });
+          }
+          if (bytesWritten > bufferByteLimit) {
+            throw new rpcErrors.ErrorRPCMessageLength();
+          }
+        } else {
+          // Wait for parser to end
+          if (!ended) {
+            parser.end();
+            await endP.p;
+            ended = true;
+          }
+          // Pass through normal chunks
+          controller.enqueue(chunk);
+        }
+      },
+    },
+    { highWaterMark: 1 },
+  );
+}
+
 export {
   parseJSONRPCRequest,
   parseJSONRPCRequestMessage,
@@ -442,4 +509,5 @@ export {
   clientInputTransformStream,
   clientOutputTransformStream,
   getHandlerTypes,
+  parseHeadStream,
 };