aiosmtpd weirdness with custom controller and SMTP factory

To start, I'm working with aiosmtpd, and am trying to write a class that wraps around it to start up the SMTP server programmatically with StartTLS. Now, up until very recently, this code worked as expected with any handler you might pass into it, such as a basic Message handler that I wrote to adjust parameters of the message, etc. and passing that in as part of the message headers.

import asyncio

import aiosmtpd
import aiosmtpd.controller
import aiosmtpd.handlers
import aiosmtpd.smtp

import email

import regex
import logging

import ssl


EMPTYBYTES = b''
COMMASPACE = ', '
CRLF = b'\r\n'
NLCRE = regex.compile(br'\r\n|\r|\n')


class StartTLSServer(aiosmtpd.controller.Controller):
    def __init__(self, handler, ssl_cert_file, ssl_key_file, loop=None, hostname=None,
                 port=8025, *, ready_timeout=1.0, enable_SMTPUTF8=True, decode_data=False,
                 require_starttls=True, smtp_ident=None, data_size_limit=10485760,
                 smtp_timeout=300):
        context = ssl.create_default_context(ssl.Purpose.CLIENT_AUTH)
        context.load_cert_chain(ssl_cert_file, ssl_key_file)
        self.tls_context = context
        self.require_starttls = require_starttls
        self.enable_SMTPUTF8 = enable_SMTPUTF8
        self.decode_data = decode_data
        self.smtp_ident = smtp_ident
        self.data_size_limit = data_size_limit
        self.smtp_timeout = smtp_timeout
        super().__init__(handler, loop=loop, hostname=hostname, port=port,
                         ready_timeout=ready_timeout, enable_SMTPUTF8=enable_SMTPUTF8)

    def factory(self):
        return aiosmtpd.smtp.SMTP(self.handler, data_size_limit=self.data_size_limit,
                                  enable_SMTPUTF8=self.enable_SMTPUTF8,
                                  decode_data=self.decode_data,
                                  require_starttls=self.require_starttls,
                                  hostname=self.smtp_ident,
                                  ident=self.smtp_ident,
                                  tls_context=self.tls_context,
                                  timeout=self.smtp_timeout)
                                 
                                  
class MessageHandler(aiosmtpd.handlers.Message):
    def __init__(self, message_class=None, *, loop=None):
        super().__init__(message_class)
        self.loop = loop or asyncio.get_event_loop()

    async def handle_DATA(self, server, session, envelope):
        message = self.prepare_message(session, envelope)
        await self.handle_message(message)
        return '250 OK'

    def prepare_message(self, session, envelope):
        # If the server was created with decode_data True, then data will be a
        # str, otherwise it will be bytes.
        data = envelope.content
        if isinstance(data, bytes):
            message = email.message_from_bytes(data, self.message_class)
        else:
            assert isinstance(data, str), (
              'Expected str or bytes, got {}'.format(type(data)))
            message = email.message_from_string(data, self.message_class)
        message['X-Peer'] = str(session.peer)
        message['X-Envelope-MailFrom'] = envelope.mail_from
        message['X-Envelope-RcptTo'] = COMMASPACE.join(envelope.rcpt_tos)
        return message  # This is handed off to handle_message directly.

    async def handle_message(self, message):
        print(message.as_string())
        return

This resides in custom_handlers.py which is then subsequently called in testing via the Python console as follows:

>>> from custom_handlers import StartTLSServer, MessageHandler
>>> server = StartTLSServer(MessageHandler, ssl_cert_file="valid cert path", ssl_key_file="valid key path", hostname="0.0.0.0", port=25, require_starttls=True, smtp_ident="StartTLSSMTPServer01")
>>> server.start()

When I want to stop the test server, I'll simply do a server.stop() however during processing of any message, we get hard-stopped by this evil error:

Traceback (most recent call last):
  File "/home/sysadmin/.local/lib/python3.8/site-packages/aiosmtpd/smtp.py", line 728, in _handle_client
    await method(arg)
  File "/home/sysadmin/.local/lib/python3.8/site-packages/aiosmtpd/smtp.py", line 1438, in smtp_DATA
    status = await self._call_handler_hook('DATA')
  File "/home/sysadmin/.local/lib/python3.8/site-packages/aiosmtpd/smtp.py", line 465, in _call_handler_hook
    status = await hook(self, self.session, self.envelope, *args)
TypeError: handle_DATA() missing 1 required positional argument: 'envelope'

Now, I can replicate this with ANY handler passed into the SMTP factory.

However, I can't replicate this with a plain aiosmtpd with a Debugging handler, like defined in the docs:

aiosmtpd -c aiosmtpd.handlers.Debugging stdout -l 0.0.0.0:8025

... which works fine. Passing the Debugging handler into the StartTLSServer causes the same error as the custom MessageHandler class, even with the Debugging handler.

Am I missing something obvious here about my class that's exploding here in a way that is different to the programmatic usage as expected by aiosmtpd?

Solution

You are missing () in order to instantiate an object of your MessageHandler class:

>>> server = StartTLSServer(MessageHandler(), ...)

When you simply pass MessageHandler without the (), aiosmtpd will try to invoke the regular function MessageHandler.handle_DATA(...) (as opposed to the bound method function MessageHandler().handle_DATA(...)).

This regular function takes four arguments: an instance of MessageHandler as its first argument followed by the usual server, session and envelope arguments. This explains why the error message complains about a missing positional argument.

PS, do note that your handle_DATA implementation is superfluous since it is identical to the implementation in the base class aiosmtpd.handlers.Message - so you can just delete it, and it should still work just fine.