feat(middleware): Enhance middleware with timeout and encapsulation features

addievo · amydevs · commit 228f824d42e8 · 2023-10-27T16:15:32.000+11:00
- Import ClientRPCResponseResult and ClientRPCRequestParams from PK.
- Implement timeoutMiddlewareServer and timeoutMiddlewareClient.
- Integrate timeoutMiddleware into defaultMiddleware.
- Fix Jest test issues.
- Rename to RPCResponseResult and RPCRequestParams for clarity.
- Perform lint fixes and Jest tests.

fix: timeouts are now only propagated unidirectionally from the client to the server

[ci-skip]
diff --git a/src/middleware.ts b/src/middleware.ts
@@ -4,7 +4,11 @@ import type {
   JSONRPCResponse,
   JSONRPCResponseResult,
   MiddlewareFactory,
+  JSONValue,
+  JSONRPCRequestMetadata,
+  JSONRPCResponseMetadata,
 } from './types';
+import type { ContextTimed } from '@matrixai/contexts';
 import { TransformStream } from 'stream/web';
 import { JSONParser } from '@streamparser/json';
 import * as utils from './utils';
@@ -75,6 +79,80 @@ function jsonMessageToBinaryStream(): TransformStream<
   });
 }
 
+function timeoutMiddlewareServer(
+  ctx: ContextTimed,
+  _cancel: (reason?: any) => void,
+  _meta: Record<string, JSONValue> | undefined,
+) {
+  const currentTimeout = ctx.timer.delay;
+  // Flags for tracking if the first message has been processed
+  let forwardFirst = true;
+  return {
+    forward: new TransformStream<
+      JSONRPCRequest<JSONRPCRequestMetadata>,
+      JSONRPCRequest<JSONRPCRequestMetadata>
+    >({
+      transform: (chunk, controller) => {
+        controller.enqueue(chunk);
+        if (forwardFirst) {
+          forwardFirst = false;
+          let clientTimeout = chunk.metadata?.timeout;
+          if (clientTimeout === undefined) return;
+          if (clientTimeout === null) clientTimeout = Infinity;
+          if (clientTimeout < currentTimeout) ctx.timer.reset(clientTimeout);
+        }
+      },
+    }),
+    reverse: new TransformStream<
+      JSONRPCResponse<JSONRPCResponseMetadata>,
+      JSONRPCResponse<JSONRPCResponseMetadata>
+    >({
+      transform: (chunk, controller) => {
+        // Passthrough chunk, no need for server to send ctx.timeout
+        controller.enqueue(chunk);
+      },
+    }),
+  };
+}
+
+/**
+ * This adds its own timeout to the forward metadata and updates it's timeout
+ * based on the reverse metadata.
+ * @param ctx
+ * @param _cancel
+ * @param _meta
+ */
+function timeoutMiddlewareClient(
+  ctx: ContextTimed,
+  _cancel: (reason?: any) => void,
+  _meta: Record<string, JSONValue> | undefined,
+) {
+  const currentTimeout = ctx.timer.delay;
+  // Flags for tracking if the first message has been processed
+  let forwardFirst = true;
+  return {
+    forward: new TransformStream<JSONRPCRequest, JSONRPCRequest>({
+      transform: (chunk, controller) => {
+        if (forwardFirst) {
+          forwardFirst = false;
+          if (chunk == null) chunk = { jsonrpc: '2.0', method: '' };
+          if (chunk.metadata == null) chunk.metadata = {};
+          (chunk.metadata as any).timeout = currentTimeout;
+        }
+        controller.enqueue(chunk);
+      },
+    }),
+    reverse: new TransformStream<
+      JSONRPCResponse<JSONRPCResponseMetadata>,
+      JSONRPCResponse<JSONRPCResponseMetadata>
+    >({
+      transform: (chunk, controller) => {
+        controller.enqueue(chunk); // Passthrough chunk, no need for client to set ctx.timeout
+      },
+    }),
+  };
+}
+
 /**
  * This function is a factory for creating a pass-through streamPair. It is used
  * as the default middleware for the middleware wrappers.
@@ -116,12 +194,14 @@ function defaultServerMiddlewareWrapper(
     >();
 
     const middleMiddleware = middlewareFactory(ctx, cancel, meta);
+    const timeoutMiddleware = timeoutMiddlewareServer(ctx, cancel, meta);
 
-    const forwardReadable = inputTransformStream.readable.pipeThrough(
-      middleMiddleware.forward,
-    ); // Usual middleware here
+    const forwardReadable = inputTransformStream.readable
+      .pipeThrough(timeoutMiddleware.forward) // Timeout middleware here
+      .pipeThrough(middleMiddleware.forward); // Usual middleware here
     const reverseReadable = outputTransformStream.readable
       .pipeThrough(middleMiddleware.reverse) // Usual middleware here
+      .pipeThrough(timeoutMiddleware.reverse) // Timeout middleware here
       .pipeThrough(jsonMessageToBinaryStream());
 
     return {
@@ -172,13 +252,15 @@ const defaultClientMiddlewareWrapper = (
       JSONRPCRequest
     >();
 
+    const timeoutMiddleware = timeoutMiddlewareClient(ctx, cancel, meta);
     const middleMiddleware = middleware(ctx, cancel, meta);
     const forwardReadable = inputTransformStream.readable
+      .pipeThrough(timeoutMiddleware.forward)
       .pipeThrough(middleMiddleware.forward) // Usual middleware here
       .pipeThrough(jsonMessageToBinaryStream());
-    const reverseReadable = outputTransformStream.readable.pipeThrough(
-      middleMiddleware.reverse,
-    ); // Usual middleware here
+    const reverseReadable = outputTransformStream.readable
+      .pipeThrough(middleMiddleware.reverse)
+      .pipeThrough(timeoutMiddleware.reverse); // Usual middleware here
 
     return {
       forward: {
@@ -199,4 +281,6 @@ export {
   defaultMiddleware,
   defaultServerMiddlewareWrapper,
   defaultClientMiddlewareWrapper,
+  timeoutMiddlewareClient,
+  timeoutMiddlewareServer,
 };
diff --git a/src/types.ts b/src/types.ts
@@ -38,7 +38,7 @@ type JSONRPCRequestMessage<T extends JSONValue = JSONValue> = {
    *  SHOULD NOT contain fractional parts [2]
    */
   id: string | number | null;
