Repository: aio-libs/aiosmtpd
Pull request #45 (2017-03-21 18:30:26)

Refinements to PR #39
#39 refinements.


Patch:
diff --git a/aiosmtpd/controller.py b/aiosmtpd/controller.py
index b34f29d..4efbbf2 100644
--- a/aiosmtpd/controller.py
+++ b/aiosmtpd/controller.py
@@ -10,7 +10,6 @@ from public import public
 class Controller:
     def __init__(self, handler, loop=None, hostname='::0', port=8025):
         self.handler = handler
-        self.server = None
         self.hostname = '::0' if hostname is None else hostname
         self.port = port
         self.loop = asyncio.new_event_loop() if loop is None else loop
diff --git a/aiosmtpd/docs/NEWS.rst b/aiosmtpd/docs/NEWS.rst
index c697eeb..53e6a8d 100644
--- a/aiosmtpd/docs/NEWS.rst
+++ b/aiosmtpd/docs/NEWS.rst
@@ -4,6 +4,8 @@
 
 1.0a5 (201X-XX-XX)
 ==================
+* Deprecate handler `process_message()` methods.  Use the new asynchronous
+  `handle_DATA()` methods, which take a session and an envelope object.
 * Added the ``STARTTLS`` extension.  Given by Konstantin Volkov.
 * Minor changes to the way the ``Debugging`` handler prints ``mail_options``
   and ``rcpt_options`` (although the latter is still not support in ``SMTP``).
diff --git a/aiosmtpd/handlers.py b/aiosmtpd/handlers.py
index 2c405e7..36a54c2 100644
--- a/aiosmtpd/handlers.py
+++ b/aiosmtpd/handlers.py
@@ -64,23 +64,23 @@ class Debugging:
                 line = line.decode('utf-8', 'replace')
             print(line, file=self.stream)
 
-    def process_message(self, peer, mailfrom, rcpttos, data, **kws):
+    @asyncio.coroutine
+    def handle_DATA(self, session, envelope):
         print('---------- MESSAGE FOLLOWS ----------', file=self.stream)
         # Yes, actually test for truthiness since it's possible for either the
         # keywords to be missing, or for their values to be empty lists.
         add_separator = False
-        mail_options = kws.get('mail_options')
-        if mail_options:
-            print('mail options:', mail_options, file=self.stream)
+        if envelope.mail_options:
+            print('mail options:', envelope.mail_options, file=self.stream)
             add_separator = True
         # rcpt_options are not currently support by the SMTP class.
-        rcpt_options = kws.get('rcpt_options')
-        if rcpt_options:                            # pragma: nocover
+        rcpt_options = envelope.rcpt_options
+        if any(rcpt_options):                            # pragma: nocover
             print('rcpt options:', rcpt_options, file=self.stream)
             add_separator = True
         if add_separator:
             print(file=self.stream)
-        self._print_message_content(peer, data)
+        self._print_message_content(session.peer, envelope.data)
         print('------------ END MESSAGE ------------', file=self.stream)
 
 
@@ -90,8 +90,9 @@ class Proxy:
         self._hostname = remote_hostname
         self._port = remote_port
 
-    def process_message(self, peer, mailfrom, rcpttos, data, **kws):
-        lines = data.splitlines(keepends=True)
+    @asyncio.coroutine
+    def handle_DATA(self, session, envelope):
+        lines = envelope.data.splitlines(keepends=True)
         # Look for the last header
         i = 0
         ending = CRLF
@@ -100,19 +101,19 @@ class Proxy:
                 ending = line
                 break
             i += 1
-        lines.insert(i, 'X-Peer: %s%s' % (peer[0], ending))
+        lines.insert(i, 'X-Peer: %s%s' % (session.peer[0], ending))
         data = EMPTYSTRING.join(lines)
-        refused = self._deliver(mailfrom, rcpttos, data)
+        refused = self._deliver(envelope.mail_from, envelope.rcpt_tos, data)
         # TBD: what to do with refused addresses?
         log.info('we got some refusals: %s', refused)
 
