Repository: aio-libs/aiosmtpd
Pull request #77 (2017-04-30 20:16:58)

Proxy can decode bytes or not
Most of this branch comes from #74 with additional coverage/test refactoring and some coding style fixes by me.


Patch:
diff --git a/aiosmtpd/handlers.py b/aiosmtpd/handlers.py
index d1fd871..d9bc9a6 100644
--- a/aiosmtpd/handlers.py
+++ b/aiosmtpd/handlers.py
@@ -6,6 +6,7 @@ Pass in an instance of one of these classes, or derive your own, to provide
 your own handling of messages.  Implement only the methods you care about.
 """
 
+import re
 import sys
 import asyncio
 import logging
@@ -13,13 +14,13 @@ import mailbox
 import smtplib
 
 from email import message_from_bytes, message_from_string
-from email.feedparser import NLCRE
 from public import public
 
 
-EMPTYSTRING = ''
+EMPTYBYTES = b''
 COMMASPACE = ', '
-CRLF = '\r\n'
+CRLF = b'\r\n'
+NLCRE = re.compile(br'\r\n|\r|\n')
 log = logging.getLogger('mail.debug')
 
 
@@ -93,7 +94,11 @@ class Proxy:
 
     @asyncio.coroutine
     def handle_DATA(self, server, session, envelope):
-        lines = envelope.content.splitlines(keepends=True)
+        if isinstance(envelope.content, str):
+            content = envelope.original_content
+        else:
+            content = envelope.content
+        lines = content.splitlines(keepends=True)
         # Look for the last header
         i = 0
         ending = CRLF
@@ -102,8 +107,10 @@ class Proxy:
                 ending = line
                 break
             i += 1
-        lines.insert(i, 'X-Peer: %s%s' % (session.peer[0], ending))
-        data = EMPTYSTRING.join(lines)
+        peer = session.peer[0].encode('ascii')
+        # XXX Python 3.4 does not support %-interpolation for bytes objects.
+        lines.insert(i, b'X-Peer: ' + peer + ending)
+        data = EMPTYBYTES.join(lines)
         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)


Test patch:
diff --git a/aiosmtpd/tests/test_handlers.py b/aiosmtpd/tests/test_handlers.py
index 0b0c58f..42e2634 100644
--- a/aiosmtpd/tests/test_handlers.py
+++ b/aiosmtpd/tests/test_handlers.py
@@ -23,6 +23,18 @@ class DecodingController(Controller):
         return Server(self.handler, decode_data=True)
 
 
+class DataHandler:
+    def __init__(self):
+        self.content = None
+        self.original_content = None
+
+    @asyncio.coroutine
+    def handle_DATA(self, server, session, envelope):
+        self.content = envelope.content
+        self.original_content = envelope.original_content
+        return '250 OK'
+
+
 class TestDebugging(unittest.TestCase):
     def setUp(self):
         self.stream = StringIO()
@@ -151,18 +163,6 @@ Testing
 """.format(peer))
 
 
-class DataHandler:
-    def __init__(self):
-        self.content = None
-        self.original_content = None
-
-    @asyncio.coroutine
-    def handle_DATA(self, server, session, envelope):
-        self.content = envelope.content
-        self.original_content = envelope.original_content
-        return '250 OK'
-
-
 class TestMessage(unittest.TestCase):
     def test_message(self):
         # In this test, the message content comes in as a bytes.
@@ -401,13 +401,70 @@ class TestCLI(unittest.TestCase):
 
 class TestProxy(unittest.TestCase):
     def setUp(self):
