From 671dff6d6f4ec309abe73cec7234fdc6463032a8 Mon Sep 17 00:00:00 2001
From: nicktrn <55853254+nicktrn@users.noreply.github.com>
Date: Mon, 25 Mar 2024 18:59:05 +0000
Subject: [PATCH] add secure flag to zod connection

---
 apps/coordinator/.env.example         |  3 ++-
 apps/coordinator/src/index.ts         |  2 ++
 apps/docker-provider/.env.example     |  1 +
 apps/kubernetes-provider/.env.example |  1 +
 packages/core-apps/src/provider.ts    |  3 +++
 packages/core/src/v3/zodSocket.ts     | 16 +++++++++++++---
 6 files changed, 22 insertions(+), 4 deletions(-)

diff --git a/apps/coordinator/.env.example b/apps/coordinator/.env.example
index 472939071..77377ab3c 100644
--- a/apps/coordinator/.env.example
+++ b/apps/coordinator/.env.example
@@ -1,3 +1,4 @@
 HTTP_SERVER_PORT=8020
 PLATFORM_ENABLED=true
-PLATFORM_WS_PORT=3030
\ No newline at end of file
+PLATFORM_WS_PORT=3030
+SECURE_CONNECTION=false
\ No newline at end of file
diff --git a/apps/coordinator/src/index.ts b/apps/coordinator/src/index.ts
index feb2722ba..5cd35c200 100644
--- a/apps/coordinator/src/index.ts
+++ b/apps/coordinator/src/index.ts
@@ -26,6 +26,7 @@ const PLATFORM_ENABLED = ["1", "true"].includes(process.env.PLATFORM_ENABLED ??
 const PLATFORM_HOST = process.env.PLATFORM_HOST || "127.0.0.1";
 const PLATFORM_WS_PORT = process.env.PLATFORM_WS_PORT || 3030;
 const PLATFORM_SECRET = process.env.PLATFORM_SECRET || "coordinator-secret";
+const SECURE_CONNECTION = ["1", "true"].includes(process.env.SECURE_CONNECTION ?? "true");
 
 const logger = new SimpleLogger(`[${NODE_NAME}]`);
 
@@ -365,6 +366,7 @@ class TaskCoordinator {
       namespace: "coordinator",
       host: PLATFORM_HOST,
       port: Number(PLATFORM_WS_PORT),
+      secure: SECURE_CONNECTION,
       clientMessages: CoordinatorToPlatformMessages,
       serverMessages: PlatformToCoordinatorMessages,
       authToken: PLATFORM_SECRET,
diff --git a/apps/docker-provider/.env.example b/apps/docker-provider/.env.example
index 075b6823a..2d24f79c8 100644
--- a/apps/docker-provider/.env.example
+++ b/apps/docker-provider/.env.example
@@ -2,6 +2,7 @@ HTTP_SERVER_PORT=8050
 
 PLATFORM_WS_PORT=3030
 PLATFORM_SECRET=provider-secret
+SECURE_CONNECTION=false
 
 # Use this if you are on macOS
 # COORDINATOR_HOST="host.docker.internal"
diff --git a/apps/kubernetes-provider/.env.example b/apps/kubernetes-provider/.env.example
index 10949a807..f21ee29ba 100644
--- a/apps/kubernetes-provider/.env.example
+++ b/apps/kubernetes-provider/.env.example
@@ -2,6 +2,7 @@ HTTP_SERVER_PORT=8060
 
 PLATFORM_WS_PORT=3030
 PLATFORM_SECRET=provider-secret
+SECURE_CONNECTION=false
 
 # Use this if you are on macOS
 # COORDINATOR_HOST="host.docker.internal"
diff --git a/packages/core-apps/src/provider.ts b/packages/core-apps/src/provider.ts
index d304b4fa5..1c5ae6064 100644
--- a/packages/core-apps/src/provider.ts
+++ b/packages/core-apps/src/provider.ts
@@ -18,6 +18,7 @@ const MACHINE_NAME = process.env.MACHINE_NAME || "local";
 const PLATFORM_HOST = process.env.PLATFORM_HOST || "127.0.0.1";
 const PLATFORM_WS_PORT = process.env.PLATFORM_WS_PORT || 3030;
 const PLATFORM_SECRET = process.env.PLATFORM_SECRET || "provider-secret";
+const SECURE_CONNECTION = ["1", "true"].includes(process.env.SECURE_CONNECTION ?? "true");
 
 const logger = new SimpleLogger(`[${MACHINE_NAME}]`);
 
@@ -87,6 +88,7 @@ export class ProviderShell implements Provider {
       namespace: "shared-queue",
       host: PLATFORM_HOST,
       port: Number(PLATFORM_WS_PORT),
+      secure: SECURE_CONNECTION,
       clientMessages: ClientToSharedQueueMessages,
       serverMessages: SharedQueueToClientMessages,
       authToken: PLATFORM_SECRET,
@@ -138,6 +140,7 @@ export class ProviderShell implements Provider {
       namespace: "provider",
       host: PLATFORM_HOST,
       port: Number(PLATFORM_WS_PORT),
+      secure: SECURE_CONNECTION,
       clientMessages: ProviderToPlatformMessages,
       serverMessages: PlatformToProviderMessages,
       authToken: PLATFORM_SECRET,
diff --git a/packages/core/src/v3/zodSocket.ts b/packages/core/src/v3/zodSocket.ts
index 294cdc89f..0898ce210 100644
--- a/packages/core/src/v3/zodSocket.ts
+++ b/packages/core/src/v3/zodSocket.ts
@@ -263,7 +263,8 @@ interface ZodSocketConnectionOptions<
   TServerMessages extends ZodSocketMessageCatalogSchema,
 > {
   host: string;
-  port: number;
+  port?: number;
+  secure?: boolean;
   namespace: string;
   clientMessages: TClientMessages;
   serverMessages: TServerMessages;
@@ -302,15 +303,24 @@ export class ZodSocketConnection<
   #logger: StructuredLogger;
 
   constructor(opts: ZodSocketConnectionOptions<TClientMessages, TServerMessages>) {
-    this.socket = io(`ws://${opts.host}:${opts.port}/${opts.namespace}`, {
+    const uri = `${opts.secure ? "wss" : "ws"}://${opts.host}:${
+      opts.port ?? (opts.secure ? "443" : "80")
+    }/${opts.namespace}`;
+
+    const logger = new SimpleStructuredLogger(opts.namespace, LogLevel.info);
+    logger.log("new zod socket", { uri });
+
+    this.socket = io(uri, {
       transports: ["websocket"],
       auth: {
         token: opts.authToken,
       },
       extraHeaders: opts.extraHeaders,
+      reconnectionDelay: 500,
+      reconnectionDelayMax: 1000,
     });
 
-    this.#logger = new SimpleStructuredLogger(opts.namespace, LogLevel.info, {
+    this.#logger = logger.child({
       socketId: this.socket.id,
     });