policy daemon: close all connections and deal with timeouts.

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():

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...")