From 5b1c2217856503f29bf455a44ca23e32c43a2008 Mon Sep 17 00:00:00 2001
From: Paul Sokolovsky <pfalcon@users.sourceforge.net>
Date: Sat, 9 Apr 2016 16:03:38 +0300
Subject: [PATCH] extmod/modwebsocket: Add option for blocking writes to
 non-blk sockets.

This is strange asymmetry which is sometimes needed, e.g. for WebREPL: we
want to process only available input and no more; but for output, we want
to get rid of all of it, because there's no other place to buffer/store
it. This asymmetry is akin to CPython's asyncio asymmetry, where reads are
asynchronous, but writes are synchronous (asyncio doesn't expect them to
block, instead expects there to be (unlimited) buffering for any sync write
to completely immediately).
---
 extmod/modwebsocket.c | 29 ++++++++++++++++++++++++-----
 1 file changed, 24 insertions(+), 5 deletions(-)

diff --git a/extmod/modwebsocket.c b/extmod/modwebsocket.c
index 6cd4f515d..9b0c19d6d 100644
--- a/extmod/modwebsocket.c
+++ b/extmod/modwebsocket.c
@@ -37,6 +37,7 @@
 #if MICROPY_PY_WEBSOCKET
 
 enum { FRAME_HEADER, FRAME_OPT, PAYLOAD };
+enum { BLOCKING_WRITE = 1 };
 
 typedef struct _mp_obj_websocket_t {
     mp_obj_base_t base;
@@ -48,10 +49,11 @@ typedef struct _mp_obj_websocket_t {
     byte mask_pos;
     byte buf_pos;
     byte buf[6];
+    byte opts;
 } mp_obj_websocket_t;
 
 STATIC mp_obj_t websocket_make_new(const mp_obj_type_t *type, size_t n_args, size_t n_kw, const mp_obj_t *args) {
-    mp_arg_check_num(n_args, n_kw, 1, 1, false);
+    mp_arg_check_num(n_args, n_kw, 1, 2, false);
     mp_obj_websocket_t *o = m_new_obj(mp_obj_websocket_t);
     o->base.type = type;
     o->sock = args[0];
@@ -59,6 +61,10 @@ STATIC mp_obj_t websocket_make_new(const mp_obj_type_t *type, size_t n_args, siz
     o->to_recv = 2;
     o->mask_pos = 0;
     o->buf_pos = 0;
+    o->opts = 0;
+    if (n_args > 1 && args[1] == mp_const_true) {
+        o->opts |= BLOCKING_WRITE;
+    }
     return o;
 }
 
@@ -157,11 +163,24 @@ STATIC mp_uint_t websocket_write(mp_obj_t self_in, const void *buf, mp_uint_t si
     assert(size < 126);
     byte header[] = {0x81, size};
 
-    mp_uint_t out_sz = mp_stream_writeall(self->sock, header, sizeof(header), errcode);
-    if (out_sz == MP_STREAM_ERROR) {
-        return MP_STREAM_ERROR;
+    mp_obj_t dest[3];
+    if (self->opts & BLOCKING_WRITE) {
+        mp_load_method(self->sock, MP_QSTR_setblocking, dest);
+        dest[2] = mp_const_true;
+        mp_call_method_n_kw(1, 0, dest);
     }
-    return mp_stream_writeall(self->sock, buf, size, errcode);
+
+    mp_uint_t out_sz = mp_stream_writeall(self->sock, header, sizeof(header), errcode);
+    if (out_sz != MP_STREAM_ERROR) {
+        out_sz = mp_stream_writeall(self->sock, buf, size, errcode);
+    }
+
+    if (self->opts & BLOCKING_WRITE) {
+        dest[2] = mp_const_false;
+        mp_call_method_n_kw(1, 0, dest);
+    }
+
+    return out_sz;
 }
 
 STATIC const mp_map_elem_t websocket_locals_dict_table[] = {