Lots more TLS tests and graceful shutdown fixes

Author: njsmith

trio/ssl.py

@@ -220,6 +220,7 @@ class SSLStream(_Stream):
     def __init__(
             self, transport_stream, sslcontext, *, max_bytes=32 * 1024, **kwargs):
         self.transport_stream = transport_stream
+        self._exc = None
         self._bufsize = max_bytes
         self._outgoing = _stdlib_ssl.MemoryBIO()
         self._incoming = _stdlib_ssl.MemoryBIO()
@@ -234,7 +235,8 @@ def __init__(
         self._inner_recv_lock = _sync.Lock()
 
         # These are used to make sure that our caller doesn't attempt to make
-        # multiple concurrent calls to send_all/wait_send_all_might_not_block or to receive_some.
+        # multiple concurrent calls to send_all/wait_send_all_might_not_block
+        # or to receive_some.
         self._outer_send_lock = _UnLock(
             _core.ResourceBusyError,
             "another task is currently sending data on this SSLStream")
@@ -263,11 +265,16 @@ def __setattr__(self, name, value):
     def __dir__(self):
         return super().__dir__() + list(self._forwarded)
 
+    def _check_status(self):
+        if self._exc is not None:
+            raise self._exc
+
     # This is probably the single trickiest function in trio. It has lots of
     # comments, though, just make sure to think carefully if you ever have to
     # touch it. The big comment at the top of this file will help explain
     # too.
     async def _retry(self, fn, *args, ignore_want_read=False):
+        print("doing", fn)
         await _core.yield_if_cancelled()
         yielded = False
         try:
@@ -288,11 +295,11 @@ async def _retry(self, fn, *args, ignore_want_read=False):
                 # might come in and mess with it while we're suspended), and
                 # we don't want to yield *before* starting the operation that
                 # will help us make progress, because then someone else might
-                # come in and
+                # come in and leapfrog us.
 
                 # Call the SSLObject method, and get its result.
                 #
-                # NB: despite what the docs, say SSLWantWriteError can't
+                # NB: despite what the docs say, SSLWantWriteError can't
                 # happen – "Writes to memory BIOs will always succeed if
                 # memory is available: that is their size can grow
                 # indefinitely."
@@ -303,14 +310,16 @@ async def _retry(self, fn, *args, ignore_want_read=False):
                     ret = fn(*args)
                 except _stdlib_ssl.SSLWantReadError:
                     want_read = True
-                except _stdlib_ssl.SSLError as exc:
+                except (SSLError, CertificateError) as exc:
+                    self._exc = _streams.BrokenStreamError
                     raise _streams.BrokenStreamError from exc
                 else:
                     finished = True
                 if ignore_want_read:
                     want_read = False
                     finished = True
                 to_send = self._outgoing.read()
+                print(bool(to_send), want_read)
 
                 # Outputs from the above code block are:
                 #
@@ -373,7 +382,13 @@ async def _retry(self, fn, *args, ignore_want_read=False):
                     # NOTE: This relies on the lock being strict FIFO fair!
                     async with self._inner_send_lock:
                         yielded = True
-                        await self.transport_stream.send_all(to_send)
+                        try:
+                            await self.transport_stream.send_all(to_send)
+                        except:
+                            # Some unknown amount of our data got sent, and we
+                            # don't know how much. This stream is doomed.
+                            self._exc = _streams.BrokenStreamError
+                            raise
                 elif want_read:
                     # It's possible that someone else is already blocked in
                     # transport_stream.receive_some. If so then we want to
@@ -427,22 +442,24 @@ async def do_handshake(self):
         immediately without doing anything (except executing a checkpoint).
 
         """
-        if self.transport_stream is None:
+        try:
+            self._check_status()
+        except:
             await _core.yield_briefly()
-            raise _streams.ClosedStreamError
+            raise
         await self._handshook.ensure(checkpoint=True)
 
     # Most things work if we don't explicitly force do_handshake to be called
-    # before calling receive_some or send_all, because openssl will automatically
-    # perform the handshake on the first SSL_{read,write} call. BUT, allowing
-    # openssl to do this will disable Python's hostname checking!!! See:
+    # before calling receive_some or send_all, because openssl will
+    # automatically perform the handshake on the first SSL_{read,write}
+    # call. BUT, allowing openssl to do this will disable Python's hostname
+    # checking!!! See:
     #   https://bugs.python.org/issue30141
     # So we *definitely* have to make sure that do_handshake is called
     # before doing anything else.
     async def receive_some(self, max_bytes):
         async with self._outer_recv_lock:
-            if self.transport_stream is None:
-                raise _streams.ClosedStreamError
+            self._check_status()
             await self._handshook.ensure(checkpoint=False)
             max_bytes = _operator.index(max_bytes)
             if max_bytes < 1:
@@ -451,8 +468,7 @@ async def receive_some(self, max_bytes):
 
     async def send_all(self, data):
         async with self._outer_send_lock:
-            if self.transport_stream is None:
-                raise _streams.ClosedStreamError
+            self._check_status()
             await self._handshook.ensure(checkpoint=False)
             # SSLObject interprets write(b"") as an EOF for some reason, which
             # is not what we want.
@@ -471,65 +487,82 @@ async def send_all(self, data):
     # maybe it's actually better to error out...?
     async def unwrap(self):
         async with self._outer_recv_lock, self._outer_send_lock:
-            if self.transport_stream is None:
-                raise _streams.ClosedStreamError
+            self._check_status()
             await self._handshook.ensure(checkpoint=False)
             await self._retry(self._ssl_object.unwrap)
             transport_stream = self.transport_stream
             self.transport_stream = None
+            self._exc = _streams.ClosedStreamError
             return (transport_stream, self._incoming.read())
 
     def forceful_close(self):
-        if self.transport_stream is not None:
+        if self._exc is not _streams.ClosedStreamError:
             self.transport_stream.forceful_close()
-            self.transport_stream = None
+            self._exc = _streams.ClosedStreamError
 
     async def graceful_close(self):
-        transport_stream = self.transport_stream
-        if transport_stream is None:
+        if self._exc is _streams.ClosedStreamError:
+            await _core.yield_briefly()
+            return
+        if self._exc is _streams.BrokenStreamError:
+            self.forceful_close()
             await _core.yield_briefly()
             return
         try:
-            # If we haven't even started the handshake, then we can (and must)
-            # skip the SSL-level shutdown.
-            if self._handshook.started:
-                # But if the handshake is in progress, wait for it to finish.
-                await self._handshook.ensure(checkpoint=False)
-                # Then we want to call SSL_shutdown *once*, to send a
-                # close_notify but *not* wait for the response (because we're
-                # closing the socket anyway, so there's no point in waiting).
-                # Subtlety: SSLObject.unwrap will immediately call it a second
-                # time, and the second time will raise SSLWantReadError
-                # because there hasn't been time for the other side to respond
-                # yet. (Unless they spontaneously sent a close_notify before
-                # we called this, and it's either already been processed or
-                # gets pulled out of the buffer by Python's second call.) So
-                # the way to do what we want is to to ignore SSLWantReadError
-                # on this call.
-                try:
-                    await self._retry(
-                        self._ssl_object.unwrap, ignore_want_read=True)
-                except _streams.BrokenStreamError:
-                    # It's okay if the stream is broken and we can't send our
-                    # goodbye message, because we're cutting off the
-                    # connection anyway...
-                    pass
+            await self._handshook.ensure(checkpoint=False)
+            # Here, we call SSL_shutdown *once*, because we want to send a
+            # close_notify but *not* wait for the other side to send back a
+            # response. In principle it would be more polite to wait for the
+            # other side to reply with their own close_notify. However, if
+            # they aren't paying attention (e.g., if they're just sending
+            # data and not receiving) then we will never notice our
+            # close_notify and we'll be waiting forever. Eventually we'll time
+            # out (hopefully), but it's still kind of nasty. And we can't
+            # require the other side to always be receiving, because (a)
+            # backpressure is kind of important, and (b) I bet there are
+            # broken TLS implementations out there that don't receive all the
+            # time. (Like e.g. anyone using Python ssl in synchronous mode.)
+            #
+            # The send-then-immediately-close behavior is explicitly allowed
+            # by the TLS specs, so we're ok on that.
+            #
+            # Subtlety: SSLObject.unwrap will immediately call it a second
+            # time, and the second time will raise SSLWantReadError because
+            # there hasn't been time for the other side to respond
+            # yet. (Unless they spontaneously sent a close_notify before we
+            # called this, and it's either already been processed or gets
+            # pulled out of the buffer by Python's second call.) So the way to
+            # do what we want is to ignore SSLWantReadError on this call.
+            #
+            # Also, because the other side might have already sent
+            # close_notify and closed their connection then it's possible that
+            # our attempt to send close_notify will raise
+            # BrokenStreamError. This is totally legal, and in fact can happen
+            # with two well-behaved trio programs talking to each other, so we
+            # don't want to raise an error. So we suppress BrokenStreamError
+            # here. (This is safe, because literally the only thing this call
+            # to _retry will do is send the close_notify alert, so that's
+            # surely where the error comes from.)
+            try:
+                await self._retry(
+                    self._ssl_object.unwrap, ignore_want_read=True)
+            except _streams.BrokenStreamError:
+                pass
             # Close the underlying stream
-            await transport_stream.graceful_close()
+            await self.transport_stream.graceful_close()
         except:
-            transport_stream.forceful_close()
+            self.transport_stream.forceful_close()
             raise
         finally:
-            self.transport_stream = None
+            self._exc = _streams.ClosedStreamError
 
     async def wait_send_all_might_not_block(self):
         # This method's implementation is deceptively simple.
         #
         # First, we take the outer send lock, because of trio's standard
         # semantics that wait_send_all_might_not_block and send_all conflict.
         async with self._outer_send_lock:
-            if self.transport_stream is None:
-                raise _streams.ClosedStreamError
+            self._check_status()
             # Then we take the inner send lock. We know that no other tasks
             # are calling self.send_all or self.wait_send_all_might_not_block,
             # because we have the outer_send_lock. But! There might be another

trio/testing.py

@@ -506,13 +506,16 @@ async def simple_check_wait_send_all_might_not_block(scope):
                           nursery.cancel_scope)
             nursery.spawn(do_receive_some, 1)
 
-        # closing the r side
-        await do_graceful_close(r)
+        # closing the r side leads to BrokenStreamError on the s side
+        # (eventually)
+        async def expect_broken_stream_on_send():
+            with _assert_raises(_streams.BrokenStreamError):
+                while True:
+                    await do_send_all(b"x" * 100)
 
-        # ...leads to BrokenStreamError on the s side (eventually)
-        with _assert_raises(_streams.BrokenStreamError):
-            while True:
-                await do_send_all(b"x" * 100)
+        async with _core.open_nursery() as nursery:
+            nursery.spawn(expect_broken_stream_on_send)
+            nursery.spawn(do_graceful_close, r)
 
         # once detected, the stream stays broken
         with _assert_raises(_streams.BrokenStreamError):
@@ -555,6 +558,7 @@ async def receive_send_then_close():
             await wait_all_tasks_blocked()
             await checked_receive_1(b"y")
             await checked_receive_1(b"")
+            await do_graceful_close(r)
 
         async with _core.open_nursery() as nursery:
             nursery.spawn(send_then_close)
@@ -621,8 +625,9 @@ async def expect_cancelled(afn, *args):
                 nursery.spawn(expect_cancelled, do_send_all, b"x")
                 nursery.spawn(expect_cancelled, do_receive_some, 1)
 
-        await do_graceful_close(s)
-        await do_graceful_close(r)
+        async with _core.open_nursery() as nursery:
+            nursery.spawn(do_graceful_close, s)
+            nursery.spawn(do_graceful_close, r)
 
     # check wait_send_all_might_not_block, if we can
     if clogged_stream_maker is not None:
@@ -739,9 +744,13 @@ async def receiver(s, data, seed):
             nursery.spawn(receiver, s1, test_data[::-1], 2)
             nursery.spawn(receiver, s2, test_data, 3)
 
-        await s1.graceful_close()
-        assert await s2.receive_some(10) == b""
-        await s2.graceful_close()
+        async def expect_receive_some_empty():
+            assert await s2.receive_some(10) == b""
+            await s2.graceful_close()
+
+        async with _core.open_nursery() as nursery:
+            nursery.spawn(expect_receive_some_empty)
+            nursery.spawn(s1.graceful_close)
 
 
 async def check_half_closeable_stream(stream_maker, clogged_stream_maker):
@@ -968,19 +977,23 @@ class MemoryReceiveStream(_abc.ReceiveStream):
     Args:
       receive_some_hook: An async function, or None. Called from
           :meth:`receive_some`. Can do whatever you like.
+      close_hook: A synchronous function, or None. Called from
+          :meth:`forceful_close`. Can do whatever you like.
 
     .. attribute:: receive_some_hook
+                   close_hook
 
-       The :attr:`receive_some_hook` is also exposed as an attribute on the
-       object, and you can change it at any time.
+       Both hooks are also exposed as attributes on the object, and you can
+       change them at any time.
 
     """
-    def __init__(self, receive_some_hook=None):
+    def __init__(self, receive_some_hook=None, close_hook=None):
         self._lock = _util.UnLock(
             _core.ResourceBusyError, "another task is using this stream")
         self._incoming = _UnboundedByteQueue()
         self._closed = False
         self.receive_some_hook = receive_some_hook
+        self.close_hook = close_hook
 
     async def receive_some(self, max_bytes):
         """Calls the :attr:`receive_some_hook` (if any), and then retrieves
@@ -1012,7 +1025,8 @@ def forceful_close(self):
         except _core.WouldBlock:
             pass
         self._incoming.close()
-        self._closed = True
+        if self.close_hook is not None:
+            self.close_hook()
 
     def put_data(self, data):
         """Appends the given data to the internal buffer.