-    def _deliver(self, mailfrom, rcpttos, data):
+    def _deliver(self, mail_from, rcpt_tos, data):
         refused = {}
         try:
             s = smtplib.SMTP()
             s.connect(self._hostname, self._port)
             try:
-                refused = s.sendmail(mailfrom, rcpttos, data)
+                refused = s.sendmail(mail_from, rcpt_tos, data)
             finally:
                 s.quit()
         except smtplib.SMTPRecipientsRefused as e:
@@ -125,7 +126,7 @@ class Proxy:
             # exception code.
             errcode = getattr(e, 'smtp_code', -1)
             errmsg = getattr(e, 'smtp_error', 'ignore')
-            for r in rcpttos:
+            for r in rcpt_tos:
                 refused[r] = (errcode, errmsg)
         return refused
 
@@ -138,7 +139,8 @@ class Sink:
             parser.error('Sink handler does not accept arguments')
         return cls()
 
-    def process_message(self, peer, mailfrom, rcpttos, data, **kws):
+    @asyncio.coroutine
+    def handle_DATA(self, session, envelope):
         pass                                        # pragma: nocover
 
 
@@ -147,25 +149,27 @@ class Message:
     def __init__(self, message_class=None):
         self.message_class = message_class
 
-    def prepare_message(self, peer, mailfrom, rcpttos, data, **kws):
+    @asyncio.coroutine
+    def handle_DATA(self, session, envelope):
+        envelope = self.prepare_message(session, envelope)
+        self.handle_message(envelope)
+
+    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.data
         if isinstance(data, bytes):
             message = message_from_bytes(data, self.message_class)
         else:
             assert isinstance(data, str), (
               'Expected str or bytes, got {}'.format(type(data)))
             message = message_from_string(data, self.message_class)
-        message['X-Peer'] = str(peer)
-        message['X-MailFrom'] = mailfrom
-        message['X-RcptTo'] = COMMASPACE.join(rcpttos)
+        message['X-Peer'] = str(session.peer)
+        message['X-MailFrom'] = envelope.mail_from
+        message['X-RcptTo'] = COMMASPACE.join(envelope.rcpt_tos)
 
         return message
 
-    def process_message(self, peer, mailfrom, rcpttos, data, **kws):
-        message = self.prepare_message(peer, mailfrom, rcpttos, data, **kws)
-        self.handle_message(message)
-
     def handle_message(self, message):
         raise NotImplementedError                   # pragma: nocover
 
@@ -177,12 +181,12 @@ class AsyncMessage(Message):
         self.loop = loop or asyncio.get_event_loop()
 
     @asyncio.coroutine
-    def process_message(self, peer, mailfrom, rcpttos, data, *, loop, **kws):
-        message = self.prepare_message(peer, mailfrom, rcpttos, data, **kws)
-        yield from self.handle_message(message, loop=loop)
+    def handle_DATA(self, session, envelope):
+        message = self.prepare_message(session, envelope)
+        yield from self.handle_message(message)
 
     @asyncio.coroutine
-    def handle_message(self, message, *, loop):
+    def handle_message(self, message):
         raise NotImplementedError                   # pragma: nocover
 
 
diff --git a/aiosmtpd/smtp.py b/aiosmtpd/smtp.py
index 146260e..a3394b3 100644
--- a/aiosmtpd/smtp.py
+++ b/aiosmtpd/smtp.py
@@ -6,6 +6,8 @@ import collections
 from email._header_value_parser import get_addr_spec, get_angle_addr
 from email.errors import HeaderParseError
 from public import public
+from warnings import warn
+
 
 try:
     import ssl
@@ -26,6 +28,24 @@ DATA_SIZE_DEFAULT = 33554432
 EMPTYBYTES = b''
 
 
+class Session:
+    def __init__(self, loop):
+        self.peer = None
+        self.ssl = None
+        self.host_name = None
+        self.extended_smtp = False
+        self.loop = loop
+
+
+class Envelope:
+    def __init__(self):
+        self.mail_from = None
+        self.mail_options = None
+        self.data = None
+        self.rcpt_tos = []
+        self.rcpt_options = []
+
+
 @public
 class SMTP(asyncio.StreamReaderProtocol):
     command_size_limit = 512
