feat: negative timeoutTime paramaters will now throw errors

chore: add tests to test negative timeout values. fix: 0 `timeoutTime` value no longer throws in `RPCClient` fix: `RPCServer.start` now throws when passed a handler with a negative `timeout` [ci-skip]
MatrixAI · Oct 30, 2023 · 1ebf53c · 1ebf53c
1 parent 669c00b
commit 1ebf53c
Show file tree

Hide file tree

Showing 5 changed files with 153 additions and 63 deletions.
diff --git a/src/RPCClient.ts b/src/RPCClient.ts
@@ -103,6 +103,9 @@ class RPCClient<M extends ClientManifest> {
     idGen?: IdGen;
     toError?: ToError;
   }) {
+    if (timeoutTime < 0) {
+      throw new errors.ErrorRPCInvalidTimeout();
+    }
     this.idGen = idGen;
     this.callerTypes = utils.getHandlerTypes(manifest);
     this.streamFactory = streamFactory;

diff --git a/src/RPCServer.ts b/src/RPCServer.ts
@@ -106,6 +106,9 @@ class RPCServer {
     fromError?: FromError;
     replacer?: (key: string, value: any) => any;
   }) {
+    if (timeoutTime < 0) {
+      throw new errors.ErrorRPCInvalidTimeout();
+    }
     this.idGen = idGen;
     this.middlewareFactory = middlewareFactory;
     this.timeoutTime = timeoutTime;
@@ -126,58 +129,68 @@ class RPCServer {
     manifest: ServerManifest;
   }): Promise<void> {
     this.logger.info(`Start ${this.constructor.name}`);
-    for (const [key, manifestItem] of Object.entries(manifest)) {
-      if (manifestItem instanceof RawHandler) {
-        this.registerRawStreamHandler(
-          key,
-          manifestItem.handle,
-          manifestItem.timeout,
-        );
-        continue;
-      }
-      if (manifestItem instanceof DuplexHandler) {
-        this.registerDuplexStreamHandler(
-          key,
-          // Bind the `this` to the generator handler to make the container available
-          manifestItem.handle.bind(manifestItem),
-          manifestItem.timeout,
-        );
-        continue;
-      }
-      if (manifestItem instanceof ServerHandler) {
-        this.registerServerStreamHandler(
-          key,
-          // Bind the `this` to the generator handler to make the container available
-          manifestItem.handle.bind(manifestItem),
-          manifestItem.timeout,
-        );
-        continue;
-      }
-      if (manifestItem instanceof ClientHandler) {
-        this.registerClientStreamHandler(
-          key,
-          manifestItem.handle,
-          manifestItem.timeout,
-        );
-        continue;
-      }
-      if (manifestItem instanceof ClientHandler) {
-        this.registerClientStreamHandler(
-          key,
-          manifestItem.handle,
-          manifestItem.timeout,
-        );
-        continue;
-      }
-      if (manifestItem instanceof UnaryHandler) {
-        this.registerUnaryHandler(
-          key,
-          manifestItem.handle,
-          manifestItem.timeout,
-        );
-        continue;
+    try {
+      for (const [key, manifestItem] of Object.entries(manifest)) {
+        if (manifestItem.timeout != null && manifestItem.timeout < 0) {
+          throw new errors.ErrorRPCInvalidHandlerTimeout();
+        }
+        if (manifestItem instanceof RawHandler) {
+          this.registerRawStreamHandler(
+            key,
+            manifestItem.handle,
+            manifestItem.timeout,
+          );
+          continue;
+        }
+        if (manifestItem instanceof DuplexHandler) {
+          this.registerDuplexStreamHandler(
+            key,
+            // Bind the `this` to the generator handler to make the container available
+            manifestItem.handle.bind(manifestItem),
+            manifestItem.timeout,
+          );
+          continue;
+        }
+        if (manifestItem instanceof ServerHandler) {
+          this.registerServerStreamHandler(
+            key,
+            // Bind the `this` to the generator handler to make the container available
+            manifestItem.handle.bind(manifestItem),
+            manifestItem.timeout,
+          );
+          continue;
+        }
+        if (manifestItem instanceof ClientHandler) {
+          this.registerClientStreamHandler(
+            key,
+            manifestItem.handle,
+            manifestItem.timeout,
+          );
+          continue;
+        }
+        if (manifestItem instanceof ClientHandler) {
+          this.registerClientStreamHandler(
+            key,
+            manifestItem.handle,
+            manifestItem.timeout,
+          );
+          continue;
+        }
+        if (manifestItem instanceof UnaryHandler) {
+          this.registerUnaryHandler(
+            key,
+            manifestItem.handle,
+            manifestItem.timeout,
+          );
+          continue;
+        }
+        utils.never();
       }
-      utils.never();
+    } catch (e) {
+      // No need to clean up streams, as streams can only be handled after RPCServer has been started.
+      this.handlerMap.clear();
+      this.defaultTimeoutMap.clear();
+      throw e;
     }
     this.logger.info(`Started ${this.constructor.name}`);
   }

diff --git a/src/errors.ts b/src/errors.ts
@@ -27,6 +27,14 @@ class ErrorRPCCallerFailed<T> extends ErrorRPC<T> {
   static description = 'Failed to call stream';
 }
 
+class ErrorRPCInvalidTimeout<T> extends ErrorRPC<T> {
+  static description = 'Invalid timeout provided';
+}
+
+class ErrorRPCInvalidHandlerTimeout<T> extends ErrorRPC<T> {
+  static description = 'Invalid handler timeout provided';
+}
+
 abstract class ErrorRPCProtocol<T> extends ErrorRPC<T> {
   static error = 'RPC Protocol Error';
   code: number;
@@ -257,6 +265,8 @@ export {
   ErrorRPCConnectionLocal,
   ErrorRPCConnectionPeer,
   ErrorRPCConnectionKeepAliveTimeOut,
+  ErrorRPCInvalidTimeout,
+  ErrorRPCInvalidHandlerTimeout,
   ErrorRPCConnectionInternal,
   ErrorMissingHeader,
   ErrorHandlerAborted,

diff --git a/tests/RPCClient.test.ts b/tests/RPCClient.test.ts
@@ -1148,6 +1148,28 @@ describe(`${RPCClient.name}`, () => {
       },
       { numRuns: 5 },
     );
+    testProp(
+      'RPCClient constructor should throw when passed a negative timeoutTime',
+      [fc.integer({ max: -1 })],
+      async (timeoutTime) => {
+        const streamPair: RPCStream<Uint8Array, Uint8Array> = {
+          cancel: () => {},
+          meta: undefined,
+          readable: new ReadableStream(),
+          writable: new WritableStream(),
+        };
+        const constructorF = () =>
+          new RPCClient({
+            timeoutTime,
+            streamFactory: () => Promise.resolve(streamPair),
+            manifest: {},
+            logger,
+            idGen,
+          });
+
+        expect(constructorF).toThrowError(rpcErrors.ErrorRPCInvalidTimeout);
+      },
+    );
     testProp(
       'Check that ctx is provided to the middleware and that the middleware can reset the timer',
       [specificMessageArb],
@@ -1220,12 +1242,9 @@ describe(`${RPCClient.name}`, () => {
           timeoutTime: higherTimeoutTime,
         });
 
-        await rpcClient.duplexStreamCaller<JSONRPCParams, JSONRPCResult>(
-          methodName,
-          {
-            timer: lowerTimeoutTime,
-          },
-        );
+        await rpcClient.duplexStreamCaller<JSONRPCParams, JSONRPCResult>(methodName, {
+          timer: lowerTimeoutTime,
+        });
 
         const ctx = await ctxP;
         expect(ctx.timer.delay).toBe(lowerTimeoutTime);
@@ -1255,12 +1274,9 @@ describe(`${RPCClient.name}`, () => {
           timeoutTime: lowerTimeoutTime,
         });
 
-        await rpcClient.duplexStreamCaller<JSONRPCParams, JSONRPCResult>(
-          methodName,
-          {
-            timer: higherTimeoutTime,
-          },
-        );
+        await rpcClient.duplexStreamCaller<JSONRPCParams, JSONRPCResult>(methodName, {
+          timer: higherTimeoutTime,
+        });
 
         const ctx = await ctxP;
         expect(ctx.timer.delay).toBe(higherTimeoutTime);

diff --git a/tests/RPCServer.test.ts b/tests/RPCServer.test.ts
@@ -1082,6 +1082,54 @@ describe(`${RPCServer.name}`, () => {
     await expect(ctx.timer).toReject();
     await rpcServer.stop({ force: true });
   });
+  testProp(
+    'RPCServer constructor should throw when passed a negative timeoutTime',
+    [fc.integer({ max: -1 })],
+    async (timeoutTime) => {
+      const constructorF = () =>
+        new RPCServer({
+          timeoutTime,
+          logger,
+          idGen,
+        });
+
+      expect(constructorF).toThrowError(rpcErrors.ErrorRPCInvalidTimeout);
+    },
+  );
+  testProp(
+    'RPCServer.start should throw when passed a handler with negative timeout',
+    [fc.integer({ max: -1 })],
+    async (timeoutTime) => {
+      const waitProm = promise();
+      const ctxLongProm = promise<ContextTimed>();
+
+      class TestMethodArbitraryTimeout extends UnaryHandler {
+        timeout = timeoutTime;
+        public handle = async (
+          input: JSONRPCParams,
+          _cancel,
+          _meta,
+          ctx_,
+        ): Promise<JSONRPCResult> => {
+          ctxLongProm.resolveP(ctx_);
+          await waitProm.p;
+          return input;
+        };
+      }
+      const rpcServer = new RPCServer({
+        logger,
+        idGen,
+      });
+
+      await expect(
+        rpcServer.start({
+          manifest: {
+            testArbitrary: new TestMethodArbitraryTimeout({}),
+          },
+        }),
+      ).rejects.toBeInstanceOf(rpcErrors.ErrorRPCInvalidHandlerTimeout);
+    },
+  );
   testProp(
     'middleware can update timeout timer',
     [specificMessageArb],