extmod/modwebsocket: Handle CLOSE control frame.

This fixes situation when clients hangs waiting for disconnect and does
so only on timeout.
This commit is contained in:
Paul Sokolovsky 2016-04-27 12:49:30 +03:00
parent 351ec6d4ab
commit 480c212009

View File

@ -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;
}
}