diff --git a/src/judas_server/backend/backend_server.py b/src/judas_server/backend/backend_server.py index 0411a15..df28f55 100644 --- a/src/judas_server/backend/backend_server.py +++ b/src/judas_server/backend/backend_server.py @@ -7,7 +7,7 @@ import socket import threading import time -from judas_protocol import Message +from judas_protocol import Category, ControlAction, Message from judas_server.backend.client import Client @@ -78,43 +78,13 @@ class BackendServer: """ conn, addr = sock.accept() self.logger.info(f"[+] Accepted connection from {addr}") - conn.setblocking(False) - # wait for hello message to get mac_id - - conn.settimeout(5) - try: - message = conn.recv(1024) - if not message: - self.logger.error(f"[-] No data received from {addr}") - conn.close() - return - except socket.timeout: - self.logger.error(f"[-] Timeout waiting for hello from {addr}") - conn.close() - return - conn.settimeout(None) - - message = message.split(b"\n")[0] # get first line only - message = Message.from_bytes(message) - - mac_id = message.payload.get("mac", None) - if mac_id is None: - self.logger.error( - f"[-] No mac_id provided by {addr}, closing connection" - ) - conn.close() - return - - client = Client(id_=mac_id, addr=addr, socket=conn) - self.clients[mac_id] = client - - self._send_ack(client) + client = Client(mac_id=None, addr=addr, socket=conn) events = selectors.EVENT_READ | selectors.EVENT_WRITE self.selector.register(conn, events, data=client) - self.logger.info(f"[+] Registered client {client}") + self.logger.info(f"[+] Registered client {client}, HELLO pending...") def _disconnect(self, client: Client) -> None: """Disconnect a client and clean up resources. @@ -122,13 +92,52 @@ class BackendServer: Args: sock (socket.socket): The client socket to disconnect. """ - self.logger.info(f"[-] Disconnecting {client}") + self.logger.info(f"[-] Disconnecting {client}...") try: self.selector.unregister(client.socket) except Exception as e: self.logger.error(f"Error unregistering client {client}: {e}") + client.disconnect() + def _send_outbound( + self, sock: socket.socket, client: Client, data: bytes + ) -> None: + """Queue data to be sent to a client. + + Args: + client (Client): The client to send data to. + data (bytes): The data to send. + """ + self.logger.debug(f"[>] Sending data to {client}: {client.outbound!r}") + sent = sock.send(client.outbound) + + client.outbound = client.outbound[sent:] + + def _receive_inbound( + self, sock: socket.socket, client: Client, packet_size: int = 4096 + ) -> None: + """Receive data from a client socket. + + Args: + sock (socket.socket): The client socket to receive data from. + client (Client): The client object. + packet_size (int): The maximum amount of data to be received at once. + Returns: + bytes: The received data. + """ + recv_data = sock.recv(1024) + if recv_data: + self.logger.debug( + f"[<] Received data from {client}: {recv_data!r}" + ) + client.inbound += recv_data + + # set last seen + client.last_seen = time.time() + else: + self._disconnect(client) + def _handle_connection( self, key: selectors.SelectorKey, mask: int ) -> None: @@ -143,12 +152,47 @@ class BackendServer: try: if mask & selectors.EVENT_READ: - recv_data = sock.recv(1024) - if recv_data: - self.logger.debug( - f"[<] Received data from {client}: {recv_data!r}" - ) - client.inbound += recv_data + self._receive_inbound(sock, client) + if client.inbound: + if client.mac_id is None: + # expect HELLO message + try: + msg = Message.from_bytes(client.inbound) + if ( + msg.category == Category.CONTROL + and msg.action == ControlAction.HELLO + and msg.payload.get("mac") is not None + ): + client.mac_id = msg.payload["mac"] + if ( + client.mac_id in self.clients + and self.clients[client.mac_id].status + == "connected" + ): + old_client: Client = self.clients[ + client.mac_id + ] + self.logger.warning( + f"Client {client.mac_id} is already connected from {old_client.addr}, disconnecting old client..." + ) + self._disconnect(old_client) + # TODO: tell client not to reconnect + self.clients[client.mac_id] = client + self.logger.info( + f"[+] Registered new client {client}" + ) + else: + self.logger.error( + f"Expected HELLO message from {client}, got {msg}" + ) + self._disconnect(client) + return + except Exception as e: + self.logger.error( + f"Failed to parse HELLO message from {client}: {e}" + ) + self._disconnect(client) + return while b"\n" in client.inbound: line, client.inbound = client.inbound.split(b"\n", 1) self.logger.info( @@ -156,21 +200,12 @@ class BackendServer: ) self._send_ack(client) - - # set last seen - client.last_seen = time.time() else: self._disconnect(client) if mask & selectors.EVENT_WRITE: if client.outbound: - self.logger.debug( - f"[>] Sending data to {client}: {client.outbound!r}" - ) - sent = sock.send(client.outbound) - - client.outbound = client.outbound[sent:] - # TODO: wait for ACK from client + self._send_outbound(sock, client, client.outbound) except ConnectionResetError as e: self.logger.error(f"Connection reset by {client}, disconnect: {e}") @@ -213,7 +248,7 @@ class BackendServer: self.logger.warning(f"Client {client_id} not found") return None return { - "id": client.id, + "id": client.mac_id, "addr": client.addr, "last_seen": client.last_seen, "status": client.status, diff --git a/src/judas_server/backend/client.py b/src/judas_server/backend/client.py index dd35fde..b680250 100644 --- a/src/judas_server/backend/client.py +++ b/src/judas_server/backend/client.py @@ -17,12 +17,13 @@ class Client: """Represents a client.""" def __init__( - self, id_: str, addr: tuple[str, int], socket: socket.socket + self, mac_id: str | None, addr: tuple[str, int], socket: socket.socket ) -> None: """Initialize the client. Args: - id_ (str): The unique identifier for the client. + mac_id (str | None): The unique identifier for the client. + Can be None if not yet assigned. addr (tuple[str, int]): The (IP, port) address of the client. socket (socket.socket): The socket object for communication. """ @@ -31,7 +32,7 @@ class Client: ) self.logger.debug(f"Initializing Client {addr}...") - self.id: str = id_ + self.mac_id: str | None = mac_id self.last_seen: float = 0.0 # unix timestanp of last inbound message self.status: ClientStatus = ClientStatus.CONNECTED @@ -41,10 +42,10 @@ class Client: self.outbound: bytes = b"" def __str__(self) -> str: - return f"Client({self.id} ({self.addr[0]}:{self.addr[1]}))" + return f"Client({self.mac_id} ({self.addr[0]}:{self.addr[1]}))" def __repr__(self) -> str: - return f"Client({self.id}, {self.addr})" + return f"Client({self.mac_id}, {self.addr})" def disconnect(self) -> None: """Disconnect the client and close the socket."""