Repository: aio-libs/aiosmtpd

Issue:
Misbehaving clients are never disconnected
It's possible to exhaust resources on the server in a number of ways:

* Connect and never send a valid command
* Connect and repeatedly send NOOP
* Connect and waste time in the SMTP transaction, e.g. send multiple HELO, reset, 

Eventually Python will run out of file handles and exceptions will occur in code that tries to open any file, socket, etc.

Comments:

Patch:
diff --git a/README.rst b/README.rst
index da032ab..a5514fc 100644
--- a/README.rst
+++ b/README.rst
@@ -184,6 +184,24 @@ versions on your system by using ``pyenv``. General steps:
    documentation about this option)
 
 
+``housekeep.py``
+-------------------
+
+If you ever need to 'reset' your repo, you can use the ``housekeep.py`` utility
+like so::
+
+    $ python housekeep.py superclean
+
+It is `strongly` recommended to NOT do superclean too often, though.
+Every time you invoke ``superclean``,
+tox will have to recreate all its testenvs,
+and this will make testing `much` longer to finish.
+
+``superclean`` is typically only needed when you switch branches,
+or if you want to really ensure that artifacts from previous testing sessions
+won't interfere with your next testing sessions.
+
+
 Contents
 ========
 
diff --git a/aiosmtpd/docs/NEWS.rst b/aiosmtpd/docs/NEWS.rst
index 4b0cab1..49abd45 100644
--- a/aiosmtpd/docs/NEWS.rst
+++ b/aiosmtpd/docs/NEWS.rst
@@ -9,6 +9,7 @@
 Added
 -----
 * Test for ``SMTP.__init__`` behavior after taking out code that edits TLS Context
+* Implement mechanism to limit the number of commands sent (Closes #145)
 
 Fixed/Improved
 --------------
diff --git a/aiosmtpd/docs/smtp.rst b/aiosmtpd/docs/smtp.rst
index 84694c5..2d3a3eb 100644
--- a/aiosmtpd/docs/smtp.rst
+++ b/aiosmtpd/docs/smtp.rst
@@ -107,7 +107,9 @@ SMTP API
    decode_data=False, hostname=None, ident=None, tls_context=None, \
    require_starttls=False, timeout=300, auth_required=False, \
    auth_require_tls=True, auth_exclude_mechanism=None, auth_callback=None, \
-   loop=None)
+   command_call_limit=None, loop=None)
+
+   **Parameters**
 
    :boldital:`handler` is an instance of a :ref:`handler <handlers>` class.
 
@@ -158,9 +160,51 @@ SMTP API
    must return a ``bool`` that indicates whether the client's authentication
    attempt is accepted/successful or not.
 
+   :boldital:`command_call_limit` if not ``None`` sets the maximum time a certain SMTP
+   command can be invoked.
+   This is to prevent DoS due to malicious client connecting and never disconnecting,
+   due to continual sending of SMTP commands to prevent timeout.
+
+   The handling differs based on the type:
+
+   .. highlights::
+
+      If :attr:`command_call_limit` is of type ``int``,
+      then the value is the call limit for ALL SMTP commands.
+
+      If :attr:`command_call_limit` is of type ``dict``,
+      it must be a ``Dict[str, int]``
+      (the type of the values will be enforced).
+      The keys will be the SMTP Command to set the limit for,
+      the values will be the call limit per SMTP Command.
+
+      .. highlights::
+
+         A special key of ``"*"`` is used to set the 'default' call limit for commands not
+         explicitly declared in :attr:`command_call_limit`.
+         If ``"*"`` is not given,
+         then the 'default' call limit will be set to ``aiosmtpd.smtp.CALL_LIMIT_DEFAULT``
+
+      Other types -- or a ``Dict`` whose any value is not an ``int`` -- will raise a
+      ``TypeError`` exception.
+
+      Examples::
+
+          # All commands have a limit of 10 calls
+          SMTP(..., command_call_limit=10)
+
+          # Commands RCPT and NOOP have their own limits; others have an implicit limit
+          # of 20 (CALL_LIMIT_DEFAULT)
+          SMTP(..., command_call_limit={"RCPT": 30, "NOOP": 5})
+
+          # Commands RCPT and NOOP have their own limits; others set to 3
+          SMTP(..., command_call_limit={"RCPT": 20, "NOOP": 10, "*": 3})
+
    :boldital:`loop` is the asyncio event loop to use.  If not given,
    :meth:`asyncio.new_event_loop()` is called to create the event loop.
 