@@ -57,9 +77,6 @@ class SMTP(asyncio.StreamReaderProtocol):
                     "True at the same time")
             decode_data = False
         self._decode_data = decode_data
-        self._set_rset_state()
-        self.seen_greeting = ''
-        self.extended_smtp = False
         self.command_size_limits.clear()
         if hostname:
             self.hostname = hostname
@@ -74,8 +91,16 @@ class SMTP(asyncio.StreamReaderProtocol):
         self.require_starttls = tls_context and require_starttls
         self._tls_handshake_failed = False
         self._tls_protocol = None
+        self.session = None
+        self.envelope = None
         self.transport = None
 
+    def _create_session(self):
+        return Session(self.loop)
+
+    def _create_envelope(self):
+        return Envelope()
+
     @property
     def max_command_size_limit(self):
         try:
@@ -84,6 +109,10 @@ class SMTP(asyncio.StreamReaderProtocol):
             return self.command_size_limit
 
     def connection_made(self, transport):
+        # Reset state due to rfc3207 part 4.2.
+        self._set_rset_state()
+        self.session = self._create_session()
+        self.session.peer = transport.get_extra_info('peername')
         is_instance = (_has_ssl and
                        isinstance(transport, sslproto._SSLProtocolTransport))
         if self.transport is not None and is_instance:   # pragma: nossl
@@ -91,26 +120,21 @@ class SMTP(asyncio.StreamReaderProtocol):
             self._reader._transport = transport
             self._writer._transport = transport
             self.transport = transport
-            # Reset state due to rfc3207 part 4.2.
-            self._set_rset_state()
             self.seen_greeting = ''
             # Do SSL certificate checking as rfc3207 part 4.1 says.
             # Why _extra is protected attribute?
-            extra = self._tls_protocol._extra
+            self.session.ssl = self._tls_protocol._extra
             if hasattr(self.event_handler, 'handle_tls_handshake'):
                 auth = self.event_handler.handle_tls_handshake(
-                    extra['ssl_object'],
-                    extra['peercert'],
-                    extra['cipher'])
+                    self.session.ssl)
                 self._tls_handshake_failed = not auth
             else:
                 self._tls_handshake_failed = False
             self._over_ssl = True
         else:
             super().connection_made(transport)
-            self.peer = transport.get_extra_info('peername')
             self.transport = transport
