void ws_broadcast(ws_s* ws, char* data, size_t size, uint8_t is_text) { // Copy the message to a broadcast data-packet struct ws_data* msg = malloc(sizeof(*msg) + size); msg->size = size; memcpy(msg->data, data, size); // Asynchronously calls `ws_get_broadcast` for each of the websockets // (except this one) // and calls `free_wsdata` once all the broadcasts were perfomed. websocket_each(ws, ws_get_broadcast, msg, free_wsdata); // echos the data to the current websocket websocket_write(ws, data, size, 1); }
STATIC mp_uint_t websocket_read(mp_obj_t self_in, void *buf, mp_uint_t size, int *errcode) { mp_obj_websocket_t *self = self_in; const mp_stream_p_t *stream_p = mp_get_stream_raise(self->sock, MP_STREAM_OP_READ); while (1) { if (self->to_recv != 0) { mp_uint_t out_sz = stream_p->read(self->sock, self->buf + self->buf_pos, self->to_recv, errcode); if (out_sz == 0 || out_sz == MP_STREAM_ERROR) { return out_sz; } self->buf_pos += out_sz; self->to_recv -= out_sz; if (self->to_recv != 0) { *errcode = EAGAIN; return MP_STREAM_ERROR; } } switch (self->state) { case FRAME_HEADER: { // TODO: Split frame handling below is untested so far, so conservatively disable it assert(self->buf[0] & 0x80); // "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) 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 self->ws_flags = (self->ws_flags & FRAME_OPCODE_MASK) | (self->buf[0] & ~FRAME_OPCODE_MASK); } else { self->ws_flags = self->buf[0]; } // Reset mask in case someone will use "simplified" protocol // without masks. memset(self->mask, 0, sizeof(self->mask)); int to_recv = 0; size_t sz = self->buf[1] & 0x7f; if (sz == 126) { // Msg size is next 2 bytes to_recv += 2; } else if (sz == 127) { // Msg size is next 8 bytes assert(0); } if (self->buf[1] & 0x80) { // Next 4 bytes is mask to_recv += 4; } self->buf_pos = 0; self->to_recv = to_recv; self->msg_sz = sz; // May be overriden by FRAME_OPT if (to_recv != 0) { self->state = FRAME_OPT; } else { if (frame_type >= FRAME_CLOSE) { self->state = CONTROL; } else { self->state = PAYLOAD; } } continue; } case FRAME_OPT: { if ((self->buf_pos & 3) == 2) { // First two bytes are message length self->msg_sz = (self->buf[0] << 8) | self->buf[1]; } if (self->buf_pos >= 4) { // Last 4 bytes is mask memcpy(self->mask, self->buf + self->buf_pos - 4, 4); } self->buf_pos = 0; if ((self->last_flags & FRAME_OPCODE_MASK) >= FRAME_CLOSE) { self->state = CONTROL; } else { self->state = PAYLOAD; } continue; } 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); out_sz = stream_p->read(self->sock, buf, sz, errcode); if (out_sz == 0 || out_sz == MP_STREAM_ERROR) { return out_sz; } sz = out_sz; for (byte *p = buf; sz--; p++) { *p ^= self->mask[self->mask_pos++ & 3]; } 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; } } if (out_sz != 0) { return out_sz; } // Empty (data) frame received is not EOF continue; } } } }
void ws_shutdown(ws_s* ws) { websocket_write(ws, "Shutting Down", 13, 1); }
/* the broadcast "task" performed by `Websocket.each` */ void ws_get_broadcast(ws_s* ws, void* arg) { struct ws_data* data = arg; websocket_write(ws, data->data, data->size, 1); // echo }
void ws_echo(ws_s* ws, char* data, size_t size, uint8_t is_text) { // echos the data to the current websocket websocket_write(ws, data, size, 1); }