Merge pull request #4 from CodeShellDev/main

wait for client

Author: CodeShellDev

internals/server/http.go

@@ -43,7 +43,12 @@ func httpHandler(w http.ResponseWriter, req *http.Request) {
     w.Header().Set("Content-Type", "application/json")
     json.NewEncoder(w).Encode(resp)
 
-	client := getClient(clientID)
+	client, err := WaitForClient(clientID, time.Duration(15 * time.Second))
+
+	if err != nil {
+		logger.Error("Could not get client: ", err.Error())
+		return
+	}
 
 	reachable, err := tryPing(client, body.IP, 
 		func() (bool, error) {

internals/server/websocket.go

@@ -1,6 +1,7 @@
 package server
 
 import (
+	"errors"
 	"net/http"
 	"sync"
 	"time"
@@ -61,6 +62,24 @@ func getClient(id string) *websocket.Conn {
 	return clients[id]
 }
 
+func WaitForClient(clientID string, timeout time.Duration) (*websocket.Conn, error) {
+    deadline := time.Now().Add(timeout)
+
+    for time.Now().Before(deadline) {
+        clientsMutex.Lock()
+        conn, exists := clients[clientID]
+        clientsMutex.Unlock()
+
+        if exists && conn != nil {
+            return conn, nil
+        }
+
+        time.Sleep(100 * time.Millisecond)
+    }
+
+    return nil, errors.New("Timed out waiting for client")
+}
+
 func sendToClient(client *websocket.Conn, data map[string]any) error {
 	return client.WriteJSON(data)
 }