From c5bbaee450f821400c292b85123399e442b653cb Mon Sep 17 00:00:00 2001 From: Antoine Nguyen Date: Mon, 15 Jun 2020 12:21:40 +0200 Subject: [PATCH] Simpler code and python 3.5 compat. --- modoboa/policyd/core.py | 75 +++++++------------ .../management/commands/policy_daemon.py | 10 +-- modoboa/policyd/tests.py | 8 ++ 3 files changed, 38 insertions(+), 55 deletions(-) diff --git a/modoboa/policyd/core.py b/modoboa/policyd/core.py index 0b5b7d165..fa3c2cc60 100644 --- a/modoboa/policyd/core.py +++ b/modoboa/policyd/core.py @@ -42,10 +42,10 @@ async def wait_for(dt): await asyncio.sleep(remaining) -async def run_at(dt, coro): +async def run_at(dt, coro, *args): """Run coroutine at given datetime.""" await wait_for(dt) - return await coro() + return await coro(*args) def get_local_config(): @@ -69,7 +69,7 @@ async def notify_limit_reached(ltype, name): } # We're going to execute sync code so we need an executor executor = concurrent.futures.ThreadPoolExecutor(max_workers=3) - loop = asyncio.get_running_loop() + loop = asyncio.get_event_loop() futures = [ loop.run_in_executor(executor, get_local_config), loop.run_in_executor(executor, get_notification_recipients) @@ -96,8 +96,7 @@ async def decrement_limit(rclient, ltype, name): new_counter = await rclient.hincrby(constants.REDIS_HASHNAME, name, -1) if new_counter <= 0: logger.info("Limit reached for {} {}".format(ltype, name)) - loop = asyncio.get_running_loop() - loop.create_task(notify_limit_reached(ltype, name)) + asyncio.ensure_future(notify_limit_reached(ltype, name)) async def apply_policies(attributes): @@ -129,45 +128,26 @@ async def apply_policies(attributes): return SUCCESS_ACTION -async def process_messages(queue): - """Pull from queue and process messages.""" - while True: - transport, message = await queue.get() - result = await apply_policies(message) - transport.write(b"action=" + result + b"\n\n") - queue.task_done() +async def handle_connection(reader, writer): + """Coroutine to handle a new connection to the server.""" + data = await reader.read(1024) + attributes = {} + for line in data.decode().split("\n"): + if not line: + continue + try: + name, value = line.split("=") + except ValueError: + continue + attributes[name] = value - -class APDProtocol(asyncio.Protocol): - """Access Policy Delegation protocol (postfix).""" - - def __init__(self, *args, **kwargs): - """Store queue.""" - self.queue = kwargs.pop("queue") - super().__init__(*args, **kwargs) - - def connection_made(self, transport): - peername = transport.get_extra_info("peername") - logger.info("Connection from {}".format(peername)) - self.transport = transport - - def data_received(self, data): - attributes = {} - for line in data.decode().split("\n"): - if not line: - continue - try: - name, value = line.split("=") - except ValueError: - continue - attributes[name] = value - - # action = SUCCESS_ACTION - state = attributes.get("protocol_state") - if state == "RCPT": - self.queue.put_nowait((self.transport, attributes)) - else: - self.transport.write(b"action=" + SUCCESS_ACTION + b"\n\n") + action = SUCCESS_ACTION + state = attributes.get("protocol_state") + if state == "RCPT": + action = await apply_policies(attributes) + writer.write(b"action=" + action + b"\n\n") + await writer.drain() + writer.close() def get_next_execution_dt(): @@ -195,7 +175,7 @@ async def reset_counters(): logger.info("Resetting all counters") # We're going to execute sync code so we need an executor executor = concurrent.futures.ThreadPoolExecutor(max_workers=3) - loop = asyncio.get_running_loop() + loop = asyncio.get_event_loop() futures = [ loop.run_in_executor(executor, get_domains_to_reset), loop.run_in_executor(executor, get_mailboxes_to_reset) @@ -208,12 +188,11 @@ async def reset_counters(): rclient.hset( constants.REDIS_HASHNAME, mb.full_address, mb.message_limit) # reschedule - loop = asyncio.get_running_loop() - loop.create_task(run_at(get_next_execution_dt(), reset_counters)) + asyncio.ensure_future(run_at(get_next_execution_dt(), reset_counters)) -def start_reset_counters_coro(loop): +def start_reset_counters_coro(): """Start coroutine.""" first_time = (timezone.now() + relativedelta(days=1)).replace( hour=0, minute=0, second=0) - loop.create_task(run_at(first_time, reset_counters)) + asyncio.ensure_future(run_at(first_time, reset_counters)) diff --git a/modoboa/policyd/management/commands/policy_daemon.py b/modoboa/policyd/management/commands/policy_daemon.py index b10137433..3064b069c 100644 --- a/modoboa/policyd/management/commands/policy_daemon.py +++ b/modoboa/policyd/management/commands/policy_daemon.py @@ -30,18 +30,14 @@ class Command(BaseCommand): def handle(self, *args, **options): """Entry point.""" - queue = asyncio.Queue() loop = asyncio.get_event_loop() - coro = loop.create_server( - lambda: core.APDProtocol(queue=queue), - options["host"], options["port"] + coro = asyncio.start_server( + core.handle_connection, options["host"], options["port"], loop=loop ) server = loop.run_until_complete(coro) - loop.create_task(core.process_messages(queue)) - # Schedule reset task - core.start_reset_counters_coro(loop) + core.start_reset_counters_coro() for signame in {'SIGINT', 'SIGTERM'}: loop.add_signal_handler( diff --git a/modoboa/policyd/tests.py b/modoboa/policyd/tests.py index a337474a6..39c3ebda8 100644 --- a/modoboa/policyd/tests.py +++ b/modoboa/policyd/tests.py @@ -113,9 +113,11 @@ sasl_username=user@test.com self.rclient.hget(constants.REDIS_HASHNAME, domain.name), domain.message_limit - 1 ) + s.close() # This one will be accepted and will trigger a notification # but I can't check it yet + s = self.connect_to_daemon() s.send(b"""protocol_state=RCPT sasl_username=user@test.com @@ -123,8 +125,10 @@ sasl_username=user@test.com time.sleep(0.1) res = s.recv(1024) self.assertEqual(res, b"action=dunno\n\n") + s.close() # This one will be denied + s = self.connect_to_daemon() s.send(b"""protocol_state=RCPT sasl_username=user@test.com @@ -151,9 +155,11 @@ sasl_username=user@test.com self.rclient.hget(constants.REDIS_HASHNAME, account.email), account.mailbox.message_limit - 1 ) + s.close() # This one will be accepted and will trigger a notification # but I can't check it yet + s = self.connect_to_daemon() s.send(b"""protocol_state=RCPT sasl_username=user@test.com @@ -161,8 +167,10 @@ sasl_username=user@test.com time.sleep(0.1) res = s.recv(1024) self.assertEqual(res, b"action=dunno\n\n") + s.close() # This one will be denied + s = self.connect_to_daemon() s.send(b"""protocol_state=RCPT sasl_username=user@test.com