Websocket Termination Fixed

services/client/client.go

@@ -1,6 +1,7 @@
 package client
 
 import (
+    "time"
 	"context"
 	"github.com/gorilla/websocket"
 	"github.com/tiagorlampert/CHAOS/internal/utils"
@@ -43,3 +44,53 @@ type Service interface {
 	SendCommand(ctx context.Context, input SendCommandInput) (SendCommandOutput, error)
 	BuildClient(BuildClientBinaryInput) (string, error)
 }
+
+// function to handle reconnection
+func (c *Client) reconnect() error {
+    backoff := time.Second
+    for {
+        c.conn, _, err := websocket.DefaultDialer.Dial(c.url, nil)
+        if err == nil {
+            return nil
+        }
+        if backoff > time.Minute {
+            return err
+        }
+        time.Sleep(backoff)
+        backoff *= 2
+    }
+}
+
+// Modify the main loop to handle unexpected termination
+func (c *Client) run() {
+    for {
+        err := c.readPump()
+        if err != nil {
+            if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
+                // Log the error here
+                log.Printf("Unexpected websocket closure: %v", err)
+            }
+            // Attempt to reconnect
+            if err := c.reconnect(); err != nil {
+                log.Printf("Failed to reconnect: %v", err)
+                return
+            }
+            // Reset any necessary client state here
+        }
+    }
+}
+
+// Modify the health check and device request functions
+func (c *Client) sendHealthCheck() error {
+    if c.conn == nil {
+        return errors.New("Connection is not established")
+    }
+    return c.conn.WriteMessage(websocket.TextMessage, []byte("GET /health"))
+}
+
+func (c *Client) sendDeviceRequest() error {
+    if c.conn == nil {
+        return errors.New("Connection is not established")
+    }
+    return c.conn.WriteMessage(websocket.TextMessage, []byte("POST /device"))
+}