diff options
Diffstat (limited to 'src')
| -rw-r--r-- | src/mailman/mta/tests/test_connection.py | 47 | ||||
| -rw-r--r-- | src/mailman/testing/mta.py | 45 |
2 files changed, 83 insertions, 9 deletions
diff --git a/src/mailman/mta/tests/test_connection.py b/src/mailman/mta/tests/test_connection.py index 3c4114634..7e4d556e1 100644 --- a/src/mailman/mta/tests/test_connection.py +++ b/src/mailman/mta/tests/test_connection.py @@ -22,7 +22,7 @@ import unittest from mailman.config import config from mailman.mta.connection import Connection from mailman.testing.layers import SMTPLayer -from smtplib import SMTPAuthenticationError +from smtplib import SMTP, SMTPAuthenticationError class TestConnection(unittest.TestCase): @@ -57,3 +57,48 @@ Subject: aardvarks """) self.assertEqual(self.layer.smtpd.get_authentication_credentials(), 'AHRlc3R1c2VyAHRlc3RwYXNz') + + +class TestConnectionCount(unittest.TestCase): + layer = SMTPLayer + + def setUp(self): + self.connection = Connection( + config.mta.smtp_host, int(config.mta.smtp_port), 0) + self.msg_text = """\ +From: anne@example.com +To: bart@example.com +Subject: aardvarks + +""" + + def test_count_0(self): + # So far, no connections. + self.assertEqual(SMTPLayer.smtpd.get_connection_count(), 0) + + def test_count_1(self): + self.connection.sendmail( + 'anne@example.com', ['bart@example.com'], self.msg_text) + self.assertEqual(SMTPLayer.smtpd.get_connection_count(), 1) + + def test_count_2(self): + self.connection.sendmail( + 'anne@example.com', ['bart@example.com'], self.msg_text) + self.connection.quit() + self.connection.sendmail( + 'cate@example.com', ['dave@example.com'], self.msg_text) + self.connection.quit() + self.assertEqual(SMTPLayer.smtpd.get_connection_count(), 2) + + def test_count_reset(self): + self.connection.sendmail( + 'anne@example.com', ['bart@example.com'], self.msg_text) + self.connection.quit() + self.connection.sendmail( + 'cate@example.com', ['dave@example.com'], self.msg_text) + self.connection.quit() + # Issue the fake SMTP command to reset the count. + client = SMTP() + client.connect(config.mta.smtp_host, int(config.mta.smtp_port)) + client.docmd('RSET') + self.assertEqual(SMTPLayer.smtpd.get_connection_count(), 0) diff --git a/src/mailman/testing/mta.py b/src/mailman/testing/mta.py index a4ae01763..f7a6a9806 100644 --- a/src/mailman/testing/mta.py +++ b/src/mailman/testing/mta.py @@ -52,6 +52,7 @@ class ConnectionCountingHandler(MessageHandler): def __init__(self, msg_queue): super().__init__() self._msg_queue = msg_queue + self.connection_count = 0 def handle_message(self, message): self._msg_queue.put(message) @@ -62,11 +63,18 @@ class ConnectionCountingSMTP(SMTP): super().__init__(handler, *args, **kws) self._auth_response = None self._waiting_for_auth_response = False - self._connection_count = 0 self._oob_queue = oob_queue self._err_queue = err_queue self._last_error = None + def connection_made(self, transport): + super().connection_made(transport) + # We can't keep the connection count on self here because the + # controller (via the factory() method) will create a new instance of + # this class for every connection. The handler instance is always the + # same though, so it's fine to stash this value away there. + self.event_handler.connection_count += 1 + @asyncio.coroutine def smtp_AUTH(self, arg): """Record that the AUTH occurred.""" @@ -92,20 +100,41 @@ class ConnectionCountingSMTP(SMTP): @asyncio.coroutine def smtp_EHLO(self, arg): - yield from super().smtp_EHLO(arg) - # If the upcall succeeded, this flag will be set. In that case, also - # push an AUTH PLAIN response, which the superclass doesn't do. - ## if self.extended_smtp: - ## yield from self.push('250 AUTH PLAIN') + if not arg: + yield from self.push('501 Syntax: EHLO hostname') + return + # See issue #21783 for a discussion of this behavior. + if self.seen_greeting: + yield from self.push('503 Duplicate HELO/EHLO') + return + self._set_rset_state() + self.seen_greeting = arg + self.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) + self.command_size_limits['MAIL'] += 26 + if not self._decode_data: + yield from self.push('250-8BITMIME') + if self.enable_SMTPUTF8: + yield from self.push('250-SMTPUTF8') + self.command_size_limits['MAIL'] += 10 + yield from self.push('250-HELP') + yield from self.push('250 AUTH PLAIN') @asyncio.coroutine def smtp_STAT(self, arg): """Cause the server to send statistics to its controller.""" # Do not count the connection caused by the STAT connect. - self._connection_count -= 1 - self._oob_queue.put(self._connection_count) + self.event_handler.connection_count -= 1 + self._oob_queue.put(self.event_handler.connection_count) yield from self.push('250 Ok') + @asyncio.coroutine + def smtp_RSET(self, arg): + yield from super().smtp_RSET(arg) + self.event_handler.connection_count = 0 + def _next_error(self, command): """Return the next error for the SMTP command, if there is one. |