fix: WebSocketServer can now update TLS certs while running

* Related #540 [ci skip]
MatrixAI · Aug 18, 2023 · f7d044b · f7d044b
1 parent 919f3f4
commit f7d044b
Show file tree

Hide file tree

Showing 4 changed files with 71 additions and 58 deletions.
diff --git a/src/PolykeyAgent.ts b/src/PolykeyAgent.ts
@@ -739,10 +739,7 @@ class PolykeyAgent {
             keyPrivatePem: keysUtils.privateKeyToPEM(data.keyPair.privateKey),
             certChainPem: await this.certManager.getCertPEMsChainPEM(),
           };
-          // FIXME: Can we even support updating TLS config anymore?
-          //  We would need to shut down the Websocket server and re-create it with the new config.
-          //  Right now graceful shutdown is not supported.
-          // this.grpcServerClient.setTLSConfig(tlsConfig);
+          this.webSocketServerClient.setTlsConfig(tlsConfig);
           this.nodeConnectionManager.updateTlsConfig(tlsConfig);
           this.logger.info(`${KeyRing.name} change propagated`);
         },

diff --git a/src/websockets/WebSocketServer.ts b/src/websockets/WebSocketServer.ts
@@ -1,5 +1,6 @@
 import type { TLSConfig } from '../network/types';
 import type { IncomingMessage, ServerResponse } from 'http';
+import type tls from 'tls';
 import https from 'https';
 import { startStop, status } from '@matrixai/async-init';
 import Logger from '@matrixai/logger';
@@ -24,7 +25,6 @@ class WebSocketServer extends EventTarget {
    * @param obj
    * @param obj.connectionCallback -
    * @param obj.tlsConfig - TLSConfig containing the private key and cert chain used for TLS.
-   * @param obj.basePath - Directory path used for storing temp cert files for starting the `uWebsocket` server.
    * @param obj.host - Listen address to bind to.
    * @param obj.port - Listen port to bind to.
    * @param obj.maxIdleTimeout - Timeout time for when the connection is cleaned up after no activity.
@@ -38,7 +38,6 @@ class WebSocketServer extends EventTarget {
   static async createWebSocketServer({
     connectionCallback,
     tlsConfig,
-    basePath,
     host,
     port,
     maxIdleTimeout = 120,
@@ -48,7 +47,6 @@ class WebSocketServer extends EventTarget {
   }: {
     connectionCallback: ConnectionCallback;
     tlsConfig: TLSConfig;
-    basePath?: string;
     host?: string;
     port?: number;
     maxIdleTimeout?: number;
@@ -66,7 +64,6 @@ class WebSocketServer extends EventTarget {
     await wsServer.start({
       connectionCallback,
       tlsConfig,
-      basePath,
       host,
       port,
     });
@@ -106,7 +103,6 @@ class WebSocketServer extends EventTarget {
     connectionCallback,
   }: {
     tlsConfig: TLSConfig;
-    basePath?: string;
     host?: string;
     port?: number;
     connectionCallback?: ConnectionCallback;
@@ -212,6 +208,15 @@ class WebSocketServer extends EventTarget {
     return this._host;
   }
 
+  @startStop.ready(new webSocketErrors.ErrorWebSocketServerNotRunning())
+  public setTlsConfig(tlsConfig: TLSConfig): void {
+    const tlsServer = this.server as tls.Server;
+    tlsServer.setSecureContext({
+      key: tlsConfig.keyPrivatePem,
+      cert: tlsConfig.certChainPem,
+    });
+  }
+
   /**
    * Handles the creation of the `ReadableWritablePair` and provides it to the
    * StreamPair handler.

diff --git a/src/websockets/WebSocketStream.ts b/src/websockets/WebSocketStream.ts
@@ -6,7 +6,6 @@ import type {
 import type * as ws from 'ws';
 import type Logger from '@matrixai/logger';
 import type { NodeIdEncoded } from '../ids/types';
-import type { JSONValue } from '../types';
 import { WritableStream, ReadableStream } from 'stream/web';
 import * as webSocketErrors from './errors';
 import * as utilsErrors from '../utils/errors';
@@ -297,7 +296,7 @@ class WebSocketStream implements ReadableWritablePair<Uint8Array, Uint8Array> {
     return this._endedProm;
   }
 
-  get meta(): Record<string, JSONValue> {
+  get meta() {
     // Spreading to avoid modifying the data
     return {
       ...this.metadata,