-            log.info('Peer: %s', repr(self.peer))
+            log.info('Peer: %s', repr(self.session.peer))
             # Process the client's requests.
             self.connection_closed = False
             self._handler_coroutine = self.loop.create_task(
@@ -128,8 +152,7 @@ class SMTP(asyncio.StreamReaderProtocol):
 
     def _set_post_data_state(self):
         """Reset state variables to their post-DATA state."""
-        self.mailfrom = None
-        self.rcpttos = []
+        self.envelope = self._create_envelope()
         self.require_SMTPUTF8 = False
 
     def _set_rset_state(self):
@@ -172,7 +195,7 @@ class SMTP(asyncio.StreamReaderProtocol):
                     command = line[:i].upper()
                     arg = line[i+1:].strip()
                 max_sz = (self.command_size_limits[command]
-                          if self.extended_smtp
+                          if self.session.extended_smtp
                           else self.command_size_limit)
                 if len(line) > max_sz:
                     yield from self.push('500 Error: line too long')
@@ -208,11 +231,12 @@ class SMTP(asyncio.StreamReaderProtocol):
             yield from self.push('501 Syntax: HELO hostname')
             return
         # See issue #21783 for a discussion of this behavior.
-        if self.seen_greeting:
+        if self.session.host_name:
             yield from self.push('503 Duplicate HELO/EHLO')
             return
         self._set_rset_state()
-        self.seen_greeting = hostname
+        self.session.host_name = hostname
+        self.session.extended_smtp = False
         yield from self.push('250 %s' % self.hostname)
 
     @asyncio.coroutine
@@ -231,12 +255,12 @@ class SMTP(asyncio.StreamReaderProtocol):
             yield from self.push('501 Syntax: EHLO hostname')
             return
         # See issue #21783 for a discussion of this behavior.
-        if self.seen_greeting:
+        if self.session.host_name:
             yield from self.push('503 Duplicate HELO/EHLO')
             return
         self._set_rset_state()
-        self.seen_greeting = arg
-        self.extended_smtp = True
+        self.session.host_name = arg
+        self.session.extended_smtp = True
         yield from self.push('250-%s' % self.hostname)
         if self.data_size_limit:
             yield from self.push('250-SIZE %s' % self.data_size_limit)
@@ -341,12 +365,12 @@ class SMTP(asyncio.StreamReaderProtocol):
                 yield from self.push('250 Syntax: HELO hostname')
             elif lc_arg == 'MAIL':
                 msg = '250 Syntax: MAIL FROM: <address>'
-                if self.extended_smtp:
+                if self.session.extended_smtp:
                     msg += extended
                 yield from self.push(msg)
             elif lc_arg == 'RCPT':
                 msg = '250 Syntax: RCPT TO: <address>'
-                if self.extended_smtp:
+                if self.session.extended_smtp:
                     msg += extended
                 yield from self.push(msg)
             elif lc_arg == 'DATA':
@@ -386,12 +410,12 @@ class SMTP(asyncio.StreamReaderProtocol):
 
     @asyncio.coroutine
     def smtp_MAIL(self, arg):
-        if not self.seen_greeting:
+        if not self.session.host_name:
             yield from self.push('503 Error: send HELO first')
             return
         log.debug('===> MAIL %s', arg)
         syntaxerr = '501 Syntax: MAIL FROM: <address>'
-        if self.extended_smtp:
+        if self.session.extended_smtp:
             syntaxerr += ' [SP <mail-parameters>]'
         if arg is None:
             yield from self.push(syntaxerr)
@@ -401,14 +425,14 @@ class SMTP(asyncio.StreamReaderProtocol):
         if not address:
             yield from self.push(syntaxerr)
             return
-        if not self.extended_smtp and params:
+        if not self.session.extended_smtp and params:
             yield from self.push(syntaxerr)
             return
-        if self.mailfrom:
+        if self.envelope.mail_from:
             yield from self.push('503 Error: nested MAIL command')
             return
-        self.mail_options = params.upper().split()
-        params = self._getparams(self.mail_options)
+        mail_options = params.upper().split()
+        params = self._getparams(mail_options)
         if params is None:
             yield from self.push(syntaxerr)
             return
@@ -439,21 +463,22 @@ class SMTP(asyncio.StreamReaderProtocol):
             yield from self.push(
                 '555 MAIL FROM parameters not recognized or not implemented')
             return
-        self.mailfrom = address
-        log.info('sender: %s', self.mailfrom)
+        self.envelope.mail_from = address
+        self.envelope.mail_options = mail_options
+        log.info('sender: %s', address)
         yield from self.push('250 OK')
 
     @asyncio.coroutine
     def smtp_RCPT(self, arg):
-        if not self.seen_greeting:
+        if not self.session.host_name:
             yield from self.push('503 Error: send HELO first')
             return
         log.debug('===> RCPT %s', arg)
-        if not self.mailfrom:
+        if not self.envelope.mail_from:
             yield from self.push('503 Error: need MAIL command')
             return
         syntaxerr = '501 Syntax: RCPT TO: <address>'
-        if self.extended_smtp:
+        if self.session.extended_smtp:
             syntaxerr += ' [SP <mail-parameters>]'
         if arg is None:
             yield from self.push(syntaxerr)
@@ -463,11 +488,11 @@ class SMTP(asyncio.StreamReaderProtocol):
         if not address:
             yield from self.push(syntaxerr)
             return
-        if not self.extended_smtp and params:
+        if not self.session.extended_smtp and params:
             yield from self.push(syntaxerr)
             return
-        self.rcpt_options = params.upper().split()
-        params = self._getparams(self.rcpt_options)
+        rcpt_options = params.upper().split()
+        params = self._getparams(rcpt_options)
         if params is None:
             yield from self.push(syntaxerr)
             return
@@ -476,8 +501,9 @@ class SMTP(asyncio.StreamReaderProtocol):
             yield from self.push(
                 '555 RCPT TO parameters not recognized or not implemented')
             return
-        self.rcpttos.append(address)
-        log.info('recips: %s', self.rcpttos)
+        self.envelope.rcpt_tos.append(address)
+        self.envelope.rcpt_options.append(rcpt_options)
+        log.info('recip: %s', address)
         yield from self.push('250 OK')
 
     @asyncio.coroutine
@@ -496,10 +522,10 @@ class SMTP(asyncio.StreamReaderProtocol):
 
     @asyncio.coroutine
     def smtp_DATA(self, arg):
-        if not self.seen_greeting:
+        if not self.session.host_name:
             yield from self.push('503 Error: send HELO first')
             return
-        if not self.rcpttos:
+        if not self.envelope.rcpt_tos:
             yield from self.push('503 Error: need RCPT command')
             return
         if arg:
@@ -535,19 +561,20 @@ class SMTP(asyncio.StreamReaderProtocol):
         received_data = EMPTYBYTES.join(data)
         if self._decode_data:
             received_data = received_data.decode('utf-8')
-        args = (self.peer, self.mailfrom, self.rcpttos, received_data)
-        kwargs = {}
-        if not self._decode_data:
-            kwargs = {
-                'mail_options': self.mail_options,
-                'rcpt_options': self.rcpt_options,
-                }
-        kwargs.update({'loop': self.loop})
-        if asyncio.iscoroutinefunction(self.event_handler.process_message):
-            status = yield from self.event_handler.process_message(
-                *args, **kwargs)
+        self.envelope.data = received_data
+        # Call the new API first if it's implemented.
+        if hasattr(self.event_handler, 'handle_DATA'):
+            args = (self.session, self.envelope)
+            status = yield from self.event_handler.handle_DATA(*args)
         else:
-            status = self.event_handler.process_message(*args, **kwargs)
+            warn('Use handler.handle_DATA instead of .process_message()',
+                 DeprecationWarning)
+            args = (self.session.peer, self.envelope.mail_from,
+                    self.envelope.rcpt_tos, self.envelope.data)
+            if asyncio.iscoroutinefunction(self.event_handler.process_message):
+                status = yield from self.event_handler.process_message(*args)
+            else:
+                status = self.event_handler.process_message(*args)
         self._set_post_data_state()
         if status:
             yield from self.push(status)


Test patch:
diff --git a/aiosmtpd/tests/test_handlers.py b/aiosmtpd/tests/test_handlers.py
index ac517e0..a3096a3 100644
--- a/aiosmtpd/tests/test_handlers.py
+++ b/aiosmtpd/tests/test_handlers.py
@@ -15,6 +15,7 @@ from smtplib import SMTP, SMTPRecipientsRefused
 from tempfile import TemporaryDirectory
 from unittest.mock import call, patch
 
+
 CRLF = '\r\n'
 
 
@@ -46,6 +47,8 @@ Testing
         text = self.stream.getvalue()
         self.assertMultiLineEqual(text, """\
 ---------- MESSAGE FOLLOWS ----------
+mail options: ['SIZE=102']
+
 From: Anne Person <anne@example.com>
 To: Bart Person <bart@example.com>
 Subject: A test
@@ -77,7 +80,6 @@ Subject: A test
 Testing
 """)
         text = self.stream.getvalue()
-        # This includes mail and rcpt options because decode_data=False.
         self.assertMultiLineEqual(text, """\
 ---------- MESSAGE FOLLOWS ----------
 mail options: ['SIZE=102']
@@ -214,7 +216,7 @@ class TestAsyncMessage(unittest.TestCase):
 
         class MessageHandler(AsyncMessage):
             @asyncio.coroutine
-            def handle_message(handler_self, message, loop):
+            def handle_message(handler_self, message):
                 self.handled_message = message
 
         self.handler = MessageHandler()
@@ -458,3 +460,71 @@ Testing
                          {'bart@example.com': (-1, 'ignore')}),
                     ]
                 )
+
+
+class CapturingServer(Server):
+    def __init__(self, *args, **kws):
+        self.warnings = None
+        super().__init__(*args, **kws)
+
+    @asyncio.coroutine
+    def smtp_DATA(self, arg):
+        with patch('aiosmtpd.smtp.warn') as mock:
+            yield from super().smtp_DATA(arg)
+        self.warnings = mock.call_args_list
+
+
+class CapturingController(Controller):
+    def factory(self):
+        self.server = CapturingServer(self.handler)
+        return self.server
+
+
+class DeprecatedHandler:
+    def process_message(self, peer, mailfrom, rcpttos, data, **kws):
+        pass
+
+
+class AsyncDeprecatedHandler:
+    @asyncio.coroutine
+    def process_message(self, peer, mailfrom, rcpttos, data, **kws):
+        pass
+
+
+class TestDeprecation(unittest.TestCase):
+    # handler.process_message() is deprecated.
+    def test_deprecation(self):
+        controller = CapturingController(DeprecatedHandler())
+        controller.start()
+        self.addCleanup(controller.stop)
+        with SMTP(controller.hostname, controller.port) as client:
+            client.sendmail('anne@example.com', ['bart@example.com'], """\
+From: Anne Person <anne@example.com>
+To: Bart Person <bart@example.com>
+Subject: A test
+
+Testing
+""")
+        self.assertEqual(len(controller.server.warnings), 1)
+        self.assertEqual(
+            controller.server.warnings[0],
+            call('Use handler.handle_DATA instead of .process_message()',
+                 DeprecationWarning))
+
+    def test_deprecation_async(self):
+        controller = CapturingController(AsyncDeprecatedHandler())
+        controller.start()
+        self.addCleanup(controller.stop)
+        with SMTP(controller.hostname, controller.port) as client:
+            client.sendmail('anne@example.com', ['bart@example.com'], """\
+From: Anne Person <anne@example.com>
+To: Bart Person <bart@example.com>
+Subject: A test
+
+Testing
+""")
+        self.assertEqual(len(controller.server.warnings), 1)
+        self.assertEqual(
+            controller.server.warnings[0],
+            call('Use handler.handle_DATA instead of .process_message()',
+                 DeprecationWarning))
diff --git a/aiosmtpd/tests/test_smtp.py b/aiosmtpd/tests/test_smtp.py
index b4e58a0..840f23e 100644
--- a/aiosmtpd/tests/test_smtp.py
+++ b/aiosmtpd/tests/test_smtp.py
@@ -31,8 +31,9 @@ class ReceivingHandler:
     def __init__(self):
         self.box = []
 
-    def process_message(self, *args, **kws):
-        self.box.append(args)
+    @asyncio.coroutine
+    def handle_DATA(self, session, envelope):
+        self.box.append(envelope.data)
 
 
 class SizedController(Controller):
@@ -64,7 +65,8 @@ class CustomIdentController(Controller):
 class ErroringHandler:
     error = None
 
-    def process_message(self, peer, mailfrom, rcpttos, data, **kws):
+    @asyncio.coroutine
+    def handle_DATA(self, session, envelope):
         return '499 Could not accept the message'
 
     @asyncio.coroutine
@@ -105,7 +107,7 @@ class TestProtocol(unittest.TestCase):
         except asyncio.CancelledError:
             pass
         assert len(handler.box) == 1
-        assert handler.box[0][3] == data
+        assert handler.box[0] == data
 
 
 class TestSMTP(unittest.TestCase):
@@ -671,8 +673,7 @@ Testing
             mail = CRLF.join(['Test', '.', 'mail'])
             client.sendmail('anne@example.com', ['bart@example.com'], mail)
             self.assertEqual(len(handler.box), 1)
-            mail = handler.box[0]
-            self.assertEqual(mail[3], 'Test\r\n.\r\nmail')
+            self.assertEqual(handler.box[0], 'Test\r\n.\r\nmail')
 
     def test_unexpected_errors(self):
         class ErrorSMTP(Server):