X-Git-Url: http://www.average.org/gitweb/?a=blobdiff_plain;f=loctrkd%2Fcollector.py;h=6fc606a14ecabcbe41c6aba94c69ca7ce5bc55f4;hb=1f1c34f9f42911958be5b282bd9092c1ed3de2a5;hp=c6261d705bfb9886fadf42dbb659dd12860754bf;hpb=dbdf9d63af31770ad57302e16b17a2fdc526773f;p=loctrkd.git diff --git a/loctrkd/collector.py b/loctrkd/collector.py index c6261d7..6fc606a 100644 --- a/loctrkd/collector.py +++ b/loctrkd/collector.py @@ -14,7 +14,7 @@ from socket import ( ) from struct import pack from time import time -from typing import Any, cast, Dict, List, Optional, Tuple, Union +from typing import Any, cast, Dict, List, Optional, Set, Tuple, Union import zmq from . import common @@ -27,16 +27,16 @@ MAXBUFFER: int = 4096 class ProtoModule: class Stream: - @staticmethod - def enframe(buffer: bytes) -> bytes: - ... - def recv(self, segment: bytes) -> List[Union[bytes, str]]: ... def close(self) -> bytes: ... + @staticmethod + def enframe(buffer: bytes, imei: Optional[str] = None) -> bytes: + ... + @staticmethod def probe_buffer(buffer: bytes) -> bool: ... @@ -61,10 +61,6 @@ class ProtoModule: def proto_of_message(packet: bytes) -> str: ... - @staticmethod - def proto_by_name(name: str) -> int: - ... - pmods: List[ProtoModule] = [] @@ -72,7 +68,7 @@ pmods: List[ProtoModule] = [] class Client: """Connected socket to the terminal plus buffer and metadata""" - def __init__(self, sock: socket, addr: Tuple[str, int]) -> None: + def __init__(self, sock: socket, addr: Any) -> None: self.sock = sock self.addr = addr self.pmod: Optional[ProtoModule] = None @@ -91,7 +87,7 @@ class Client: "%d bytes in buffer on close: %s", len(rest), rest[:64].hex() ) - def recv(self) -> Optional[List[Tuple[float, Tuple[str, int], bytes]]]: + def recv(self) -> Optional[List[Tuple[float, Any, bytes]]]: """Read from the socket and parse complete messages""" try: segment = self.sock.recv(MAXBUFFER) @@ -139,9 +135,9 @@ class Client: return msgs def send(self, buffer: bytes) -> None: - assert self.stream is not None + assert self.stream is not None and self.pmod is not None try: - self.sock.send(self.stream.enframe(buffer)) + self.sock.send(self.pmod.enframe(buffer, imei=self.imei)) except OSError as e: log.error( "Sending to fd %d (IMEI %s): %s", @@ -156,25 +152,32 @@ class Clients: self.by_fd: Dict[int, Client] = {} self.by_imei: Dict[str, Client] = {} - def add(self, clntsock: socket, clntaddr: Tuple[str, int]) -> int: + def fds(self) -> Set[int]: + return set(self.by_fd.keys()) + + def add(self, clntsock: socket, clntaddr: Any) -> int: fd = clntsock.fileno() log.info("Start serving fd %d from %s", fd, clntaddr) self.by_fd[fd] = Client(clntsock, clntaddr) return fd def stop(self, fd: int) -> None: + if fd not in self.by_fd: + log.debug("Fd %d is not served, ingore stop", fd) + return clnt = self.by_fd[fd] log.info("Stop serving fd %d (IMEI %s)", clnt.sock.fileno(), clnt.imei) clnt.close() - if clnt.imei: + if clnt.imei and self.by_imei[clnt.imei] == clnt: # could be replaced del self.by_imei[clnt.imei] del self.by_fd[fd] def recv( self, fd: int - ) -> Optional[ - List[Tuple[ProtoModule, Optional[str], float, Tuple[str, int], bytes]] - ]: + ) -> Optional[List[Tuple[ProtoModule, Optional[str], float, Any, bytes]]]: + if fd not in self.by_fd: + log.debug("Client at fd %d gone, ingore event", fd) + return None clnt = self.by_fd[fd] msgs = clnt.recv() if msgs is None: @@ -189,18 +192,11 @@ class Clients: clnt.imei = imei oldclnt = self.by_imei.get(clnt.imei) if oldclnt is not None: - log.info( - "Orphaning fd %d with the same IMEI", - oldclnt.sock.fileno(), - ) + oldfd = oldclnt.sock.fileno() + log.info("Removing stale connection on fd %d", oldfd) oldclnt.imei = None + self.stop(oldfd) self.by_imei[clnt.imei] = clnt - else: - log.warning( - "Login message from %s: %s, but client imei unfilled", - peeraddr, - packet, - ) result.append((clnt.pmod, clnt.imei, when, peeraddr, packet)) log.debug( "Received from %s (IMEI %s): %s", @@ -243,11 +239,11 @@ def runserver(conf: ConfigParser, handle_hibernate: bool = True) -> None: poller.register(zpull, flags=zmq.POLLIN) poller.register(tcpfd, flags=zmq.POLLIN) clients = Clients() + pollingfds: Set[int] = set() try: while True: - tosend = [] - topoll = [] - tostop = [] + tosend: List[Resp] = [] + toadd: List[Tuple[socket, Any]] = [] events = poller.poll(1000) for sk, fl in events: if sk is zpull: @@ -261,12 +257,12 @@ def runserver(conf: ConfigParser, handle_hibernate: bool = True) -> None: elif sk == tcpfd: clntsock, clntaddr = tcpl.accept() clntsock.setsockopt(SOL_SOCKET, SO_KEEPALIVE, 1) - topoll.append((clntsock, clntaddr)) + toadd.append((clntsock, clntaddr)) elif fl & zmq.POLLIN: received = clients.recv(sk) if received is None: log.debug("Terminal gone from fd %d", sk) - tostop.append(sk) + clients.stop(sk) else: for pmod, imei, when, peeraddr, packet in received: proto = pmod.proto_of_message(packet) @@ -288,7 +284,7 @@ def runserver(conf: ConfigParser, handle_hibernate: bool = True) -> None: sk, imei, ) - tostop.append(sk) + clients.stop(sk) respmsg = pmod.inline_response(packet) if respmsg is not None: tosend.append( @@ -310,12 +306,13 @@ def runserver(conf: ConfigParser, handle_hibernate: bool = True) -> None: packet=zmsg.packet, ).packed ) - for fd in tostop: + for fd in pollingfds - clients.fds(): poller.unregister(fd) # type: ignore - clients.stop(fd) - for clntsock, clntaddr in topoll: + for clntsock, clntaddr in toadd: fd = clients.add(clntsock, clntaddr) + for fd in clients.fds() - pollingfds: poller.register(fd, flags=zmq.POLLIN) + pollingfds = clients.fds() except KeyboardInterrupt: zpub.close() zpull.close()