-        self.stream = StringIO()
+        # There are two controllers and two SMTPd's running here.  The
+        # "upstream" one listens on port 9025 and is connected to a "data
+        # handler" which captures the messages it receives.  The second -and
+        # the one under test here- listens on port 9024 and proxies to the one
+        # on port 9025.  Because we need to set the decode_data flag
+        # differently for each different test, the controller of the proxy is
+        # created in the individual tests, not in the setup.
+        self.upstream = DataHandler()
+        upstream_controller = Controller(self.upstream, port=9025)
+        upstream_controller.start()
+        self.addCleanup(upstream_controller.stop)
+        self.proxy = Proxy(upstream_controller.hostname, 9025)
+        self.source = """\
+From: Anne Person <anne@example.com>
+To: Bart Person <bart@example.com>
+Subject: A test
+
+Testing
+"""
+        # The upstream SMTPd will always receive the content as bytes
+        # delimited with CRLF.
+        self.expected = CRLF.join([
+            'From: Anne Person <anne@example.com>',
+            'To: Bart Person <bart@example.com>',
+            'Subject: A test',
+            'X-Peer: ::1',
+            '',
+            'Testing']).encode('ascii')
+
+    def test_deliver_bytes(self):
+        with ExitStack() as resources:
+            controller = Controller(self.proxy, port=9024)
+            controller.start()
+            resources.callback(controller.stop)
+            client = resources.enter_context(
+                SMTP(*(controller.hostname, controller.port)))
+            client.sendmail(
+                'anne@example.com', ['bart@example.com'], self.source)
+            client.quit()
+        self.assertEqual(self.upstream.content, self.expected)
+        self.assertEqual(self.upstream.original_content, self.expected)
+
+    def test_deliver_str(self):
+        with ExitStack() as resources:
+            controller = DecodingController(self.proxy, port=9024)
+            controller.start()
+            resources.callback(controller.stop)
+            client = resources.enter_context(
+                SMTP(*(controller.hostname, controller.port)))
+            client.sendmail(
+                'anne@example.com', ['bart@example.com'], self.source)
+            client.quit()
+        self.assertEqual(self.upstream.content, self.expected)
+        self.assertEqual(self.upstream.original_content, self.expected)
+
+
+class TestProxyMocked(unittest.TestCase):
+    def setUp(self):
         handler = Proxy('localhost', 9025)
         controller = DecodingController(handler)
         controller.start()
         self.addCleanup(controller.stop)
         self.address = (controller.hostname, controller.port)
-        self.message = """\
+        self.source = """\
 From: Anne Person <anne@example.com>
 To: Bart Person <bart@example.com>
 Subject: A test
@@ -415,27 +472,6 @@ Subject: A test
 Testing
 """
 
-    def test_deliver(self):
-        with ExitStack() as resources:
-            mock = resources.enter_context(
-                patch('aiosmtpd.handlers.smtplib.SMTP'))
-            client = resources.enter_context(SMTP(*self.address))
-            client.sendmail(
-                'anne@example.com', ['bart@example.com'], self.message)
-            client.quit()
-            mock().connect.assert_called_once_with('localhost', 9025)
-            # SMTP always fixes eols, so it must be always CRLF as delimiter
-            msg = CRLF.join([
-                'From: Anne Person <anne@example.com>',
-                'To: Bart Person <bart@example.com>',
-                'Subject: A test',
-                'X-Peer: ::1',
-                '',
-                'Testing'])
-            mock().sendmail.assert_called_once_with(
-                'anne@example.com', ['bart@example.com'], msg)
-            mock().quit.assert_called_once_with()
-
     def test_recipients_refused(self):
         with ExitStack() as resources:
             log_mock = resources.enter_context(patch('aiosmtpd.handlers.log'))
@@ -446,7 +482,7 @@ Testing
                 })
             client = resources.enter_context(SMTP(*self.address))
             client.sendmail(
-                'anne@example.com', ['bart@example.com'], self.message)
+                'anne@example.com', ['bart@example.com'], self.source)
             client.quit()
             # The log contains information about what happened in the proxy.
             self.assertEqual(
@@ -464,7 +500,7 @@ Testing
             mock().sendmail.side_effect = OSError
             client = resources.enter_context(SMTP(*self.address))
             client.sendmail(
-                'anne@example.com', ['bart@example.com'], self.message)
+                'anne@example.com', ['bart@example.com'], self.source)
             client.quit()
             # The log contains information about what happened in the proxy.
             self.assertEqual(