From 6b4d4a25ceb1b3f016f9e4954bc630e6a0b7167f Mon Sep 17 00:00:00 2001
From: Eric Poulsen <eric@zyxod.com>
Date: Fri, 21 Jul 2017 10:30:35 -0700
Subject: [PATCH] extmod/modussl_mbedtls: Implement non-blocking SSL sockets.

---
 extmod/modussl_mbedtls.c | 25 +++++++++++++++++++------
 1 file changed, 19 insertions(+), 6 deletions(-)

diff --git a/extmod/modussl_mbedtls.c b/extmod/modussl_mbedtls.c
index 8c90c2cf46..597eaee73d 100644
--- a/extmod/modussl_mbedtls.c
+++ b/extmod/modussl_mbedtls.c
@@ -29,6 +29,7 @@
 
 #include <stdio.h>
 #include <string.h>
+#include <errno.h> // needed because mp_is_nonblocking_error uses system error codes
 
 #include "py/nlr.h"
 #include "py/runtime.h"
@@ -83,6 +84,9 @@ int _mbedtls_ssl_send(void *ctx, const byte *buf, size_t len) {
 
     int out_sz = sock_stream->write(sock, buf, len, &err);
     if (out_sz == MP_STREAM_ERROR) {
+        if (mp_is_nonblocking_error(err)) {
+            return MBEDTLS_ERR_SSL_WANT_WRITE;
+        }
         return -err;
     } else {
         return out_sz;
@@ -97,6 +101,9 @@ int _mbedtls_ssl_recv(void *ctx, byte *buf, size_t len) {
 
     int out_sz = sock_stream->read(sock, buf, len, &err);
     if (out_sz == MP_STREAM_ERROR) {
+        if (mp_is_nonblocking_error(err)) {
+            return MBEDTLS_ERR_SSL_WANT_READ;
+        }
         return -err;
     } else {
         return out_sz;
@@ -199,6 +206,9 @@ STATIC mp_uint_t socket_read(mp_obj_t o_in, void *buf, mp_uint_t size, int *errc
     if (ret >= 0) {
         return ret;
     }
+    if (ret == MBEDTLS_ERR_SSL_WANT_READ) {
+        ret = MP_EWOULDBLOCK;
+    }
     *errcode = ret;
     return MP_STREAM_ERROR;
 }
@@ -210,17 +220,20 @@ STATIC mp_uint_t socket_write(mp_obj_t o_in, const void *buf, mp_uint_t size, in
     if (ret >= 0) {
         return ret;
     }
+    if (ret == MBEDTLS_ERR_SSL_WANT_WRITE) {
+        ret = MP_EWOULDBLOCK;
+    }
     *errcode = ret;
     return MP_STREAM_ERROR;
 }
 
 STATIC mp_obj_t socket_setblocking(mp_obj_t self_in, mp_obj_t flag_in) {
-    // Currently supports only blocking mode
-    (void)self_in;
-    if (!mp_obj_is_true(flag_in)) {
-        mp_not_implemented("");
-    }
-    return mp_const_none;
+    mp_obj_ssl_socket_t *o = MP_OBJ_TO_PTR(self_in);
+    mp_obj_t sock = o->sock;
+    mp_obj_t dest[3];
+    mp_load_method(sock, MP_QSTR_setblocking, dest);
+    dest[2] = flag_in;
+    return mp_call_method_n_kw(1, 0, dest);
 }
 STATIC MP_DEFINE_CONST_FUN_OBJ_2(socket_setblocking_obj, socket_setblocking);