take micropython tricks for selectability of ssl sockets

Author: jepler

shared-module/ssl/SSLSocket.c

@@ -22,6 +22,8 @@
 
 #include "mbedtls/version.h"
 
+#define MP_STREAM_POLL_RDWR (MP_STREAM_POLL_RD | MP_STREAM_POLL_WR)
+
 #if defined(MBEDTLS_ERROR_C)
 #include "../../lib/mbedtls_errors/mp_mbedtls_errors.c"
 #endif
@@ -220,6 +222,7 @@ ssl_sslsocket_obj_t *common_hal_ssl_sslcontext_wrap_socket(ssl_sslcontext_obj_t
     o->base.type = &ssl_sslsocket_type;
     o->ssl_context = self;
     o->sock_obj = socket;
+    o->poll_mask = 0;
 
     mp_load_method(socket, MP_QSTR_accept, o->accept_args);
     mp_load_method(socket, MP_QSTR_bind, o->bind_args);
@@ -331,6 +334,7 @@ ssl_sslsocket_obj_t *common_hal_ssl_sslcontext_wrap_socket(ssl_sslcontext_obj_t
 }
 
 mp_uint_t common_hal_ssl_sslsocket_recv_into(ssl_sslsocket_obj_t *self, uint8_t *buf, uint32_t len) {
+    self->poll_mask = 0;
     int ret = mbedtls_ssl_read(&self->ssl, buf, len);
     DEBUG_PRINT("recv_into mbedtls_ssl_read() -> %d\n", ret);
     if (ret == MBEDTLS_ERR_SSL_PEER_CLOSE_NOTIFY) {
@@ -342,17 +346,24 @@ mp_uint_t common_hal_ssl_sslsocket_recv_into(ssl_sslsocket_obj_t *self, uint8_t
         DEBUG_PRINT("returning %d\n", ret);
         return ret;
     }
+    if (ret == MBEDTLS_ERR_SSL_WANT_WRITE) {
+        self->poll_mask = MP_STREAM_POLL_WR;
+    }
     DEBUG_PRINT("raising errno [error case] %d\n", ret);
     mbedtls_raise_error(ret);
 }
 
 mp_uint_t common_hal_ssl_sslsocket_send(ssl_sslsocket_obj_t *self, const uint8_t *buf, uint32_t len) {
+    self->poll_mask = 0;
     int ret = mbedtls_ssl_write(&self->ssl, buf, len);
     DEBUG_PRINT("send mbedtls_ssl_write() -> %d\n", ret);
     if (ret >= 0) {
         DEBUG_PRINT("returning %d\n", ret);
         return ret;
     }
+    if (ret == MBEDTLS_ERR_SSL_WANT_READ) {
+        self->poll_mask = MP_STREAM_POLL_RD;
+    }
     DEBUG_PRINT("raising errno [error case] %d\n", ret);
     mbedtls_raise_error(ret);
 }
@@ -449,10 +460,29 @@ void common_hal_ssl_sslsocket_settimeout(ssl_sslsocket_obj_t *self, mp_obj_t tim
     ssl_socket_settimeout(self, timeout_obj);
 }
 
-static bool poll_common(ssl_sslsocket_obj_t *self, int mode) {
+static bool poll_common(ssl_sslsocket_obj_t *self, uintptr_t arg) {
+    // Take into account that the library might have buffered data already
+    int has_pending = 0;
+    if (arg & MP_STREAM_POLL_RD) {
+        has_pending = mbedtls_ssl_check_pending(&self->ssl);
+        if (has_pending) {
+            // Shortcut if we only need to read and we have buffered data, no need to go to the underlying socket
+            return true;
+        }
+    }
+
+    // If the library signaled us that it needs reading or writing, only
+    // check that direction
+    if (self->poll_mask && (arg & MP_STREAM_POLL_RDWR)) {
+        arg = (arg & ~MP_STREAM_POLL_RDWR) | self->poll_mask;
+    }
+
+    // If direction the library needed is available, return a fake
+    // result to the caller so that it reenters a read or a write to
+    // allow the handshake to progress
     const mp_stream_p_t *stream_p = mp_get_stream_raise(self->sock_obj, MP_STREAM_OP_IOCTL);
     int errcode;
-    mp_int_t ret = stream_p->ioctl(self->sock_obj, MP_STREAM_POLL, mode, &errcode);
+    mp_int_t ret = stream_p->ioctl(self->sock_obj, MP_STREAM_POLL, arg, &errcode);
     return ret != 0;
 }
 

shared-module/ssl/SSLSocket.h

@@ -29,6 +29,7 @@ typedef struct ssl_sslsocket_obj {
     mbedtls_x509_crt cacert;
     mbedtls_x509_crt cert;
     mbedtls_pk_context pkey;
+    uintptr_t poll_mask;
     bool closed;
     mp_obj_t accept_args[2];
     mp_obj_t bind_args[3];