+   **Attributes**
+
    .. py:attribute:: line_length_limit
 
       The maximum line length, in octets (not characters; one UTF-8 character
diff --git a/aiosmtpd/lmtp.py b/aiosmtpd/lmtp.py
index 5a7ba40..a9ba9d7 100644
--- a/aiosmtpd/lmtp.py
+++ b/aiosmtpd/lmtp.py
@@ -4,7 +4,7 @@ from public import public
 
 @public
 class LMTP(SMTP):
-    show_smtp_greeting = False
+    show_smtp_greeting: bool = False
 
     @syntax('LHLO hostname')
     async def smtp_LHLO(self, arg):
diff --git a/aiosmtpd/smtp.py b/aiosmtpd/smtp.py
index 5447640..0381184 100644
--- a/aiosmtpd/smtp.py
+++ b/aiosmtpd/smtp.py
@@ -59,11 +59,13 @@ __all__ = [
     "AuthMechanismType",
     "MISSING",
 ]  # Will be added to by @public
-__version__ = '1.2.3a5'
+__version__ = '1.2.3a6'
 __ident__ = 'Python SMTP {}'.format(__version__)
 log = logging.getLogger('mail.log')
 
 
+BOGUS_LIMIT = 5
+CALL_LIMIT_DEFAULT = 20
 DATA_SIZE_DEFAULT = 2**25  # Where does this number come from, I wonder...
 EMPTY_BARR = bytearray()
 EMPTYBYTES = b''
@@ -151,6 +153,11 @@ def login_always_fail(mechanism, login, password):
     return False
 
 
+def is_int(o):
+    return isinstance(o, int)
+
+
+@public
 class TLSSetupException(Exception):
     pass
 
@@ -181,6 +188,7 @@ class SMTP(asyncio.StreamReaderProtocol):
             auth_require_tls=True,
             auth_exclude_mechanism: Optional[Iterable[str]] = None,
             auth_callback: AuthCallbackType = None,
+            command_call_limit: Union[int, Dict[str, int], None] = None,
             loop=None,
     ):
         self.__ident__ = ident or __ident__
@@ -256,6 +264,23 @@ class SMTP(asyncio.StreamReaderProtocol):
             if m.startswith("smtp_")
         }
 
+        if command_call_limit is None:
+            self._enforce_call_limit = False
+        else:
+            self._enforce_call_limit = True
+            if isinstance(command_call_limit, int):
+                self._call_limit_base = {}
+                self._call_limit_default: int = command_call_limit
+            elif isinstance(command_call_limit, dict):
+                if not all(map(is_int, command_call_limit.values())):
+                    raise TypeError("All command_call_limit values must be int")
+                self._call_limit_base = command_call_limit
+                self._call_limit_default: int = command_call_limit.get(
+                    "*", CALL_LIMIT_DEFAULT
+                )
+            else:
+                raise TypeError("command_call_limit must be int or Dict[str, int]")
+
     def _create_session(self):
         return Session(self.loop)
 
@@ -378,6 +403,15 @@ class SMTP(asyncio.StreamReaderProtocol):
     async def _handle_client(self):
         log.info('%r handling connection', self.session.peer)
         await self.push('220 {} {}'.format(self.hostname, self.__ident__))
+        if self._enforce_call_limit:
+            call_limit = collections.defaultdict(
+                lambda x=self._call_limit_default: x,
+                self._call_limit_base
+            )
+        else:
+            # Not used, but this silences code inspection tools
+            call_limit = {}
+        bogus_budget = BOGUS_LIMIT
         while self.transport is not None:   # pragma: nobranch
             try:
                 try:
@@ -454,8 +488,31 @@ class SMTP(asyncio.StreamReaderProtocol):
                     # RFC3207 part 4
                     await self.push('530 Must issue a STARTTLS command first')
                     continue
+
+                if self._enforce_call_limit:
+                    budget = call_limit[command]
+                    if budget < 1:
+                        log.warning(
+                            "%r over limit for %s", self.session.peer, command
+                        )
+                        await self.push(
+                            f"421 4.7.0 {command} sent too many times"
+                        )
+                        self.transport.close()
+                        continue
+                    call_limit[command] = budget - 1
+
                 method = self._smtp_methods.get(command)
                 if method is None:
