From 25cef26251f3c3fa129960c41b3e51b9bd0190f2 Mon Sep 17 00:00:00 2001
From: Owen <owen@txv.io>
Date: Tue, 19 Aug 2025 21:29:56 -0700
Subject: [PATCH] Fix ws reconnect and change create site

---
 server/hybridServer.ts                        |  6 +++-
 server/routers/site/createSite.ts             |  6 ++--
 server/routers/ws/client.ts                   | 35 +++++++++++++++----
 .../settings/sites/[niceId]/SiteInfoCard.tsx  |  2 +-
 4 files changed, 39 insertions(+), 10 deletions(-)

diff --git a/server/hybridServer.ts b/server/hybridServer.ts
index e38ca088..2cd04e0d 100644
--- a/server/hybridServer.ts
+++ b/server/hybridServer.ts
@@ -139,9 +139,13 @@ export async function createHybridClientServer() {
         logger.error("Failed to connect:", error);
     }
 
-    client.sendMessageInterval(
+    // Store the ping interval stop function for cleanup if needed
+    const stopPingInterval = client.sendMessageInterval(
         "remoteExitNode/ping",
         { timestamp: Date.now() / 1000 },
         60000
     ); // send every minute
+
+    // Return client and cleanup function for potential use
+    return { client, stopPingInterval };
 }
diff --git a/server/routers/site/createSite.ts b/server/routers/site/createSite.ts
index 66af0b1f..3a4dd885 100644
--- a/server/routers/site/createSite.ts
+++ b/server/routers/site/createSite.ts
@@ -207,7 +207,7 @@ export async function createSite(
         await db.transaction(async (trx) => {
             let newSite: Site;
 
-            if (exitNodeId) {
+            if ((type == "wireguard" || type == "newt") && exitNodeId) {
                 // we are creating a site with an exit node (tunneled)
                 if (!subnet) {
                     return next(
@@ -264,12 +264,14 @@ export async function createSite(
                 [newSite] = await trx
                     .insert(sites)
                     .values({
+                        exitNodeId: exitNodeId,
                         orgId,
                         name,
                         niceId,
                         address: updatedAddress || null,
                         type,
-                        dockerSocketEnabled: type == "newt",
+                        dockerSocketEnabled: false,
+                        online: true,
                         subnet: "0.0.0.0/0"
                     })
                     .returning();
diff --git a/server/routers/ws/client.ts b/server/routers/ws/client.ts
index fda1e62c..13b5d0da 100644
--- a/server/routers/ws/client.ts
+++ b/server/routers/ws/client.ts
@@ -145,11 +145,11 @@ export class WebSocketClient extends EventEmitter {
   }
 
   private async connectWithRetry(): Promise<void> {
-    if (this.isConnecting) return;
+    if (this.isConnecting || this.isConnected) return;
     
     this.isConnecting = true;
     
-    while (this.shouldReconnect && !this.isConnected) {
+    while (this.shouldReconnect && !this.isConnected && this.isConnecting) {
       try {
         await this.establishConnection();
         this.isConnecting = false;
@@ -157,7 +157,7 @@ export class WebSocketClient extends EventEmitter {
       } catch (error) {
         logger.error(`Failed to connect: ${error}. Retrying in ${this.reconnectInterval}ms...`);
         
-        if (!this.shouldReconnect) {
+        if (!this.shouldReconnect || !this.isConnecting) {
           this.isConnecting = false;
           return;
         }
@@ -172,6 +172,13 @@ export class WebSocketClient extends EventEmitter {
   }
 
   private async establishConnection(): Promise<void> {
+    // Clean up any existing connection before establishing a new one
+    if (this.conn) {
+      this.conn.removeAllListeners();
+      this.conn.close();
+      this.conn = null;
+    }
+
     // Parse the base URL to determine protocol and hostname
     const baseURL = new URL(this.baseURL);
     const wsProtocol = baseURL.protocol === 'https:' ? 'wss' : 'ws';
@@ -217,9 +224,8 @@ export class WebSocketClient extends EventEmitter {
         if (this.conn === null) {
           // Connection failed during establishment
           reject(error);
-        } else {
-          this.handleDisconnect();
         }
+        // Don't call handleDisconnect here as the 'close' event will handle it
       });
 
       conn.on('pong', () => {
@@ -232,6 +238,12 @@ export class WebSocketClient extends EventEmitter {
   }
 
   private startPingMonitor(): void {
+    // Clear any existing ping timer to prevent duplicates
+    if (this.pingTimer) {
+      clearInterval(this.pingTimer);
+      this.pingTimer = null;
+    }
+    
     this.pingTimer = setInterval(() => {
       if (this.conn && this.conn.readyState === WebSocket.OPEN) {
         this.conn.ping();
@@ -246,6 +258,11 @@ export class WebSocketClient extends EventEmitter {
   }
 
   private handleDisconnect(): void {
+    // Prevent multiple disconnect handlers from running simultaneously
+    if (!this.isConnected && !this.isConnecting) {
+      return;
+    }
+
     this.setConnected(false);
     this.isConnecting = false;
     
@@ -259,6 +276,12 @@ export class WebSocketClient extends EventEmitter {
       this.pingTimeoutTimer = null;
     }
 
+    // Clear any existing reconnect timer to prevent multiple reconnection attempts
+    if (this.reconnectTimer) {
+      clearTimeout(this.reconnectTimer);
+      this.reconnectTimer = null;
+    }
+
     if (this.conn) {
       this.conn.removeAllListeners();
       this.conn = null;
@@ -269,7 +292,7 @@ export class WebSocketClient extends EventEmitter {
     // Reconnect if needed
     if (this.shouldReconnect) {
       // Add a small delay before starting reconnection to prevent immediate retry
-      setTimeout(() => {
+      this.reconnectTimer = setTimeout(() => {
         this.connectWithRetry();
       }, 1000);
     }
diff --git a/src/app/[orgId]/settings/sites/[niceId]/SiteInfoCard.tsx b/src/app/[orgId]/settings/sites/[niceId]/SiteInfoCard.tsx
index 36ab1727..5eed91c5 100644
--- a/src/app/[orgId]/settings/sites/[niceId]/SiteInfoCard.tsx
+++ b/src/app/[orgId]/settings/sites/[niceId]/SiteInfoCard.tsx
@@ -66,7 +66,7 @@ export default function SiteInfoCard({}: SiteInfoCardProps) {
                         </InfoSectionContent>
                     </InfoSection>
 
-                    {env.flags.enableClients && (
+                    {env.flags.enableClients && site.type == "newt" && (
                         <InfoSection>
                             <InfoSectionTitle>Address</InfoSectionTitle>
                             <InfoSectionContent>