Branch data Line data Source code
1 : : /*
2 : : * This file is part of the MicroPython project, http://micropython.org/
3 : : *
4 : : * The MIT License (MIT)
5 : : *
6 : : * Copyright (c) 2016 Paul Sokolovsky
7 : : *
8 : : * Permission is hereby granted, free of charge, to any person obtaining a copy
9 : : * of this software and associated documentation files (the "Software"), to deal
10 : : * in the Software without restriction, including without limitation the rights
11 : : * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
12 : : * copies of the Software, and to permit persons to whom the Software is
13 : : * furnished to do so, subject to the following conditions:
14 : : *
15 : : * The above copyright notice and this permission notice shall be included in
16 : : * all copies or substantial portions of the Software.
17 : : *
18 : : * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
19 : : * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
20 : : * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
21 : : * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
22 : : * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
23 : : * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
24 : : * THE SOFTWARE.
25 : : */
26 : :
27 : : #include <stdio.h>
28 : : #include <stdint.h>
29 : : #include <string.h>
30 : :
31 : : #include "py/runtime.h"
32 : : #include "py/stream.h"
33 : : #include "extmod/modwebsocket.h"
34 : :
35 : : #if MICROPY_PY_WEBSOCKET
36 : :
37 : : enum { FRAME_HEADER, FRAME_OPT, PAYLOAD, CONTROL };
38 : :
39 : : enum { BLOCKING_WRITE = 0x80 };
40 : :
41 : : typedef struct _mp_obj_websocket_t {
42 : : mp_obj_base_t base;
43 : : mp_obj_t sock;
44 : : uint32_t msg_sz;
45 : : byte mask[4];
46 : : byte state;
47 : : byte to_recv;
48 : : byte mask_pos;
49 : : byte buf_pos;
50 : : byte buf[6];
51 : : byte opts;
52 : : // Copy of last data frame flags
53 : : byte ws_flags;
54 : : // Copy of current frame flags
55 : : byte last_flags;
56 : : } mp_obj_websocket_t;
57 : :
58 : : static mp_uint_t websocket_write(mp_obj_t self_in, const void *buf, mp_uint_t size, int *errcode);
59 : : static mp_uint_t websocket_write_raw(mp_obj_t self_in, const byte *header, int hdr_sz, const void *buf, mp_uint_t size, int *errcode);
60 : :
61 : 22 : 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) {
62 : 22 : mp_arg_check_num(n_args, n_kw, 1, 2, false);
63 : 22 : mp_get_stream_raise(args[0], MP_STREAM_OP_READ | MP_STREAM_OP_WRITE | MP_STREAM_OP_IOCTL);
64 : 22 : mp_obj_websocket_t *o = mp_obj_malloc(mp_obj_websocket_t, type);
65 : 22 : o->sock = args[0];
66 : 22 : o->state = FRAME_HEADER;
67 : 22 : o->to_recv = 2;
68 : 22 : o->mask_pos = 0;
69 : 22 : o->buf_pos = 0;
70 : 22 : o->opts = FRAME_TXT;
71 [ - + - - ]: 22 : if (n_args > 1 && args[1] == mp_const_true) {
72 : 0 : o->opts |= BLOCKING_WRITE;
73 : : }
74 : 22 : return MP_OBJ_FROM_PTR(o);
75 : : }
76 : :
77 : 14 : static mp_uint_t websocket_read(mp_obj_t self_in, void *buf, mp_uint_t size, int *errcode) {
78 : 14 : mp_obj_websocket_t *self = MP_OBJ_TO_PTR(self_in);
79 : 14 : const mp_stream_p_t *stream_p = mp_get_stream(self->sock);
80 : 40 : while (1) {
81 [ + + ]: 40 : if (self->to_recv != 0) {
82 : 22 : mp_uint_t out_sz = stream_p->read(self->sock, self->buf + self->buf_pos, self->to_recv, errcode);
83 [ + - ]: 22 : if (out_sz == 0 || out_sz == MP_STREAM_ERROR) {
84 : : return out_sz;
85 : : }
86 : 22 : self->buf_pos += out_sz;
87 : 22 : self->to_recv -= out_sz;
88 [ - + ]: 22 : if (self->to_recv != 0) {
89 : 0 : *errcode = MP_EAGAIN;
90 : 0 : return MP_STREAM_ERROR;
91 : : }
92 : : }
93 : :
94 [ - + + + ]: 40 : switch (self->state) {
95 : 36 : case FRAME_HEADER: {
96 : : // TODO: Split frame handling below is untested so far, so conservatively disable it
97 [ - + ]: 18 : assert(self->buf[0] & 0x80);
98 : :
99 : : // "Control frames MAY be injected in the middle of a fragmented message."
100 : : // So, they must be processed before data frames (and not alter
101 : : // self->ws_flags)
102 : 18 : byte frame_type = self->buf[0];
103 : 18 : self->last_flags = frame_type;
104 : 18 : frame_type &= FRAME_OPCODE_MASK;
105 : :
106 [ + + ]: 18 : if ((self->buf[0] & FRAME_OPCODE_MASK) == FRAME_CONT) {
107 : : // Preserve previous frame type
108 : 2 : self->ws_flags = (self->ws_flags & FRAME_OPCODE_MASK) | (self->buf[0] & ~FRAME_OPCODE_MASK);
109 : : } else {
110 : 16 : self->ws_flags = self->buf[0];
111 : : }
112 : :
113 : : // Reset mask in case someone will use "simplified" protocol
114 : : // without masks.
115 : 18 : memset(self->mask, 0, sizeof(self->mask));
116 : :
117 : 18 : int to_recv = 0;
118 : 18 : size_t sz = self->buf[1] & 0x7f;
119 [ + + ]: 18 : if (sz == 126) {
120 : : // Msg size is next 2 bytes
121 : : to_recv += 2;
122 [ - + ]: 16 : } else if (sz == 127) {
123 : : // Msg size is next 8 bytes
124 : 0 : assert(0);
125 : : }
126 [ + + ]: 18 : if (self->buf[1] & 0x80) {
127 : : // Next 4 bytes is mask
128 : 2 : to_recv += 4;
129 : : }
130 : :
131 : 18 : self->buf_pos = 0;
132 : 18 : self->to_recv = to_recv;
133 : 18 : self->msg_sz = sz; // May be overridden by FRAME_OPT
134 [ + + ]: 18 : if (to_recv != 0) {
135 : 4 : self->state = FRAME_OPT;
136 : : } else {
137 [ + + ]: 14 : if (frame_type >= FRAME_CLOSE) {
138 : 6 : self->state = CONTROL;
139 : : } else {
140 : 8 : self->state = PAYLOAD;
141 : : }
142 : : }
143 : 18 : continue;
144 : : }
145 : :
146 : 4 : case FRAME_OPT: {
147 [ + + ]: 4 : if ((self->buf_pos & 3) == 2) {
148 : : // First two bytes are message length
149 : 2 : self->msg_sz = (self->buf[0] << 8) | self->buf[1];
150 : : }
151 [ + + ]: 4 : if (self->buf_pos >= 4) {
152 : : // Last 4 bytes is mask
153 : 2 : memcpy(self->mask, self->buf + self->buf_pos - 4, 4);
154 : : }
155 : 4 : self->buf_pos = 0;
156 [ - + ]: 4 : if ((self->last_flags & FRAME_OPCODE_MASK) >= FRAME_CLOSE) {
157 : 0 : self->state = CONTROL;
158 : : } else {
159 : 4 : self->state = PAYLOAD;
160 : : }
161 : 4 : continue;
162 : : }
163 : :
164 : 18 : case PAYLOAD:
165 : 0 : case CONTROL: {
166 : 18 : mp_uint_t out_sz = 0;
167 [ + + ]: 18 : if (self->msg_sz == 0) {
168 : : // In case message had zero payload
169 : 6 : goto no_payload;
170 : : }
171 : :
172 : 12 : size_t sz = MIN(size, self->msg_sz);
173 : 12 : out_sz = stream_p->read(self->sock, buf, sz, errcode);
174 [ + - ]: 12 : if (out_sz == 0 || out_sz == MP_STREAM_ERROR) {
175 : : return out_sz;
176 : : }
177 : :
178 : : sz = out_sz;
179 [ + + ]: 308 : for (byte *p = buf; sz--; p++) {
180 : 296 : *p ^= self->mask[self->mask_pos++ & 3];
181 : : }
182 : :
183 : 12 : self->msg_sz -= out_sz;
184 [ + - ]: 12 : if (self->msg_sz == 0) {
185 : 18 : byte last_state;
186 : 12 : no_payload:
187 : 18 : last_state = self->state;
188 : 18 : self->state = FRAME_HEADER;
189 : 18 : self->to_recv = 2;
190 : 18 : self->mask_pos = 0;
191 : 18 : self->buf_pos = 0;
192 : :
193 : : // Handle control frame
194 [ + + ]: 18 : if (last_state == CONTROL) {
195 : 6 : byte frame_type = self->last_flags & FRAME_OPCODE_MASK;
196 [ + + ]: 6 : if (frame_type == FRAME_CLOSE) {
197 : 2 : static const byte close_resp[2] = {0x88, 0};
198 : 2 : int err;
199 : 2 : websocket_write_raw(self_in, close_resp, sizeof(close_resp), close_resp, 0, &err);
200 : 2 : return 0;
201 : : }
202 : :
203 : : // DEBUG_printf("Finished receiving ctrl message %x, ignoring\n", self->last_flags);
204 : 4 : continue;
205 : : }
206 : : }
207 : :
208 [ - + ]: 12 : if (out_sz != 0) {
209 : : return out_sz;
210 : : }
211 : : // Empty (data) frame received is not EOF
212 : 0 : continue;
213 : : }
214 : :
215 : : }
216 : : }
217 : : }
218 : :
219 : 4 : static mp_uint_t websocket_write(mp_obj_t self_in, const void *buf, mp_uint_t size, int *errcode) {
220 : 4 : mp_obj_websocket_t *self = MP_OBJ_TO_PTR(self_in);
221 [ - + ]: 4 : assert(size < 0x10000);
222 : 4 : byte header[4] = {0x80 | (self->opts & FRAME_OPCODE_MASK)};
223 : 4 : int hdr_sz;
224 [ + + ]: 4 : if (size < 126) {
225 : 2 : header[1] = size;
226 : 2 : hdr_sz = 2;
227 : : } else {
228 : 2 : header[1] = 126;
229 : 2 : header[2] = size >> 8;
230 : 2 : header[3] = size & 0xff;
231 : 2 : hdr_sz = 4;
232 : : }
233 : :
234 : 4 : return websocket_write_raw(self_in, header, hdr_sz, buf, size, errcode);
235 : : }
236 : 6 : static mp_uint_t websocket_write_raw(mp_obj_t self_in, const byte *header, int hdr_sz, const void *buf, mp_uint_t size, int *errcode) {
237 : 6 : mp_obj_websocket_t *self = MP_OBJ_TO_PTR(self_in);
238 : :
239 : 6 : mp_obj_t dest[3];
240 [ - + ]: 6 : if (self->opts & BLOCKING_WRITE) {
241 : 0 : mp_load_method(self->sock, MP_QSTR_setblocking, dest);
242 : 0 : dest[2] = mp_const_true;
243 : 0 : mp_call_method_n_kw(1, 0, dest);
244 : : }
245 : :
246 : 6 : mp_uint_t out_sz = mp_stream_write_exactly(self->sock, header, hdr_sz, errcode);
247 [ + - ]: 6 : if (*errcode == 0) {
248 : 6 : out_sz = mp_stream_write_exactly(self->sock, buf, size, errcode);
249 : : }
250 : :
251 [ - + ]: 6 : if (self->opts & BLOCKING_WRITE) {
252 : 0 : dest[2] = mp_const_false;
253 : 0 : mp_call_method_n_kw(1, 0, dest);
254 : : }
255 : :
256 [ - + ]: 6 : if (*errcode != 0) {
257 : 0 : return MP_STREAM_ERROR;
258 : : }
259 : : return out_sz;
260 : : }
261 : :
262 : 10 : static mp_uint_t websocket_ioctl(mp_obj_t self_in, mp_uint_t request, uintptr_t arg, int *errcode) {
263 : 10 : mp_obj_websocket_t *self = MP_OBJ_TO_PTR(self_in);
264 [ + + + + ]: 10 : switch (request) {
265 : 2 : case MP_STREAM_CLOSE:
266 : : // TODO: Send close signaling to the other side, otherwise it's
267 : : // abrupt close (connection abort).
268 : 2 : mp_stream_close(self->sock);
269 : 2 : return 0;
270 : 2 : case MP_STREAM_GET_DATA_OPTS:
271 : 2 : return self->ws_flags & FRAME_OPCODE_MASK;
272 : 4 : case MP_STREAM_SET_DATA_OPTS: {
273 : 4 : int cur = self->opts & FRAME_OPCODE_MASK;
274 : 4 : self->opts = (self->opts & ~FRAME_OPCODE_MASK) | (arg & FRAME_OPCODE_MASK);
275 : 4 : return cur;
276 : : }
277 : 2 : default:
278 : 2 : *errcode = MP_EINVAL;
279 : 2 : return MP_STREAM_ERROR;
280 : : }
281 : : }
282 : :
283 : : static const mp_rom_map_elem_t websocket_locals_dict_table[] = {
284 : : { MP_ROM_QSTR(MP_QSTR_read), MP_ROM_PTR(&mp_stream_read_obj) },
285 : : { MP_ROM_QSTR(MP_QSTR_readinto), MP_ROM_PTR(&mp_stream_readinto_obj) },
286 : : { MP_ROM_QSTR(MP_QSTR_readline), MP_ROM_PTR(&mp_stream_unbuffered_readline_obj) },
287 : : { MP_ROM_QSTR(MP_QSTR_write), MP_ROM_PTR(&mp_stream_write_obj) },
288 : : { MP_ROM_QSTR(MP_QSTR_ioctl), MP_ROM_PTR(&mp_stream_ioctl_obj) },
289 : : { MP_ROM_QSTR(MP_QSTR_close), MP_ROM_PTR(&mp_stream_close_obj) },
290 : : };
291 : : static MP_DEFINE_CONST_DICT(websocket_locals_dict, websocket_locals_dict_table);
292 : :
293 : : static const mp_stream_p_t websocket_stream_p = {
294 : : .read = websocket_read,
295 : : .write = websocket_write,
296 : : .ioctl = websocket_ioctl,
297 : : };
298 : :
299 : : static MP_DEFINE_CONST_OBJ_TYPE(
300 : : websocket_type,
301 : : MP_QSTR_websocket,
302 : : MP_TYPE_FLAG_NONE,
303 : : make_new, websocket_make_new,
304 : : protocol, &websocket_stream_p,
305 : : locals_dict, &websocket_locals_dict
306 : : );
307 : :
308 : : static const mp_rom_map_elem_t websocket_module_globals_table[] = {
309 : : { MP_ROM_QSTR(MP_QSTR___name__), MP_ROM_QSTR(MP_QSTR_websocket) },
310 : : { MP_ROM_QSTR(MP_QSTR_websocket), MP_ROM_PTR(&websocket_type) },
311 : : };
312 : :
313 : : static MP_DEFINE_CONST_DICT(websocket_module_globals, websocket_module_globals_table);
314 : :
315 : : const mp_obj_module_t mp_module_websocket = {
316 : : .base = { &mp_type_module },
317 : : .globals = (mp_obj_dict_t *)&websocket_module_globals,
318 : : };
319 : :
320 : : // This module should not be extensible (as it is not a CPython standard
321 : : // library nor is it necessary to override from the filesystem), however it
322 : : // has previously been known as `uwebsocket`, so by making it extensible the
323 : : // `uwebsocket` alias will continue to work.
324 : : MP_REGISTER_EXTENSIBLE_MODULE(MP_QSTR_websocket, mp_module_websocket);
325 : :
326 : : #endif // MICROPY_PY_WEBSOCKET
|