-};
+} & JSONRPCRequestMetadata;
 
 /**
  * This is the JSON RPC notification object. this is used for a request that
@@ -60,7 +60,7 @@ type JSONRPCRequestNotification<T extends JSONValue = JSONValue> = {
    *  This member MAY be omitted.
    */
   params?: T;
-};
+} & JSONRPCRequestMetadata;
 
 /**
  * This is the JSON RPC response result object. It contains the response data for a
@@ -84,7 +84,7 @@ type JSONRPCResponseResult<T extends JSONValue = JSONValue> = {
    *  it MUST be Null.
    */
   id: string | number | null;
-};
+} & JSONRPCResponseMetadata;
 
 /**
  * This is the JSON RPC response Error object. It contains any errors that have
@@ -110,6 +110,29 @@ type JSONRPCResponseError = {
   id: string | number | null;
 };
 
+type ObjectEmpty = NonNullable<unknown>;
+
+// Prevent overwriting the metadata type with `Omit<>`
+type JSONRPCRequestMetadata<T extends Record<string, JSONValue> = ObjectEmpty> =
+  {
+    metadata?: {
+      [Key: string]: JSONValue;
+    } & Partial<{
+      timeout: number | null;
+    }>;
+  } & Omit<T, 'metadata'>;
+
+// Prevent overwriting the metadata type with `Omit<>`
+type JSONRPCResponseMetadata<
+  T extends Record<string, JSONValue> = ObjectEmpty,
+> = {
+  metadata?: {
+    [Key: string]: JSONValue;
+  } & Partial<{
+    timeout: number | null;
+  }>;
+} & Omit<T, 'metadata'>;
+
 /**
  * This is a JSON RPC error object, it encodes the error data for the JSONRPCResponseError object.
  */
