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:
parent
351ec6d4ab
commit
480c212009
@ -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;
|
||||
}
|
||||
|
||||
}
|
||||
|
Loading…
Reference in New Issue
Block a user