wip adding leading response message to raw RPC calls

Author: tegefaulkes

src/rpc/RPCClient.ts

@@ -5,8 +5,8 @@ import type {
   JSONRPCRequestMessage,
   StreamFactory,
   ClientManifest,
-  RPCStream,
-} from './types';
+  RPCStream, JSONRPCResponseResult
+} from "./types";
 import type { JSONValue } from '../types';
 import type {
   JSONRPCRequest,
@@ -20,7 +20,8 @@ import { Timer } from '@matrixai/timer';
 import * as rpcUtilsMiddleware from './utils/middleware';
 import * as rpcErrors from './errors';
 import * as rpcUtils from './utils/utils';
-import { promise } from '../utils';
+import { never, promise } from "../utils";
+import { parseJSONRPCResponse } from "./utils/utils";
 
 const timerCleanupReasonSymbol = Symbol('timerCleanUpReasonSymbol');
 
@@ -260,12 +261,14 @@ class RPCClient<M extends ClientManifest> {
     const abortRaceProm = promise<never>();
     // Prevent unhandled rejection when we're done with the promise
     abortRaceProm.p.catch(() => {});
+    signal.addEventListener('abort', () => {
+      abortRaceProm.rejectP(signal.reason);
+    }, { once: true});
     let abortHandler: () => void;
     if (ctx.signal != null) {
       // Propagate signal events
       abortHandler = () => {
         abortController.abort(ctx.signal?.reason);
-        abortRaceProm.rejectP(ctx.signal?.reason);
       };
       if (ctx.signal.aborted) abortHandler();
       ctx.signal.addEventListener('abort', abortHandler);
@@ -288,7 +291,6 @@ class RPCClient<M extends ClientManifest> {
     void timer.then(
       () => {
         abortController.abort(timeoutError);
-        abortRaceProm.rejectP(timeoutError);
       },
       () => {}, // Ignore cancellation error
     );
@@ -384,29 +386,27 @@ class RPCClient<M extends ClientManifest> {
   public async rawStreamCaller(
     method: string,
     headerParams: JSONValue,
-    ctx: Partial<ContextTimed> = {},
-  ): Promise<RPCStream<Uint8Array, Uint8Array>> {
+    ctx: Partial<ContextTimedInput> = {},
+  ): Promise<RPCStream<Uint8Array, Uint8Array, Record<string, JSONValue> & {result: JSONValue, command: string}>> {
     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(() => {});
     let abortHandler: () => void;
     if (ctx.signal != null) {
       // Propagate signal events
       abortHandler = () => {
         abortController.abort(ctx.signal?.reason);
-        abortRaceProm.rejectP(ctx.signal?.reason);
       };
       if (ctx.signal.aborted) abortHandler();
       ctx.signal.addEventListener('abort', abortHandler);
     }
-    const timer =
-      ctx.timer ??
-      new Timer({
-        delay: this.streamKeepAliveTimeoutTime,
+    let timer: Timer;
+    if (!(ctx.timer instanceof Timer)) {
+      timer = new Timer({
+        delay: ctx.timer ?? this.streamKeepAliveTimeoutTime,
       });
+    } else {
+      timer = ctx.timer;
+    }
     const cleanUp = () => {
       // Clean up the timer and signal
       if (ctx.timer == null) timer.cancel(timerCleanupReasonSymbol);
@@ -416,13 +416,22 @@ class RPCClient<M extends ClientManifest> {
     void timer.then(
       () => {
         abortController.abort(timeoutError);
-        abortRaceProm.rejectP(timeoutError);
       },
       () => {},
     );
-    let rpcStream: RPCStream<Uint8Array, Uint8Array>;
-    const setupStream = async () => {
-      const rpcStream = await this.streamFactory({ signal, timer });
+    let streamCreation: [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
+      void abortProm.p.catch(() => {});
+      signal.addEventListener('abort', () => {
+        abortProm.rejectP(signal.reason);
+      }, {once: true});
+      const rpcStream = await Promise.race([
+        this.streamFactory({ signal, timer }),
+        abortProm.p,
+      ]);
       const tempWriter = rpcStream.writable.getWriter();
       const header: JSONRPCRequestMessage = {
         jsonrpc: '2.0',
@@ -432,15 +441,51 @@ class RPCClient<M extends ClientManifest> {
       };
       await tempWriter.write(Buffer.from(JSON.stringify(header)));
       tempWriter.releaseLock();
-      return rpcStream;
+      const headTransformStream = rpcUtilsMiddleware.binaryToJsonMessageStream(
+        rpcUtils.parseJSONRPCResponse,
+      );
+      void rpcStream.readable
+        // Allow us to re-use the readable after reading the first message
+        .pipeTo(headTransformStream.writable, {
+          preventClose: true,
+          preventCancel: true,
+        })
+        // 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,
+        ]);
+        if (message.done) never();
+        if ('error'in message.value) {
+          const metadata = {
+            ...(rpcStream.meta ?? {}),
+            command: method,
+          };
+          throw rpcUtils.toError(message.value.error.data, metadata);
+        }
+        leadingMessage = message.value;
+      } 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];
     };
     try {
-      rpcStream = await Promise.race([setupStream(), abortRaceProm.p]);
+      streamCreation = await setupStream();
     } finally {
       cleanUp();
     }
+    const [result, rpcStream] = streamCreation
     const metadata = {
       ...(rpcStream.meta ?? {}),
+      result,
       command: method,
     };
     return {

src/rpc/RPCServer.ts

@@ -31,6 +31,7 @@ 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';
 
@@ -344,7 +345,7 @@ class RPCServer extends EventTarget {
       });
       // Ignore any errors here, it should propagate to the ends of the stream
       void reverseMiddlewareStream.pipeTo(reverseStream).catch(() => {});
-      return middleware.reverse.readable;
+      return [undefined, middleware.reverse.readable];
     };
     this.registerRawStreamHandler(method, rawSteamHandler, timeout);
   }
@@ -565,12 +566,47 @@ class RPCServer extends EventTarget {
         timer.refresh();
       }
       this.logger.info(`Handling stream with method (${method})`);
-      const outputStream = handler(
-        [headerMessage.value, inputStream],
-        rpcStream.cancel,
-        rpcStream.meta,
-        { signal: abortController.signal, timer },
-      );
+      let handlerResult: [JSONValue | undefined, ReadableStream<Uint8Array>];
+      const headerWriter = rpcStream.writable.getWriter();
+      try {
+        handlerResult = handler(
+          [headerMessage.value, inputStream],
+          rpcStream.cancel,
+          rpcStream.meta,
+          { signal: abortController.signal, timer },
+        );
+      } catch (e) {
+        const rpcError: JSONRPCError = {
+          code: e.exitCode ?? sysexits.UNKNOWN,
+          message: e.description ?? '',
+          data: rpcUtils.fromError(e, this.sensitive),
+        };
+        const rpcErrorMessage: JSONRPCResponseError = {
+          jsonrpc: '2.0',
+          error: rpcError,
+          id: null,
+        };
+        await headerWriter.write(Buffer.from(JSON.stringify(rpcErrorMessage)))
+        // clean up and return
+        timer.cancel(cleanupReason);
+        abortController.signal.removeEventListener('abort', handleAbort);
+        graceTimer?.cancel(cleanupReason);
+        abortController.abort(new rpcErrors.ErrorRPCStreamEnded());
+        rpcStream.cancel(Error('TMP header message was an error'));
+        return;
+      }
+      const [leadingResult, outputStream] = handlerResult;
+
+      if (leadingResult !== undefined){
+        // writing leading metadata
+        const leadingMessage: JSONRPCResponseResult = {
+          jsonrpc: "2.0",
+          result: leadingResult,
+          id: null
+        };
+        await headerWriter.write(Buffer.from(JSON.stringify(leadingMessage)));
+      }
+      headerWriter.releaseLock();
       const outputStreamEndProm = outputStream
         .pipeTo(rpcStream.writable)
         .catch(() => {}); // Ignore any errors, we only care that it finished

src/rpc/handlers.ts

@@ -29,7 +29,7 @@ abstract class RawHandler<
     cancel: (reason?: any) => void,
     meta: Record<string, JSONValue> | undefined,
     ctx: ContextTimed,
-  ): ReadableStream<Uint8Array>;
+  ): [JSONValue, ReadableStream<Uint8Array>];
 }
 
 abstract class DuplexHandler<

src/rpc/types.ts

@@ -164,7 +164,7 @@ type HandlerImplementation<I, O> = (
 
 type RawHandlerImplementation = HandlerImplementation<
   [JSONRPCRequest, ReadableStream<Uint8Array>],
-  ReadableStream<Uint8Array>
+  [JSONValue | undefined, ReadableStream<Uint8Array>]
 >;
 
 type DuplexHandlerImplementation<
@@ -264,7 +264,7 @@ type DuplexCallerImplementation<
 type RawCallerImplementation = (
   headerParams: JSONValue,
   ctx?: Partial<ContextTimedInput>,
-) => Promise<RPCStream<Uint8Array, Uint8Array>>;
+) => Promise<RPCStream<Uint8Array, Uint8Array, Record<string, JSONValue> & {result: JSONValue, command: string}>>;
 
 type ConvertDuplexCaller<T> = T extends DuplexCaller<infer I, infer O>
   ? DuplexCallerImplementation<I, O>

tests/rpc/RPC.test.ts

@@ -4,7 +4,7 @@ import type { JSONValue } from '@/types';
 import { TransformStream } from 'stream/web';
 import { fc, testProp } from '@fast-check/jest';
 import Logger, { LogLevel, StreamHandler } from '@matrixai/logger';
-import { sleep } from 'ix/asynciterable/_sleep';
+import * as utils from '@/utils';
 import RPCServer from '@/rpc/RPCServer';
 import RPCClient from '@/rpc/RPCClient';
 import {
@@ -43,10 +43,10 @@ describe('RPC', () => {
       class TestMethod extends RawHandler {
         public handle(
           input: [JSONRPCRequest, ReadableStream<Uint8Array>],
-        ): ReadableStream<Uint8Array> {
+        ): [JSONValue, ReadableStream<Uint8Array>] {
           const [header_, stream] = input;
           header = header_;
-          return stream;
+          return ['some leading data', stream];
         }
       }
       const rpcServer = await RPCServer.createRPCServer({
@@ -89,12 +89,94 @@ describe('RPC', () => {
         id: null,
       };
       expect(header).toStrictEqual(expectedHeader);
+      expect(callerInterface.meta?.result).toBe('some leading data');
       expect(await outputResult).toStrictEqual(inputData);
       await pipeProm;
       await rpcServer.destroy();
       await rpcClient.destroy();
     },
   );
+  test(
+    'RPC communication with raw stream times out waiting for leading message',
+    async () => {
+      const { clientPair, serverPair } = rpcTestUtils.createTapPairs<
+        Uint8Array,
+        Uint8Array
+      >();
+      void (async () => {
+        for await (const _ of serverPair.readable) {
+          // just consume
+        }
+      })();
+
+      const rpcClient = await RPCClient.createRPCClient({
+        manifest: {
+          testMethod: new RawCaller(),
+        },
+        streamFactory: async () => {
+          return {
+            ...clientPair,
+            cancel: () => {},
+          };
+        },
+        logger,
+      });
+
+      await expect(rpcClient.methods.testMethod({
+        hello: 'world',
+      },
+        { timer: 100 },
+        )).rejects.toThrow(rpcErrors.ErrorRPCTimedOut);
+      await rpcClient.destroy();
+    },
+  );
+  test(
+    'RPC communication with raw stream, raw handler throws',
+    async () => {
+      const { clientPair, serverPair } = rpcTestUtils.createTapPairs<
+        Uint8Array,
+        Uint8Array
+      >();
+
+      class TestMethod extends RawHandler {
+        public handle(
+          input: [JSONRPCRequest, ReadableStream<Uint8Array>],
+        ): [JSONValue, ReadableStream<Uint8Array>] {
+          throw Error('some error');
+        }
+      }
+      const rpcServer = await RPCServer.createRPCServer({
+        manifest: {
+          testMethod: new TestMethod({}),
+        },
+        logger,
+      });
+      rpcServer.handleStream({
+        ...serverPair,
+        cancel: () => {},
+      });
+
+      const rpcClient = await RPCClient.createRPCClient({
+        manifest: {
+          testMethod: new RawCaller(),
+        },
+        streamFactory: async () => {
+          return {
+            ...clientPair,
+            cancel: () => {},
+          };
+        },
+        logger,
+      });
+
+      await expect(rpcClient.methods.testMethod({
+        hello: 'world',
+      })).rejects.toThrow(rpcErrors.ErrorPolykeyRemote)
+
+      await rpcServer.destroy();
+      await rpcClient.destroy();
+    },
+  );
   testProp(
     'RPC communication with duplex stream',
     [fc.array(rpcTestUtils.safeJsonValueArb, { minLength: 1 })],
@@ -466,7 +548,7 @@ describe('RPC', () => {
     const writer = callerInterface.writable.getWriter();
     await writer.write({});
     // Allow time to process buffer
-    await sleep(0);
+    await utils.sleep(0);
     await expect(writer.write({})).toReject();
     const reader = callerInterface.readable.getReader();
     await expect(reader.read()).toReject();