@@ -357,6 +380,8 @@ export type {
   JSONRPCRequestNotification,
   JSONRPCResponseResult,
   JSONRPCResponseError,
+  JSONRPCRequestMetadata,
+  JSONRPCResponseMetadata,
   JSONRPCError,
   JSONRPCRequest,
   JSONRPCResponse,
diff --git a/tests/RPC.test.ts b/tests/RPC.test.ts
@@ -870,7 +870,62 @@ describe('RPC', () => {
     },
     { numRuns: 1 },
   );
+  test('RPC server times out using client timeout', async () => {
+    // Setup server and client communication pairs
+    const { clientPair, serverPair } = rpcTestUtils.createTapPairs<
+      Uint8Array,
+      Uint8Array
+    >();
+    const { p: ctxP, resolveP: resolveCtxP } = utils.promise<ContextTimed>();
+    class TestMethod extends UnaryHandler {
+      public handle = async (
+        input: JSONValue,
+        cancel: (reason?: any) => void,
+        meta: Record<string, JSONValue> | undefined,
+        ctx: ContextTimed,
+      ): Promise<JSONValue> => {
+        const abortProm = utils.promise<never>();
+        ctx.signal.addEventListener('abort', () => {
+          resolveCtxP(ctx);
+          abortProm.resolveP(ctx.signal.reason);
+        });
+        throw await abortProm.p;
+      };
+    }
+    // Set up a client and server with matching timeout settings
+    const rpcServer = new RPCServer({
+      logger,
+      idGen,
+      handlerTimeoutTime: 150,
+    });
+    await rpcServer.start({
+      manifest: {
+        testMethod: new TestMethod({}),
+      },
+    });
+    rpcServer.handleStream({
+      ...serverPair,
+      cancel: () => {},
+    });
 
+    const rpcClient = new RPCClient({
+      manifest: {
+        testMethod: new UnaryCaller(),
+      },
+      streamFactory: async () => {
+        return {
+          ...clientPair,
+          cancel: () => {},
+        };
+      },
+      logger,
+      idGen,
+    });
+    await expect(rpcClient.methods.testMethod({}, { timer: 100 })).toReject();
+    await expect(ctxP).resolves.toHaveProperty(['timer', 'delay'], 100);
+
+    await rpcServer.stop({ force: true });
+  });
   testProp(
     'RPC Serializes and Deserializes Error',
     [rpcTestUtils.errorArb(rpcTestUtils.errorArb())],
diff --git a/tests/RPCClient.test.ts b/tests/RPCClient.test.ts
@@ -123,18 +123,20 @@ describe(`${RPCClient.name}`, () => {
     }
     await writable.close();
 
-    const expectedMessages: Array<JSONRPCRequestMessage> = messages.map((v) => {
-      const request: JSONRPCRequestMessage = {
+    const expectedMessages: Array<JSONRPCRequestMessage> = messages.map(
+      (v, i) => ({
         jsonrpc: '2.0',
         method: methodName,
         id: null,
         ...(v.result === undefined ? {} : { params: v.result }),
-      };
-      return request;
-    });
+        ...(i === 0 ? { metadata: { timeout: null } } : {}),
+      }),
+    );
+
     const outputMessages = (await outputResult).map((v) =>
       JSON.parse(v.toString()),
     );
+
     expect(outputMessages).toStrictEqual(expectedMessages);
   });
   testProp(
@@ -171,6 +173,9 @@ describe(`${RPCClient.name}`, () => {
           jsonrpc: '2.0',
           id: null,
           params,
+          metadata: {
+            timeout: null,
+          },
         }),
       );
     },
@@ -207,14 +212,16 @@ describe(`${RPCClient.name}`, () => {
       }
       await writer.close();
       expect(await output).toStrictEqual(message.result);
-      const expectedOutput = params.map((v) =>
+      const expectedOutput = params.map((v, i) =>
         JSON.stringify({
           method: methodName,
           jsonrpc: '2.0',
           id: null,
           params: v,
+          ...(i === 0 ? { metadata: { timeout: null } } : {}),
         }),
       );
+
       expect((await outputResult).map((v) => v.toString())).toStrictEqual(
         expectedOutput,
       );
@@ -249,6 +256,7 @@ describe(`${RPCClient.name}`, () => {
           jsonrpc: '2.0',
           id: null,
           params: params,
+          metadata: { timeout: null },
         }),
       );
     },
@@ -423,19 +431,19 @@ describe(`${RPCClient.name}`, () => {
       }
 
       const expectedMessages: Array<JSONRPCRequestMessage> = messages.map(
-        () => {
-          const request: JSONRPCRequestMessage = {
-            jsonrpc: '2.0',
-            method: methodName,
-            id: null,
-            params: 'one',
-          };
-          return request;
-        },
+        (_, i) => ({
+          jsonrpc: '2.0',
+          method: methodName,
+          id: null,
+          params: 'one',
+          ...(i === 0 ? { metadata: { timeout: null } } : {}),
+        }),
       );
+
       const outputMessages = (await outputResult).map((v) =>
         JSON.parse(v.toString()),
       );
+
       expect(outputMessages).toStrictEqual(expectedMessages);
     },
   );
@@ -527,6 +535,7 @@ describe(`${RPCClient.name}`, () => {
           jsonrpc: '2.0',
           id: null,
           params,
+          metadata: { timeout: null },
         }),
       );
     },
@@ -562,12 +571,13 @@ describe(`${RPCClient.name}`, () => {
       }
       expect(await output).toStrictEqual(message.result);
       await writer.close();
-      const expectedOutput = params.map((v) =>
+      const expectedOutput = params.map((v, i) =>
         JSON.stringify({
           method: 'client',
           jsonrpc: '2.0',
           id: null,
           params: v,
+          ...(i === 0 ? { metadata: { timeout: null } } : {}),
         }),
       );
       expect((await outputResult).map((v) => v.toString())).toStrictEqual(
@@ -603,6 +613,7 @@ describe(`${RPCClient.name}`, () => {
           jsonrpc: '2.0',
           id: null,
           params: params,
+          metadata: { timeout: null },
         }),
       );
     },
diff --git a/tests/RPCServer.test.ts b/tests/RPCServer.test.ts
diff --git a/tests/middleware.test.ts b/tests/middleware.test.ts
