X-Git-Url: http://average.org/gitweb/?a=blobdiff_plain;f=loctrkd%2Fcollector.py;h=788cb11ca1ee8b186838ea1ef882a5b6faf4ea83;hb=f55e1a18723a3b3aca34e9a4b8b878860f951e52;hp=26d37f33f60af39a123c8aafa043c50244125ef0;hpb=e2e977603507f9121e5c60c5b69d98fb425fa833;p=loctrkd.git diff --git a/loctrkd/collector.py b/loctrkd/collector.py index 26d37f3..788cb11 100644 --- a/loctrkd/collector.py +++ b/loctrkd/collector.py @@ -14,10 +14,11 @@ from socket import ( ) from struct import pack from time import time -from typing import Any, cast, Dict, List, Optional, Set, Tuple, Union +from typing import Any, Dict, List, Optional, Set, Tuple, Union import zmq from . import common +from .protomodule import ProtoModule from .zmsg import Bcast, Resp log = getLogger("loctrkd/collector") @@ -25,50 +26,10 @@ log = getLogger("loctrkd/collector") MAXBUFFER: int = 4096 -class ProtoModule: - class Stream: - 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: - ... - - @staticmethod - def parse_message(packet: bytes, is_incoming: bool = True) -> Any: - ... - - @staticmethod - def inline_response(packet: bytes) -> Optional[bytes]: - ... - - @staticmethod - def is_goodbye_packet(packet: bytes) -> bool: - ... - - @staticmethod - def imei_from_packet(packet: bytes) -> Optional[str]: - ... - - @staticmethod - def proto_of_message(packet: bytes) -> str: - ... - - -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 @@ -87,7 +48,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) @@ -107,11 +68,9 @@ class Client: ) return None if self.stream is None: - for pmod in pmods: - if pmod.probe_buffer(segment): - self.pmod = pmod - self.stream = pmod.Stream() - break + self.pmod = common.probe_pmod(segment) + if self.pmod is not None: + self.stream = self.pmod.Stream() if self.stream is None: log.info( "unrecognizable %d bytes of data %s from fd %d", @@ -151,15 +110,20 @@ class Clients: def __init__(self) -> None: self.by_fd: Dict[int, Client] = {} self.by_imei: Dict[str, Client] = {} - self.tostop: Set[int] = set() - 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() @@ -169,9 +133,10 @@ class Clients: 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,7 +154,7 @@ class Clients: oldfd = oldclnt.sock.fileno() log.info("Removing stale connection on fd %d", oldfd) oldclnt.imei = None - self.tostop.add(oldfd) + self.stop(oldfd) self.by_imei[clnt.imei] = clnt result.append((clnt.pmod, clnt.imei, when, peeraddr, packet)) log.debug( @@ -211,11 +176,6 @@ class Clients: def runserver(conf: ConfigParser, handle_hibernate: bool = True) -> None: - global pmods - pmods = [ - cast(ProtoModule, import_module("." + modnm, __package__)) - for modnm in conf.get("collector", "protocols").split(",") - ] # Is this https://github.com/zeromq/pyzmq/issues/1627 still not fixed?! zctx = zmq.Context() # type: ignore zpub = zctx.socket(zmq.PUB) # type: ignore @@ -233,11 +193,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 = [] - clients.tostop = set() + tosend: List[Resp] = [] + toadd: List[Tuple[socket, Any]] = [] events = poller.poll(1000) for sk, fl in events: if sk is zpull: @@ -251,12 +211,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) - clients.tostop.add(sk) + clients.stop(sk) else: for pmod, imei, when, peeraddr, packet in received: proto = pmod.proto_of_message(packet) @@ -278,7 +238,7 @@ def runserver(conf: ConfigParser, handle_hibernate: bool = True) -> None: sk, imei, ) - clients.tostop.add(sk) + clients.stop(sk) respmsg = pmod.inline_response(packet) if respmsg is not None: tosend.append( @@ -300,12 +260,13 @@ def runserver(conf: ConfigParser, handle_hibernate: bool = True) -> None: packet=zmsg.packet, ).packed ) - for fd in clients.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()