feat: static middleware for server

- Related #500 [ci skip]
MatrixAI · Feb 7, 2023 · d53f1ff · d53f1ff
1 parent f78a7f9
commit d53f1ff
Show file tree

Hide file tree

Showing 8 changed files with 160 additions and 97 deletions.
diff --git a/src/RPC/RPCClient.ts b/src/RPC/RPCClient.ts
@@ -299,13 +299,20 @@ class RPCClient<M extends Manifest> {
   }
 
   protected middleware: Array<
-    MiddlewareFactory<JsonRpcRequest<JSONValue>, JsonRpcResponse<JSONValue>>
+    MiddlewareFactory<
+      JsonRpcRequest<JSONValue>,
+      JsonRpcRequest<JSONValue>,
+      JsonRpcResponse<JSONValue>,
+      JsonRpcResponse<JSONValue>
+    >
   > = [];
 
   @ready(new rpcErrors.ErrorRpcDestroyed())
   public registerMiddleware(
     middlewareFactory: MiddlewareFactory<
       JsonRpcRequest<JSONValue>,
+      JsonRpcRequest<JSONValue>,
+      JsonRpcResponse<JSONValue>,
       JsonRpcResponse<JSONValue>
     >,
   ) {

diff --git a/src/RPC/RPCServer.ts b/src/RPC/RPCServer.ts
@@ -31,16 +31,24 @@ class RPCServer {
   static async createRPCServer({
     manifest,
     container,
+    middleware = rpcUtils.defaultMiddlewareWrapper(),
     logger = new Logger(this.name),
   }: {
     manifest: Manifest;
     container: POJO;
+    middleware?: MiddlewareFactory<
+      JsonRpcRequest<JSONValue>,
+      Uint8Array,
+      Uint8Array,
+      JsonRpcResponseResult<JSONValue>
+    >;
     logger?: Logger;
   }): Promise<RPCServer> {
     logger.info(`Creating ${this.name}`);
     const rpcServer = new this({
       manifest,
       container,
+      middleware,
       logger,
     });
     logger.info(`Created ${this.name}`);
@@ -53,14 +61,27 @@ class RPCServer {
   protected handlerMap: Map<string, RawDuplexStreamHandler> = new Map();
   protected activeStreams: Set<PromiseCancellable<void>> = new Set();
   protected events: EventTarget = new EventTarget();
+  protected middleware: MiddlewareFactory<
+    JsonRpcRequest<JSONValue>,
+    Uint8Array,
+    Uint8Array,
+    JsonRpcResponseResult<JSONValue>
+  >;
 
   public constructor({
     manifest,
     container,
+    middleware,
     logger,
   }: {
     manifest: Manifest;
     container: POJO;
+    middleware: MiddlewareFactory<
+      JsonRpcRequest<JSONValue>,
+      Uint8Array,
+      Uint8Array,
+      JsonRpcResponseResult<JSONValue>
+    >;
     logger: Logger;
   }) {
     for (const [key, manifestItem] of Object.entries(manifest)) {
@@ -85,6 +106,7 @@ class RPCServer {
       }
     }
     this.container = container;
+    this.middleware = middleware;
     this.logger = logger;
   }
 
@@ -120,23 +142,10 @@ class RPCServer {
       connectionInfo,
       ctx,
     ) => {
-      // Middleware
-      const outputTransformStream = new rpcUtils.JsonMessageToJsonStream();
-      const outputReadableSteam = outputTransformStream.readable;
-      let forwardStream = input.pipeThrough(
-        new rpcUtils.JsonToJsonMessageStream(
-          rpcUtils.parseJsonRpcRequest,
-          undefined,
-          header,
-        ),
-      );
-      let reverseStream = outputTransformStream.writable;
-      for (const middlewareFactory of this.middleware) {
-        const middleware = middlewareFactory();
-        forwardStream = forwardStream.pipeThrough(middleware.forward);
-        void middleware.reverse.readable.pipeTo(reverseStream).catch(() => {});
-        reverseStream = middleware.reverse.writable;
-      }
+      // Setting up middleware
+      const middleware = this.middleware(header);
+      const forwardStream = input.pipeThrough(middleware.forward);
+      const reverseStream = middleware.reverse.writable;
       const events = this.events;
       const outputGen = async function* (): AsyncGenerator<
         JsonRpcResponse<JSONValue>
@@ -206,7 +215,7 @@ class RPCServer {
       });
       void reverseMiddlewareStream.pipeTo(reverseStream).catch(() => {});
 
-      return outputReadableSteam;
+      return middleware.reverse.readable;
     };
 
     this.registerRawStreamHandler(method, rawSteamHandler);
@@ -344,25 +353,6 @@ class RPCServer {
   ) {
     this.events.removeEventListener(type, callback, options);
   }
-
-  protected middleware: Array<
-    MiddlewareFactory<JsonRpcRequest<JSONValue>, JsonRpcResponse<JSONValue>>
-  > = [];
-
-  @ready(new rpcErrors.ErrorRpcDestroyed())
-  public registerMiddleware(
-    middlewareFactory: MiddlewareFactory<
-      JsonRpcRequest<JSONValue>,
-      JsonRpcResponse<JSONValue>
-    >,
-  ) {
-    this.middleware.push(middlewareFactory);
-  }
-
-  @ready(new rpcErrors.ErrorRpcDestroyed())
-  public clearMiddleware() {
-    this.middleware = [];
-  }
 }
 
 export default RPCServer;
diff --git a/src/RPC/types.ts b/src/RPC/types.ts
@@ -135,9 +135,11 @@ type StreamPairCreateCallback = () => Promise<
   ReadableWritablePair<Uint8Array, Uint8Array>
 >;
 
-type MiddlewareFactory<F, R> = () => {
-  forward: ReadableWritablePair<F, F>;
-  reverse: ReadableWritablePair<R, R>;
+type MiddlewareFactory<FR, FW, RR, RW> = (
+  header?: JsonRpcRequest<JSONValue>,
+) => {
+  forward: ReadableWritablePair<FR, FW>;
+  reverse: ReadableWritablePair<RR, RW>;
 };
 
 type DuplexStreamCaller<

diff --git a/src/RPC/utils.ts b/src/RPC/utils.ts
@@ -13,6 +13,7 @@ import type {
   JsonRpcResponseResult,
   JsonRpcRequest,
   JsonRpcResponse,
+  MiddlewareFactory,
 } from 'RPC/types';
 import type { JSONValue } from '../types';
 import type { JsonValue } from 'fast-check';
@@ -637,6 +638,59 @@ function getHandlerTypes(manifest: Manifest): Record<string, HandlerType> {
   return out;
 }
 
+const defaultMiddleware: MiddlewareFactory<
+  JsonRpcRequest<JSONValue>,
+  JsonRpcRequest<JSONValue>,
+  JsonRpcResponse<JSONValue>,
+  JsonRpcResponse<JSONValue>
+> = () => {
+  return {
+    forward: new TransformStream(),
+    reverse: new TransformStream(),
+  };
+};
+
+const defaultMiddlewareWrapper = (
+  middleware: MiddlewareFactory<
+    JsonRpcRequest<JSONValue>,
+    JsonRpcRequest<JSONValue>,
+    JsonRpcResponse<JSONValue>,
+    JsonRpcResponse<JSONValue>
+  > = defaultMiddleware,
+) => {
+  return (header: JsonRpcRequest<JSONValue>) => {
+    const inputTransformStream = new JsonToJsonMessageStream(
+      parseJsonRpcRequest,
+      undefined,
+      header,
+    );
+    const outputTransformStream = new TransformStream<
+      JsonRpcResponseResult<JSONValue>,
+      JsonRpcResponseResult<JSONValue>
+    >();
+
+    const middleMiddleware = middleware(header);
+
+    const forwardReadable = inputTransformStream.readable.pipeThrough(
+      middleMiddleware.forward,
+    ); // Usual middleware here
+    const reverseReadable = outputTransformStream.readable
+      .pipeThrough(middleMiddleware.reverse) // Usual middleware here
+      .pipeThrough(new JsonMessageToJsonStream());
+
+    return {
+      forward: {
+        readable: forwardReadable,
+        writable: inputTransformStream.writable,
+      },
+      reverse: {
+        readable: reverseReadable,
+        writable: outputTransformStream.writable,
+      },
+    };
+  };
+};
+
 export {
   JsonToJsonMessageStream,
   JsonMessageToJsonStream,
@@ -657,4 +711,6 @@ export {
   QueueMergingTransformStream,
   extractFirstMessageTransform,
   getHandlerTypes,
+  defaultMiddleware,
+  defaultMiddlewareWrapper,
 };
diff --git a/src/clientRPC/utils.ts b/src/clientRPC/utils.ts
@@ -77,6 +77,8 @@ function authenticationMiddlewareServer(
   keyRing: KeyRing,
 ): MiddlewareFactory<
   JsonRpcRequest<ClientDataAndMetadata<JSONValue>>,
+  JsonRpcRequest<ClientDataAndMetadata<JSONValue>>,
+  JsonRpcResponse<ClientDataAndMetadata<JSONValue>>,
   JsonRpcResponse<ClientDataAndMetadata<JSONValue>>
 > {
   return () => {
@@ -127,6 +129,8 @@ function authenticationMiddlewareClient(
   session: Session,
 ): MiddlewareFactory<
   JsonRpcRequest<ClientDataAndMetadata<JSONValue>>,
+  JsonRpcRequest<ClientDataAndMetadata<JSONValue>>,
+  JsonRpcResponse<ClientDataAndMetadata<JSONValue>>,
   JsonRpcResponse<ClientDataAndMetadata<JSONValue>>
 > {
   return () => {