StackRduino
diff --git a/‎extmod/modwebsocket.c‎
Lines changed: 52 additions & 12 deletions b/‎extmod/modwebsocket.c‎
Lines changed: 52 additions & 12 deletions
@@ -37,7 +37,7 @@
 
 #if MICROPY_PY_WEBSOCKET
 
-enum { FRAME_HEADER, FRAME_OPT, PAYLOAD };
+enum { FRAME_HEADER, FRAME_OPT, PAYLOAD, CONTROL };
 
 enum { BLOCKING_WRITE = 0x80 };
 
@@ -52,10 +52,14 @@ typedef struct _mp_obj_websocket_t {
     byte buf_pos;
     byte buf[6];
     byte opts;
-    // Copy of current frame's flags
+    // Copy of last data frame flags
     byte ws_flags;
+    // Copy of current frame flags
+    byte last_flags;
 } mp_obj_websocket_t;
 
+STATIC mp_uint_t websocket_write(mp_obj_t self_in, const void *buf, mp_uint_t size, int *errcode);
+
 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, 2, false);
     mp_obj_websocket_t *o = m_new_obj(mp_obj_websocket_t);
@@ -97,10 +101,9 @@ STATIC mp_uint_t websocket_read(mp_obj_t self_in, void *buf, mp_uint_t size, int
                 // "Control frames MAY be injected in the middle of a fragmented message."
                 // So, they must be processed before data frames (and not alter
                 // self->ws_flags)
-                if ((self->buf[0] & FRAME_OPCODE_MASK) >= FRAME_CLOSE) {
-                    // TODO: implement
-                    assert(0);
-                }
+                byte frame_type = self->buf[0];
+                self->last_flags = frame_type;
+                frame_type &= FRAME_OPCODE_MASK;
 
                 if ((self->buf[0] & FRAME_OPCODE_MASK) == FRAME_CONT) {
                     // Preserve previous frame type
@@ -119,7 +122,7 @@ STATIC mp_uint_t websocket_read(mp_obj_t self_in, void *buf, mp_uint_t size, int
                     // Msg size is next 2 bytes
                     to_recv += 2;
                 } else if (sz == 127) {
-                    // Msg size is next 2 bytes
+                    // Msg size is next 8 bytes
                     assert(0);
                 }
                 if (self->buf[1] & 0x80) {
@@ -133,7 +136,11 @@ STATIC mp_uint_t websocket_read(mp_obj_t self_in, void *buf, mp_uint_t size, int
                 if (to_recv != 0) {
                     self->state = FRAME_OPT;
                 } else {
-                    self->state = PAYLOAD;
+                    if (frame_type >= FRAME_CLOSE) {
+                        self->state = CONTROL;
+                    } else {
+                        self->state = PAYLOAD;
+                    }
                 }
                 continue;
             }
@@ -148,13 +155,24 @@ STATIC mp_uint_t websocket_read(mp_obj_t self_in, void *buf, mp_uint_t size, int
                     memcpy(self->mask, self->buf + self->buf_pos - 4, 4);
                 }
                 self->buf_pos = 0;
-                self->state = PAYLOAD;
+                if ((self->last_flags & FRAME_OPCODE_MASK) >= FRAME_CLOSE) {
+                    self->state = CONTROL;
+                } else {
+                    self->state = PAYLOAD;
+                }
                 continue;
             }
 
-            case PAYLOAD: {
+            case PAYLOAD:
+            case CONTROL: {
+                mp_uint_t out_sz = 0;
+                if (self->msg_sz == 0) {
+                    // In case message had zero payload
+                    goto no_payload;
+                }
+
                 size_t sz = MIN(size, self->msg_sz);
-                mp_uint_t out_sz = stream_p->read(self->sock, buf, sz, errcode);
+                out_sz = stream_p->read(self->sock, buf, sz, errcode);
                 if (out_sz == 0 || out_sz == MP_STREAM_ERROR) {
                     return out_sz;
                 }
@@ -166,12 +184,34 @@ STATIC mp_uint_t websocket_read(mp_obj_t self_in, void *buf, mp_uint_t size, int
 
                 self->msg_sz -= out_sz;
                 if (self->msg_sz == 0) {
+                    byte last_state;
+no_payload:
+                    last_state = self->state;
                     self->state = FRAME_HEADER;
                     self->to_recv = 2;
                     self->mask_pos = 0;
                     self->buf_pos = 0;
+
+                    // Handle control frame
+                    if (last_state == CONTROL) {
+                        byte frame_type = self->last_flags & FRAME_OPCODE_MASK;
+                        if (frame_type == FRAME_CLOSE) {
+                            static char close_resp[2] = {0x88, 0};
+                            int err;
+                            websocket_write(self_in, close_resp, sizeof(close_resp), &err);
+                            return 0;
+                        }
+
+                        //DEBUG_printf("Finished receiving ctrl message %x, ignoring\n", self->last_flags);
+                        continue;
+                    }
                 }
-                return out_sz;
+
+                if (out_sz != 0) {
+                    return out_sz;
+                }
+                // Empty (data) frame received is not EOF
+                continue;
             }
 
         }