Repository: aio-libs/aiosmtpd
Pull request #146 (2018-06-13 12:47:50)

Add a timeout feature
Client will now be disconnected after a set period of time, which can be
adjusted.  This prevents one form of resource exhaustion when badly
behaving clients connect and never send any data.

This partially fixes issue #145 


Patch:
diff --git a/aiosmtpd/docs/smtp.rst b/aiosmtpd/docs/smtp.rst
index 529865a..1dbeca3 100644
--- a/aiosmtpd/docs/smtp.rst
+++ b/aiosmtpd/docs/smtp.rst
@@ -101,7 +101,7 @@ override to provide additional responses.
 SMTP API
 ========
 
-.. class:: SMTP(handler, *, data_size_limit=33554432, enable_SMTPUTF8=False, decode_data=False, hostname=None, ident=None, tls_context=None, require_starttls=False, loop=None)
+.. class:: SMTP(handler, *, data_size_limit=33554432, enable_SMTPUTF8=False, decode_data=False, hostname=None, ident=None, tls_context=None, require_starttls=False, timeout=300, loop=None)
 
    *handler* is an instance of a :ref:`handler <handlers>` class.
 
@@ -120,7 +120,7 @@ SMTP API
    *hostname* is the first part of the string returned in the ``220`` greeting
    response given to clients when they first connect to the server.  If not given,
    the system's fully-qualified domain name is used.
-   
+
    *ident* is the second part of the string returned in the ``220`` greeting
    response that identifies the software name and version of the SMTP server
    to the client. If not given, a default Python SMTP ident is used.
@@ -131,6 +131,10 @@ SMTP API
    and *require_starttls* should be ``True``.  See :ref:`tls` for a more in
    depth discussion on enabling ``STARTTLS``.
 
+   *timeout* is the number of seconds to wait between valid SMTP commands.
+   After this time the connection will be closed by the server.  The default
+   is 300 seconds, as per `RFC 2821`_.
+
    *loop* is the asyncio event loop to use.  If not given,
    :meth:`asyncio.new_event_loop()` is called to create the event loop.
 
@@ -253,8 +257,8 @@ advertised, and the ``STARTTLS`` command will not be accepted.
 *require_starttls* is meaningless in this case, and should be set to
 ``False``.
 
-
 .. _StreamReaderProtocol: https://docs.python.org/3/library/asyncio-stream.html#streamreaderprotocol
 .. _`RFC 3207`: http://www.faqs.org/rfcs/rfc3207.html
+.. _`RFC 2821`: https://www.ietf.org/rfc/rfc2821.txt
 .. _`asyncio transport`: https://docs.python.org/3/library/asyncio-protocol.html#asyncio-transport
 .. _ssl: https://docs.python.org/3/library/ssl.html
diff --git a/aiosmtpd/smtp.py b/aiosmtpd/smtp.py
index 681732a..441c73d 100644
--- a/aiosmtpd/smtp.py
+++ b/aiosmtpd/smtp.py
@@ -75,6 +75,7 @@ class SMTP(asyncio.StreamReaderProtocol):
                  ident=None,
                  tls_context=None,
                  require_starttls=False,
+                 timeout=300,
                  loop=None):
         self.__ident__ = ident or __ident__
         self.loop = loop if loop else make_loop()
@@ -98,6 +99,8 @@ class SMTP(asyncio.StreamReaderProtocol):
             self.tls_context.check_hostname = False
             self.tls_context.verify_mode = ssl.CERT_NONE
         self.require_starttls = tls_context and require_starttls
+        self._timeout_duration = timeout
+        self._timeout_handle = None
         self._tls_handshake_okay = True
         self._tls_protocol = None
         self._original_transport = None
@@ -131,6 +134,7 @@ class SMTP(asyncio.StreamReaderProtocol):
         self._set_rset_state()
         self.session = self._create_session()
         self.session.peer = transport.get_extra_info('peername')
+        self._reset_timeout()
         seen_starttls = (self._original_transport is not None)
         if self.transport is not None and seen_starttls:
             # It is STARTTLS connection over normal connection.
@@ -156,6 +160,7 @@ class SMTP(asyncio.StreamReaderProtocol):
 
     def connection_lost(self, error):
         log.info('%r connection lost', self.session.peer)
+        self._timeout_handle.cancel()
         # If STARTTLS was issued, then our transport is the SSL protocol
         # transport, and we need to close the original transport explicitly,
         # otherwise an unexpected eof_received() will be called *after* the
@@ -182,6 +187,21 @@ class SMTP(asyncio.StreamReaderProtocol):
             return False
         return super().eof_received()
 
+    def _reset_timeout(self):
+        if self._timeout_handle is not None:
+            self._timeout_handle.cancel()
+
+        self._timeout_handle = self.loop.call_later(
+                self._timeout_duration, self._timeout_cb)
+
+    def _timeout_cb(self):
+        log.info('%r connection timeout', self.session.peer)
+
+        # Calling close() on the transport will trigger connection_lost(),
+        # which gracefully closes the SSL transport if required and cleans
+        # up state.
+        self.transport.close()
+
     def _client_connected_cb(self, reader, writer):
         # This is redundant since we subclass StreamReaderProtocol, but I like
         # the shorter names.
@@ -279,6 +299,9 @@ class SMTP(asyncio.StreamReaderProtocol):
                     await self.push(
                         '500 Error: command "%s" not recognized' % command)
                     continue
+
+                # Received a valid command, reset the timer.
+                self._reset_timeout()
                 await method(arg)
             except asyncio.CancelledError:
                 # The connection got reset during the DATA command.


Test patch:
diff --git a/aiosmtpd/tests/test_smtp.py b/aiosmtpd/tests/test_smtp.py
index be1da96..fc3ea40 100644
--- a/aiosmtpd/tests/test_smtp.py
+++ b/aiosmtpd/tests/test_smtp.py
@@ -1,5 +1,6 @@
 """Test the SMTP protocol."""
 
+import time
 import socket
 import asyncio
 import unittest
@@ -27,6 +28,11 @@ class NoDecodeController(Controller):
         return Server(self.handler, decode_data=False)
 
 
+class TimeoutController(Controller):
+    def factory(self):
+        return Server(self.handler, timeout=0.1)
+
+
 class ReceivingHandler:
     box = None
 
@@ -1155,3 +1161,17 @@ class TestSleepingHandler(unittest.TestCase):
             client.send('HELO example.com\r\n')
             client.sock.shutdown(socket.SHUT_WR)
             self.assertRaises(SMTPServerDisconnected, client.getreply)
+
+
+class TestTimeout(unittest.TestCase):
+    def setUp(self):
+        controller = TimeoutController(Sink)
+        controller.start()
+        self.addCleanup(controller.stop)
+        self.address = (controller.hostname, controller.port)
+
+    def test_timeout(self):
+        with SMTP(*self.address) as client:
+            code, response = client.ehlo('example.com')
+            time.sleep(0.3)
+            self.assertRaises(SMTPServerDisconnected, client.getreply)