tls: Fix WantReadError/WantWriteError signaling

This lets us further simplify TLSWrappedSocket.do_handshake().

Author: Synss

src/mbedtls/_tls.pyx

@@ -84,7 +84,7 @@ cdef int buffer_read(void *ctx, unsigned char *buf, const size_t len) nogil:
     """Read from input buffer."""
     c_buf = <_tls._C_Buffers *> ctx
     if _rb.c_len(c_buf.in_ctx) == 0:
-        return _tls.MBEDTLS_ERR_SSL_WANT_WRITE
+        return _tls.MBEDTLS_ERR_SSL_WANT_READ
     return _rb.c_readinto(c_buf.in_ctx, buf, len)
 
 
@@ -1247,7 +1247,6 @@ cdef class MbedTLSBuffer:
             return 0
         if amt <= 0:
             return 0
-        # cdef size_t avail = _tls.mbedtls_ssl_get_bytes_avail(&self._ctx)
         read = _tls.mbedtls_ssl_read(&self._ctx, &buffer[0], amt)
         if read > 0:
             return read
@@ -1342,26 +1341,28 @@ cdef class MbedTLSBuffer:
     def do_handshake(self):
         if self._handshake_state is HandshakeStep.HANDSHAKE_OVER:
             raise ValueError("handshake already over")
-        self._handle_handshake_response(_tls.mbedtls_ssl_handshake_step(&self._ctx))
+        self._handle_handshake_response(
+            _tls.mbedtls_ssl_handshake_step(&self._ctx)
+        )
 
     def _renegotiate(self):
         """Initialize an SSL renegotiation on the running connection."""
         self._handle_handshake_response(_tls.mbedtls_ssl_renegotiate(&self._ctx))
 
     def _handle_handshake_response(self, ret):
-        if ret == 0:
-            return
-        elif ret == _tls.MBEDTLS_ERR_SSL_WANT_READ:
+        if ret == _tls.MBEDTLS_ERR_SSL_WANT_READ:
             raise WantReadError()
-        elif ret == _tls.MBEDTLS_ERR_SSL_WANT_WRITE:
+        if ret == _tls.MBEDTLS_ERR_SSL_WANT_WRITE:
             raise WantWriteError()
-        elif ret == _tls.MBEDTLS_ERR_SSL_HELLO_VERIFY_REQUIRED:
+        if ret == _tls.MBEDTLS_ERR_SSL_HELLO_VERIFY_REQUIRED:
             self._reset()
             raise HelloVerifyRequest()
-        else:
-            assert ret < 0
+        if ret < 0:
             self._reset()
             _exc.check_error(ret)
+        if ret == 0 and self._output_buffer:
+            raise WantWriteError
+        assert ret == 0
 
     def _get_channel_binding(self, cb_type="tls-unique"):
         return None

src/mbedtls/tls.py

@@ -186,6 +186,10 @@ def __exit__(self, *exc_info):
     def __str__(self):
         return str(self._socket)
 
+    @property
+    def _handshake_state(self):
+        return self._buffer._handshake_state
+
     # PEP 543 requires the full socket API.
 
     @property
@@ -333,19 +337,16 @@ def shutdown(self, how):
     # PEP 543 adds the following methods.
 
     def do_handshake(self):
-        while (
-            self._buffer._handshake_state is not HandshakeStep.HANDSHAKE_OVER
-        ):
+        while self._handshake_state is not HandshakeStep.HANDSHAKE_OVER:
             try:
                 self._buffer.do_handshake()
-                amt = self._socket.send(self._buffer.peek_outgoing(1024))
-                self._buffer.consume_outgoing(amt)
             except WantReadError:
-                amt = self._socket.send(self._buffer.peek_outgoing(1024))
-                self._buffer.consume_outgoing(amt)
-            except WantWriteError:
                 data = self._socket.recv(1024)
                 self._buffer.receive_from_network(data)
+            except WantWriteError:
+                in_transit = self._buffer.peek_outgoing(1024)
+                amt = self._socket.send(in_transit)
+                self._buffer.consume_outgoing(amt)
 
     def setcookieparam(self, param):
         self._buffer.setcookieparam(param)

tests/test_tls.py

@@ -1,6 +1,7 @@
 import datetime as dt
 import pickle
 import socket
+from contextlib import suppress
 
 import pytest
 
@@ -515,7 +516,8 @@ def do_io(*, src, dst, amt=1024):
 
         def do_handshake(*end_state_pair):
             for end, state in end_state_pair:
-                end.do_handshake()
+                with suppress(WantReadError, WantWriteError):
+                    end.do_handshake()
                 assert end._handshake_state is state
 
         do_handshake(