diff options
Diffstat (limited to 'extmod/modwebsocket.c')
-rw-r--r-- | extmod/modwebsocket.c | 18 |
1 files changed, 15 insertions, 3 deletions
diff --git a/extmod/modwebsocket.c b/extmod/modwebsocket.c index 344933ded3..fc5e29a05e 100644 --- a/extmod/modwebsocket.c +++ b/extmod/modwebsocket.c @@ -240,9 +240,9 @@ STATIC mp_uint_t websocket_write(mp_obj_t self_in, const void *buf, mp_uint_t si mp_call_method_n_kw(1, 0, dest); } - mp_uint_t out_sz = mp_stream_writeall(self->sock, header, hdr_sz, errcode); - if (out_sz != MP_STREAM_ERROR) { - out_sz = mp_stream_writeall(self->sock, buf, size, errcode); + mp_uint_t out_sz = mp_stream_write_exactly(self->sock, header, hdr_sz, errcode); + if (*errcode == 0) { + out_sz = mp_stream_write_exactly(self->sock, buf, size, errcode); } if (self->opts & BLOCKING_WRITE) { @@ -250,6 +250,9 @@ STATIC mp_uint_t websocket_write(mp_obj_t self_in, const void *buf, mp_uint_t si mp_call_method_n_kw(1, 0, dest); } + if (*errcode != 0) { + return MP_STREAM_ERROR; + } return out_sz; } @@ -269,10 +272,19 @@ STATIC mp_uint_t websocket_ioctl(mp_obj_t self_in, mp_uint_t request, uintptr_t } } +STATIC mp_obj_t websocket_close(mp_obj_t self_in) { + mp_obj_websocket_t *self = MP_OBJ_TO_PTR(self_in); + // TODO: Send close signaling to the other side, otherwise it's + // abrupt close (connection abort). + return mp_stream_close(self->sock); +} +STATIC MP_DEFINE_CONST_FUN_OBJ_1(websocket_close_obj, websocket_close); + STATIC const mp_map_elem_t websocket_locals_dict_table[] = { { MP_OBJ_NEW_QSTR(MP_QSTR_read), (mp_obj_t)&mp_stream_read_obj }, { MP_OBJ_NEW_QSTR(MP_QSTR_write), (mp_obj_t)&mp_stream_write_obj }, { MP_OBJ_NEW_QSTR(MP_QSTR_ioctl), (mp_obj_t)&mp_stream_ioctl_obj }, + { MP_OBJ_NEW_QSTR(MP_QSTR_close), (mp_obj_t)&websocket_close_obj }, }; STATIC MP_DEFINE_CONST_DICT(websocket_locals_dict, websocket_locals_dict_table); |