feat(reactive-rpc): 🎸 integrate WebSocket into RPC server

streamich · streamich · commit eed1662eb213 · 2023-12-28T21:14:25.000+01:00
diff --git a/src/reactive-rpc/server/http1/Http1Server.ts b/src/reactive-rpc/server/http1/Http1Server.ts
@@ -8,8 +8,8 @@ import {Router} from '../../../util/router';
 import {Printable} from '../../../util/print/types';
 import {printTree} from '../../../util/print/printTree';
 import {PayloadTooLarge} from './errors';
-import {findTokenInText} from './util';
-import {Http1ConnectionContext} from './context';
+import {findTokenInText, setCodecs} from './util';
+import {Http1ConnectionContext, WsConnectionContext} from './context';
 import {RpcCodecs} from '../../common/codec/RpcCodecs';
 import {Codecs} from '../../../json-pack/codecs/Codecs';
 import {RpcMessageCodecs} from '../../common/codec/RpcMessageCodecs';
@@ -43,9 +43,10 @@ export interface Http1EndpointDefinition {
 
 export interface WsEndpointDefinition {
   path: string;
-  maxPayload?: number;
+  maxIncomingMessage?: number;
+  maxOutgoingBackpressure?: number;
   onUpgrade?(req: http.IncomingMessage, connection: WsServerConnection): void;
-  onConnect(connection: WsServerConnection, req: http.IncomingMessage): void;
+  onConnect(ctx: WsConnectionContext, req: http.IncomingMessage): void;
 }
 
 export interface Http1ServerOpts {
@@ -150,7 +151,7 @@ export class Http1Server implements Printable {
       );
       const headers = req.headers;
       const contentType = headers['content-type'];
-      if (typeof contentType === 'string') ctx.setCodecs(contentType, codecs);
+      if (typeof contentType === 'string') setCodecs(ctx, contentType, codecs);
       const handler = match.data.handler;
       await handler(ctx);
     } catch (error) {
@@ -165,7 +166,15 @@ export class Http1Server implements Printable {
   protected wsMatcher: RouteMatcher<WsEndpointDefinition> = () => undefined;
 
   private readonly onWsUpgrade = (req: http.IncomingMessage, socket: net.Socket, head: Buffer) => {
-    const route = req.url || '';
+    const url = req.url ?? '';
+    const queryStartIndex = url.indexOf('?');
+    let path = url;
+    let query = '';
+    if (queryStartIndex >= 0) {
+      path = url.slice(0, queryStartIndex);
+      query = url.slice(queryStartIndex + 1);
+    }
+    const route = (req.method || '') + path;
     const match = this.wsMatcher(route);
     if (!match) {
       socket.end();
@@ -174,14 +183,37 @@ export class Http1Server implements Printable {
     const def = match.data;
     const headers = req.headers;
     const connection = new WsServerConnection(this.wsEncoder, socket as net.Socket, head);
+    connection.maxIncomingMessage = def.maxIncomingMessage ?? 2 * 1024 * 1024;
+    connection.maxBackpressure = def.maxOutgoingBackpressure ?? 2 * 1024 * 1024;
     if (def.onUpgrade) def.onUpgrade(req, connection);
     else {
       const secWebSocketKey = headers['sec-websocket-key'] ?? '';
       const secWebSocketProtocol = headers['sec-websocket-protocol'] ?? '';
       const secWebSocketExtensions = headers['sec-websocket-extensions'] ?? '';
       connection.upgrade(secWebSocketKey, secWebSocketProtocol, secWebSocketExtensions);
     }
-    def.onConnect(connection, req);
+    const codecs = this.codecs;
+    const ip = this.findIp(req);
+    const token = this.findToken(req);
+    const ctx = new WsConnectionContext(
+      connection,
+      path,
+      query,
+      ip,
+      token,
+      match.params,
+      new NullObject(),
+      codecs.value.json,
+      codecs.value.json,
+      codecs.messages.compact,
+    );
+    const contentType = headers['content-type'];
+    if (typeof contentType === 'string') setCodecs(ctx, contentType, codecs);
+    else {
+      const secWebSocketProtocol = headers['sec-websocket-protocol'] ?? '';
+      if (typeof secWebSocketProtocol === 'string') setCodecs(ctx, secWebSocketProtocol, codecs);
+    }
+    def.onConnect(ctx, req);
   };
 
   public ws(def: WsEndpointDefinition): void {
@@ -209,23 +241,19 @@ export class Http1Server implements Printable {
    */
   public findToken(req: http.IncomingMessage): string {
     let token: string = '';
-    let text: string = '';
     const headers = req.headers;
     let header: string | string[] | undefined;
     header = headers['authorization'];
-    text = typeof header === 'string' ? header : header?.[0] ?? '';
-    if (text) token = findTokenInText(text);
+    if (typeof header === 'string') token = findTokenInText(header);
     if (token) return token;
-    text = req.url || '';
-    if (text) token = findTokenInText(text);
+    const url = req.url;
+    if (typeof url === 'string') token = findTokenInText(url);
     if (token) return token;
     header = headers['cookie'];
-    text = typeof header === 'string' ? header : header?.[0] ?? '';
-    if (text) token = findTokenInText(text);
+    if (typeof header === 'string') token = findTokenInText(header);
     if (token) return token;
     header = headers['sec-websocket-protocol'];
-    text = typeof header === 'string' ? header : header?.[0] ?? '';
-    if (text) token = findTokenInText(text);
+    if (typeof header === 'string') token = findTokenInText(header);
     return token;
   }
 
diff --git a/src/reactive-rpc/server/http1/RpcServer.ts b/src/reactive-rpc/server/http1/RpcServer.ts
@@ -3,8 +3,8 @@ import {Printable} from '../../../util/print/types';
 import {printTree} from '../../../util/print/printTree';
 import {Http1Server} from './Http1Server';
 import {RpcError} from '../../common/rpc/caller';
-import {IncomingBatchMessage, RpcMessageBatchProcessor} from '../../common';
-import {ConnectionContext} from './context';
+import {IncomingBatchMessage, ReactiveRpcClientMessage, ReactiveRpcMessage, RpcMessageBatchProcessor, RpcMessageStreamProcessor} from '../../common';
+import {ConnectionContext, WsConnectionContext} from './context';
 import type {RpcCaller} from '../../common/rpc/caller/RpcCaller';
 import type {ServerLogger} from './types';
 
@@ -105,6 +105,61 @@ export class RpcServer implements Printable {
     });
   }
 
+  public enableWsRpc(path: string = '/rpc'): void {
+    const opts = this.opts;
+    const logger = opts.logger ?? console;
+    const caller = opts.caller;
+    this.http1.ws({
+      path,
+      maxIncomingMessage: 2 * 1024 * 1024,
+      maxOutgoingBackpressure: 2 * 1024 * 1024,
+      onConnect: (ctx: WsConnectionContext, req: http.IncomingMessage) => {
+        const connection = ctx.connection;
+        const reqCodec = ctx.reqCodec;
+        const resCodec = ctx.resCodec;
+        const msgCodec = ctx.msgCodec;
+        const encoder = resCodec.encoder;
+        const rpc = new RpcMessageStreamProcessor({
+          caller,
+          send: (messages: ReactiveRpcMessage[]) => {
+            try {
+              const writer = encoder.writer;
+              writer.reset();
+              msgCodec.encodeBatch(resCodec, messages);
+              const encoded = writer.flush();
+              connection.sendBinMsg(encoded);
+            } catch (error) {
+              logger.error('WS_SEND', error, {messages});
+              connection.close();
+            }
+          },
+          bufferSize: 1,
+          bufferTime: 0,
+        });
+        connection.onmessage = (uint8: Uint8Array, isUtf8: boolean) => {
+          let messages: ReactiveRpcClientMessage[];
+          try {
+            messages = msgCodec.decodeBatch(reqCodec, uint8) as ReactiveRpcClientMessage[]; 
+          } catch (error) {
+            logger.error('RX_RPC_DECODING', error, {codec: reqCodec.id, buf: Buffer.from(uint8).toString('base64')});
+            connection.close();
+            return;
+          }
+          try {
+            rpc.onMessages(messages, ctx);
+          } catch (error) {
+            logger.error('RX_RPC_PROCESSING', error, messages!);
+            connection.close();
+            return;
+          }
+        };
+        connection.onclose = (code: number, reason: string) => {   
+          rpc.stop();
+        };
+      },
+    });
+  }
+
   public enableDefaults(): void {
     // this.enableCors();
     this.enableHttpPing();
diff --git a/src/reactive-rpc/server/http1/context.ts b/src/reactive-rpc/server/http1/context.ts
@@ -3,9 +3,7 @@ import {listToUint8} from '../../../util/buffers/concat';
 import type * as http from 'http';
 import type {JsonValueCodec} from '../../../json-pack/codecs/types';
 import type {RpcMessageCodec} from '../../common/codec/types';
-import type {RpcCodecs} from '../../common/codec/RpcCodecs';
-
-const REGEX_CODECS_SPECIFIER = /rpc\.(\w{0,32})\.(\w{0,32})\.(\w{0,32})(?:\-(\w{0,32}))?/;
+import type {WsServerConnection} from '../ws/server/WsServerConnection';
 
 export interface ConnectionContext<Meta = Record<string, unknown>> {
   path: string;
@@ -34,67 +32,24 @@ export class Http1ConnectionContext<Meta = Record<string, unknown>> implements C
     public msgCodec: RpcMessageCodec,
   ) {}
 
-  /**
-   * @param specifier A string which may contain a codec specifier. For example:
-   *  - `rpc.rx.compact.cbor` for Rx-RPC with compact messages and CBOR values.
-   *  - `rpc.json2.verbose.json` for JSON-RPC 2.0 with verbose messages encoded as JSON.
-   */
-  public setCodecs(specifier: string, codecs: RpcCodecs): void {
-    const match = REGEX_CODECS_SPECIFIER.exec(specifier);
-    if (!match) return;
-    const [, protocol, messageFormat, request, response] = match;
-    switch (protocol) {
-      case 'rx': {
-        switch (messageFormat) {
-          case 'compact': {
-            this.msgCodec = codecs.messages.compact;
-            break;
-          }
-          case 'binary': {
-            this.msgCodec = codecs.messages.binary;
-            break;
-          }
-        }
-        break;
-      }
-      case 'json2': {
-        this.msgCodec = codecs.messages.jsonRpc2;
-        break;
-      }
-    }
-    switch (request) {
-      case 'cbor': {
-        this.resCodec = this.reqCodec = codecs.value.cbor;
-        break;
-      }
-      case 'json': {
-        this.resCodec = this.reqCodec = codecs.value.json;
-        break;
-      }
-      case 'msgpack': {
-        this.resCodec = this.reqCodec = codecs.value.msgpack;
-        break;
-      }
-    }
-    switch (response) {
-      case 'cbor': {
-        this.resCodec = codecs.value.cbor;
-        break;
-      }
-      case 'json': {
-        this.resCodec = codecs.value.json;
-        break;
-      }
-      case 'msgpack': {
-        this.resCodec = codecs.value.msgpack;
-        break;
-      }
-    }
-  }
-
   public async body(maxPayload: number): Promise<Uint8Array> {
     const list = await getBody(this.req, maxPayload);
     const bodyUint8 = listToUint8(list);
     return bodyUint8;
   }
 }
+
+export class WsConnectionContext<Meta = Record<string, unknown>> implements ConnectionContext<Meta> {
+  constructor(
+    public readonly connection: WsServerConnection,
+    public path: string,
+    public query: string,
+    public readonly ip: string,
+    public token: string,
+    public readonly params: string[] | null,
+    public readonly meta: Meta,
+    public reqCodec: JsonValueCodec,
+    public resCodec: JsonValueCodec,
+    public msgCodec: RpcMessageCodec,
+  ) {}
+}
diff --git a/src/reactive-rpc/server/http1/util.ts b/src/reactive-rpc/server/http1/util.ts
@@ -1,5 +1,7 @@
 import {PayloadTooLarge} from './errors';
+import type {ConnectionContext} from './context';
 import type * as http from 'http';
+import type {RpcCodecs} from '../../common/codec/RpcCodecs';
 
 export const getBody = (request: http.IncomingMessage, max: number): Promise<Buffer[]> => {
   return new Promise<Buffer[]>((resolve, reject) => {
@@ -32,3 +34,64 @@ export const findTokenInText = (text: string): string => {
   if (!match) return '';
   return match[1] || '';
 };
+
+
+const REGEX_CODECS_SPECIFIER = /rpc\.(\w{0,32})\.(\w{0,32})\.(\w{0,32})(?:\-(\w{0,32}))?/;
+
+/**
+ * @param specifier A string which may contain a codec specifier. For example:
+ *  - `rpc.rx.compact.cbor` for Rx-RPC with compact messages and CBOR values.
+ *  - `rpc.json2.verbose.json` for JSON-RPC 2.0 with verbose messages encoded as JSON.
+ */
+export const setCodecs = (ctx: ConnectionContext, specifier: string, codecs: RpcCodecs): void => {
+  const match = REGEX_CODECS_SPECIFIER.exec(specifier);
+  if (!match) return;
+  const [, protocol, messageFormat, request, response] = match;
+  switch (protocol) {
+    case 'rx': {
+      switch (messageFormat) {
+        case 'compact': {
+          ctx.msgCodec = codecs.messages.compact;
+          break;
+        }
+        case 'binary': {
+          ctx.msgCodec = codecs.messages.binary;
+          break;
+        }
+      }
+      break;
+    }
+    case 'json2': {
+      ctx.msgCodec = codecs.messages.jsonRpc2;
+      break;
+    }
+  }
+  switch (request) {
+    case 'cbor': {
+      ctx.resCodec = ctx.reqCodec = codecs.value.cbor;
+      break;
+    }
+    case 'json': {
+      ctx.resCodec = ctx.reqCodec = codecs.value.json;
+      break;
+    }
+    case 'msgpack': {
+      ctx.resCodec = ctx.reqCodec = codecs.value.msgpack;
+      break;
+    }
+  }
+  switch (response) {
+    case 'cbor': {
+      ctx.resCodec = codecs.value.cbor;
+      break;
+    }
+    case 'json': {
+      ctx.resCodec = codecs.value.json;
+      break;
+    }
+    case 'msgpack': {
+      ctx.resCodec = codecs.value.msgpack;
+      break;
+    }
+  }
+};
diff --git a/src/reactive-rpc/server/ws/server/WsServerConnection.ts b/src/reactive-rpc/server/ws/server/WsServerConnection.ts
