feat(hello_handler.py): remove client from pending_hello if recv'd HELLO

src/judas_server/backend/backend_server.py

@@ -22,6 +22,7 @@ if TYPE_CHECKING:
 
 class BackendServer:
     ACK_TIMEOUT: Final[float] = 5.0  # seconds
+    HELLO_TIMEOUT: Final[float] = 3.0  # seconds
 
     def __init__(self, host: str = "0.0.0.0", port: int = 3692) -> None:
         """Initialize the backend server.
@@ -58,6 +59,7 @@ class BackendServer:
         self._initialize_handlers()
 
         self.pending_acks: list[tuple[Client, Message, float]] = []
+        self.pending_hello: dict[Client, float] = {}
 
         self.running: bool = False
 
@@ -196,6 +198,8 @@ class BackendServer:
         events = selectors.EVENT_READ | selectors.EVENT_WRITE
         self.selector.register(conn, events, data=client)
 
+        self.pending_hello[client] = time.time()
+
         self.logger.info(f"[+] Registered client {client}, HELLO pending...")
 
     def _disconnect(self, client: Client) -> None:
@@ -365,6 +369,15 @@ class BackendServer:
                         self.send(client, msg)
                         self.pending_acks.remove((client, msg, timestamp))
 
+                # check pending HELLOs
+                for client, timestamp in list(self.pending_hello.items()):
+                    if time.time() - timestamp > self.HELLO_TIMEOUT:
+                        self.logger.warning(
+                            f"HELLO timeout for {client}, disconnecting..."
+                        )
+                        self._disconnect(client)
+                        del self.pending_hello[client]
+
                 time.sleep(0.001)  # prevent 100% CPU usage
 
         except Exception as e:

src/judas_server/backend/handler/ack_handler.py

@@ -0,0 +1,27 @@
+# -*- coding: utf-8 -*-
+
+from typing import TYPE_CHECKING
+
+from .base_handler import BaseHandler
+
+if TYPE_CHECKING:
+    from judas_protocol import Message
+
+    from judas_server.backend import BackendServer, Client
+
+
+class AckHandler(BaseHandler):
+    def __init__(self, backend_server: BackendServer) -> None:
+        super().__init__(backend_server)
+
+    def handle(self, client: Client, message: Message) -> None:
+        pending_acks = self.backend_server.pending_acks
+        if message.id in pending_acks:
+            del pending_acks[message.id]
+            self.logger.debug(
+                f"[*] Received ACK for message {message.id} from {client}."
+            )
+        else:
+            self.logger.warning(
+                f"[!] Received ACK for unknown (or ACK'd) message {message.id} from {client}."
+            )

src/judas_server/backend/handler/hello_handler.py

@@ -58,6 +58,7 @@ class HelloHandler(BaseHandler):
             "last_seen": client.last_seen
         }
 
+        del self.backend_server.pending_hello[client]
         self.backend_server._save_known_clients()
         client.status = ClientStatus.ONLINE