+                    log.warning("%r unrecognised: %s", self.session.peer, command)
+                    bogus_budget -= 1
+                    if bogus_budget < 1:
+                        log.warning("%r too many bogus commands", self.session.peer)
+                        await self.push(
+                            "502 5.5.1 Too many unrecognized commands, goodbye."
+                        )
+                        self.transport.close()
+                        continue
                     await self.push(
                         '500 Error: command "%s" not recognized' % command)
                     continue


Test patch:
diff --git a/aiosmtpd/tests/test_smtp.py b/aiosmtpd/tests/test_smtp.py
index 04f541a..6c18c58 100644
--- a/aiosmtpd/tests/test_smtp.py
+++ b/aiosmtpd/tests/test_smtp.py
@@ -9,7 +9,11 @@ import warnings
 from aiosmtpd.controller import Controller
 from aiosmtpd.handlers import Sink
 from aiosmtpd.smtp import (
-    MISSING, Session as SMTPSess, SMTP as Server, __ident__ as GREETING, auth_mechanism
+    CALL_LIMIT_DEFAULT,
+    MISSING,
+    Session as SMTPSess,
+    SMTP as Server,
+    __ident__ as GREETING, auth_mechanism
 )
 from aiosmtpd.testing.helpers import (
     ReceivingHandler,
@@ -803,7 +807,8 @@ class TestSMTP(unittest.TestCase):
             self.assertEqual(code, 250)
             self.assertEqual(response, b'OK')
 
-    def test_unknown_command(self):
+    @patch("logging.Logger.warning")
+    def test_unknown_command(self, mock_warning):
         with SMTP(*self.address) as client:
             code, response = client.docmd('FOOBAR')
             self.assertEqual(code, 500)
@@ -1888,3 +1893,158 @@ class TestAuthArgs(unittest.TestCase):
         self.assertRaises(ValueError, auth_mechanism, "has space")
         self.assertRaises(ValueError, auth_mechanism, "has.dot")
         self.assertRaises(ValueError, auth_mechanism, "has/slash")
+
+
+class TestLimits(unittest.TestCase):
+    @patch("logging.Logger.warning")
+    def test_all_limit_15(self, mock_warning):
+        kwargs = dict(
+            command_call_limit=15,
+        )
+        controller = Controller(Sink(), server_kwargs=kwargs)
+        self.addCleanup(controller.stop)
+        controller.start()
+        with SMTP(controller.hostname, controller.port) as client:
+            code, mesg = client.ehlo('example.com')
+            self.assertEqual(250, code)
+            for _ in range(0, 15):
+                code, mesg = client.noop()
+                self.assertEqual(250, code)
+            code, mesg = client.noop()
+            self.assertEqual(421, code)
+            self.assertEqual(b"4.7.0 NOOP sent too many times", mesg)
+            with self.assertRaises(SMTPServerDisconnected):
+                client.noop()
+
+    @patch("logging.Logger.warning")
+    def test_different_limits(self, mock_warning):
+        noop_max, expn_max = 15, 5
+        kwargs = dict(
+            command_call_limit={"NOOP": noop_max, "EXPN": expn_max},
+        )
+        controller = Controller(Sink(), server_kwargs=kwargs)
+        self.addCleanup(controller.stop)
+        controller.start()
+        with SMTP(controller.hostname, controller.port) as client:
+            code, mesg = client.ehlo('example.com')
+            self.assertEqual(250, code)
+            for _ in range(0, noop_max):
+                code, mesg = client.noop()
+                self.assertEqual(250, code)
+            code, mesg = client.noop()
+            self.assertEqual(421, code)
+            self.assertEqual(b"4.7.0 NOOP sent too many times", mesg)
+            with self.assertRaises(SMTPServerDisconnected):
+                client.noop()
+        with SMTP(controller.hostname, controller.port) as client:
+            code, mesg = client.ehlo('example.com')
+            self.assertEqual(250, code)
+            for _ in range(0, expn_max):
+                code, mesg = client.expn("alice@example.com")
+                self.assertEqual(502, code)
+            code, mesg = client.expn("alice@example.com")
+            self.assertEqual(421, code)
+            self.assertEqual(b"4.7.0 EXPN sent too many times", mesg)
+            with self.assertRaises(SMTPServerDisconnected):
+                client.noop()
+        with SMTP(controller.hostname, controller.port) as client:
+            code, mesg = client.ehlo('example.com')
+            self.assertEqual(250, code)
+            for _ in range(0, CALL_LIMIT_DEFAULT):
+                code, mesg = client.vrfy("alice@example.com")
+                self.assertEqual(252, code)
+            code, mesg = client.vrfy("alice@example.com")
+            self.assertEqual(421, code)
+            self.assertEqual(b"4.7.0 VRFY sent too many times", mesg)
+            with self.assertRaises(SMTPServerDisconnected):
+                client.noop()
+
+    @patch("logging.Logger.warning")
+    def test_different_limits_custom_default(self, mock_warning):
+        # Important: make sure default_max > CALL_LIMIT_DEFAULT
+        # Others can be set small to cut down on testing time, but must be different
+        noop_max, expn_max, default_max = 7, 5, 25
+        self.assertGreater(default_max, CALL_LIMIT_DEFAULT)
+        self.assertNotEqual(noop_max, expn_max)
+        kwargs = dict(
+            command_call_limit={"NOOP": noop_max, "EXPN": expn_max, "*": default_max},
+        )
+        controller = Controller(Sink(), server_kwargs=kwargs)
+        self.addCleanup(controller.stop)
+        controller.start()
+        with SMTP(controller.hostname, controller.port) as client:
+            code, mesg = client.ehlo('example.com')
+            self.assertEqual(250, code)
+            for _ in range(0, noop_max):
+                code, mesg = client.noop()
+                self.assertEqual(250, code)
+            code, mesg = client.noop()
+            self.assertEqual(421, code)
+            self.assertEqual(b"4.7.0 NOOP sent too many times", mesg)
+            with self.assertRaises(SMTPServerDisconnected):
+                client.noop()
+        with SMTP(controller.hostname, controller.port) as client:
+            code, mesg = client.ehlo('example.com')
+            self.assertEqual(250, code)
+            for _ in range(0, expn_max):
+                code, mesg = client.expn("alice@example.com")
+                self.assertEqual(502, code)
+            code, mesg = client.expn("alice@example.com")
+            self.assertEqual(421, code)
+            self.assertEqual(b"4.7.0 EXPN sent too many times", mesg)
+            with self.assertRaises(SMTPServerDisconnected):
+                client.noop()
+        with SMTP(controller.hostname, controller.port) as client:
+            code, mesg = client.ehlo('example.com')
+            self.assertEqual(250, code)
+            for _ in range(0, default_max):
+                code, mesg = client.vrfy("alice@example.com")
+                self.assertEqual(252, code)
+            code, mesg = client.vrfy("alice@example.com")
+            self.assertEqual(421, code)
+            self.assertEqual(b"4.7.0 VRFY sent too many times", mesg)
+            with self.assertRaises(SMTPServerDisconnected):
+                client.noop()
+
+    def test_limit_wrong_type(self):
+        kwargs = dict(
+            command_call_limit="invalid",
+        )
+        controller = Controller(Sink(), server_kwargs=kwargs)
+        self.addCleanup(controller.stop)
+        with self.assertRaises(TypeError):
+            controller.start()
+
+    def test_limit_wrong_value_type(self):
+        kwargs = dict(
+            command_call_limit={"NOOP": "invalid"},
+        )
+        controller = Controller(Sink(), server_kwargs=kwargs)
+        self.addCleanup(controller.stop)
+        with self.assertRaises(TypeError):
+            controller.start()
+
+    @patch("logging.Logger.warning")
+    def test_limit_bogus(self, mock_warning):
+        # Extreme limit.
+        kwargs = dict(
+            command_call_limit=1,
+        )
+        controller = Controller(Sink(), server_kwargs=kwargs)
+        self.addCleanup(controller.stop)
+        controller.start()
+        with SMTP(controller.hostname, controller.port) as client:
+            code, mesg = client.ehlo('example.com')
+            self.assertEqual(250, code)
+            for i in range(0, 4):
+                code, mesg = client.docmd(f"BOGUS{i}")
+                self.assertEqual(500, code)
+                expected = f"Error: command \"BOGUS{i}\" not recognized"
+                self.assertEqual(expected, mesg.decode("ascii"))
+            code, mesg = client.docmd("LASTBOGUS")
+            self.assertEqual(502, code)
+            self.assertEqual(
+                b"5.5.1 Too many unrecognized commands, goodbye.", mesg
+            )
+            with self.assertRaises(SMTPServerDisconnected):
+                client.noop()