diff --git a/modoboa/policyd/core.py b/modoboa/policyd/core.py index 63ff87a81..625346ebe 100644 --- a/modoboa/policyd/core.py +++ b/modoboa/policyd/core.py @@ -125,6 +125,9 @@ async def apply_policies(attributes): await decrement_limit(rclient, "domain", domain) if decr_user: await decrement_limit(rclient, "account", sasl_username) + rclient.close() + await rclient.wait_closed() + logger.debug("Let it pass") return SUCCESS_ACTION @@ -132,6 +135,7 @@ async def handle_connection(reader, writer): """Coroutine to handle a new connection to the server.""" action = SUCCESS_ACTION try: + logger.debug("Reading data") data = await reader.readuntil(b"\n\n") except asyncio.IncompleteReadError: pass @@ -147,10 +151,25 @@ async def handle_connection(reader, writer): attributes[name] = value state = attributes.get("protocol_state") if state == "RCPT": + logger.debug("Applying policies") action = await apply_policies(attributes) + logger.debug("Done") + logger.debug("Sending action %s", action) writer.write(b"action=" + action + b"\n\n") await writer.drain() - writer.close() + + +async def new_connection(reader, writer): + try: + await asyncio.wait_for(handle_connection(reader, writer), timeout=5) + except asyncio.TimeoutError as err: + logger.warning("Timeout received while handling connection: %s", err) + finally: + writer.close() + if hasattr(writer, "wait_closed"): + # Python 3.7+ only + await writer.wait_closed() + logger.info("exit") def get_next_execution_dt(): diff --git a/modoboa/policyd/management/commands/policy_daemon.py b/modoboa/policyd/management/commands/policy_daemon.py index 3064b069c..a454e032d 100644 --- a/modoboa/policyd/management/commands/policy_daemon.py +++ b/modoboa/policyd/management/commands/policy_daemon.py @@ -27,12 +27,14 @@ class Command(BaseCommand): """Add command line arguments.""" parser.add_argument("--host", type=str, default="127.0.0.1") parser.add_argument("--port", type=int, default=9999) + parser.add_argument("--debug", action="store_true", + help="Enable debug mode") def handle(self, *args, **options): """Entry point.""" loop = asyncio.get_event_loop() coro = asyncio.start_server( - core.handle_connection, options["host"], options["port"], loop=loop + core.new_connection, options["host"], options["port"], loop=loop ) server = loop.run_until_complete(coro) @@ -47,6 +49,10 @@ class Command(BaseCommand): logger.info("Serving on {}".format(server.sockets[0].getsockname())) + if options["debug"]: + loop.set_debug(True) + logging.getLogger('asyncio').setLevel(logging.DEBUG) + loop.run_forever() logger.info("Stopping policy daemon...")