8000 extmod/modwebsocket: Add client support. · micropython/micropython@caec282 · GitHub
[go: up one dir, main page]
Skip to content

Commit caec282

Browse files
ccooper21ccooper21
committed
extmod/modwebsocket: Add client support.
This change adds websocket client support as per RFC6455. Section 5.1 states that a websocket client must mask all frames that it sends to the server. Further, section 5.3 states the masking key is a 32-bit value chosen at random by the client unique to each frame. These requirements are now implemented. A non-compliant debug mask is also supported, allowing control of the masking key.
1 parent 1742ab2 commit caec282

File tree

3 files changed

+91
-15
lines changed

3 files changed

+91
-15
lines changed

‎extmod/modwebsocket.c

Lines changed: 76 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@
2828
#include <stdint.h>
2929
#include <string.h>
3030

31+
#include "py/objmodule.h"
3132
#include "py/runtime.h"
3233
#include "py/stream.h"
3334
#include "extmod/modwebsocket.h"
@@ -38,14 +39,18 @@ enum { FRAME_HEADER, FRAME_OPT, PAYLOAD, CONTROL };
3839

3940
enum { BLOCKING_WRITE = 0x80 };
4041

42+
enum { NO_WRITE_MASKING, NORMAL_WRITE_MASKING, DEBUG_WRITE_MASKING };
43+
4144
typedef struct _mp_obj_websocket_t {
4245
mp_obj_base_t base;
4346
mp_obj_t sock;
4447
uint32_t msg_sz;
45-
byte mask[4];
48+
byte read_mask[4];
49+
byte do_write_masking;
50+
byte debug_write_mask[4];
4651
byte state;
4752
byte to_recv;
48-
byte mask_pos;
53+
byte read_mask_pos;
4954
byte buf_pos;
5055
byte buf[6];
5156
byte opts;
@@ -58,16 +63,39 @@ typedef struct _mp_obj_websocket_t {
5863
STATIC mp_uint_t websocket_write(mp_obj_t self_in, const void *buf, mp_uint_t size, int *errcode);
5964

6065
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) {
61-
mp_arg_check_num(n_args, n_kw, 1, 2, false);
66+
static const mp_arg_t allowed_args[] = {
67+
{ MP_QSTR_sock, MP_ARG_REQUIRED | MP_ARG_OBJ, {.u_obj = MP_OBJ_NULL} },
68+
{ MP_QSTR_use_blocking_writes, MP_ARG_BOOL, {.u_bool = false} },
69+
{ MP_QSTR_is_client, MP_ARG_BOOL, {.u_bool = false} },
70+
{ MP_QSTR_debug_mask, MP_ARG_KW_ONLY | MP_ARG_OBJ, {.u_obj = MP_OBJ_NULL} },
71+
};
72+
73+
// parse args
74+
struct {
75+
mp_arg_val_t sock, use_blocking_writes, is_client, debug_mask;
76+
} arg_vals;
77+
mp_arg_parse_all_kw_array(n_args, n_kw, args,
78+
MP_ARRAY_SIZE(allowed_args), allowed_args, (mp_arg_val_t*)&arg_vals);
79+
6280
mp_obj_websocket_t *o = m_new_obj(mp_obj_websocket_t);
6381
o->base.type = type;
64-
o->sock = args[0];
82+
o->sock = arg_vals.sock.u_obj;
83+
o->do_write_masking = !arg_vals.is_client.u_bool ? NO_WRITE_MASKING : NORMAL_WRITE_MASKING;
84+
if (arg_vals.debug_mask.u_obj != MP_OBJ_NULL) {
85+
mp_buffer_info_t bufinfo;
86+
mp_get_buffer_raise(arg_vals.debug_mask.u_obj, &bufinfo, MP_BUFFER_READ);
87+
if (bufinfo.len != 4) {
88+
mp_raise_ValueError("debug mask must have length of 4");
89+
}
90+
o->do_write_masking = DEBUG_WRITE_MASKING;
91+
memcpy(o->debug_write_mask, bufinfo.buf, 4);
92+
}
6593
o->state = FRAME_HEADER;
6694
o->to_recv = 2;
67-
o->mask_pos = 0;
95+
o->read_mask_pos = 0;
6896
o->buf_pos = 0;
6997
o->opts = FRAME_TXT;
70-
if (n_args > 1 && args[1] == mp_const_true) {
98+
if (arg_vals.use_blocking_writes.u_bool) {
7199
o->opts |= BLOCKING_WRITE;
72100
}
73101
return MP_OBJ_FROM_PTR(o);
@@ -111,7 +139,7 @@ STATIC mp_uint_t websocket_read(mp_obj_t self_in, void *buf, mp_uint_t size, int
111139

112140
// Reset mask in case someone will use "simplified" protocol
113141
// without masks.
114-
memset(self->mask, 0, sizeof(self->mask));
142+
memset(self->read_mask, 0, sizeof(self->read_mask));
115143

116144
int to_recv = 0;
117145
size_t sz = self->buf[1] & 0x7f;
@@ -149,7 +177,7 @@ STATIC mp_uint_t websocket_read(mp_obj_t self_in, void *buf, mp_uint_t size, int
149177
}
150178
if (self->buf_pos >= 4) {
151179
// Last 4 bytes is mask
152-
memcpy(self->mask, self->buf + self->buf_pos - 4, 4);
180+
memcpy(self->read_mask, self->buf + self->buf_pos - 4, 4);
153181
}
154182
self->buf_pos = 0;
155183
if ((self->last_flags & FRAME_OPCODE_MASK) >= FRAME_CLOSE) {
@@ -176,7 +204,7 @@ STATIC mp_uint_t websocket_read(mp_obj_t self_in, void *buf, mp_uint_t size, int
176204

177205
sz = out_sz;
178206
for (byte *p = buf; sz--; p++) {
179-
*p ^= self->mask[self->mask_pos++ & 3];
207+
*p ^= self->read_mask[self->read_mask_pos++ & 3];
180208
}
181209

182210
self->msg_sz -= out_sz;
@@ -186,7 +214,7 @@ STATIC mp_uint_t websocket_read(mp_obj_t self_in, void *buf, mp_uint_t size, int
186214
last_state = self->state;
187215
self->state = FRAME_HEADER;
188216
self->to_recv = 2;
189-
self->mask_pos = 0;
217+
self->read_mask_pos = 0;
190218
self->buf_pos = 0;
191219

192220
// Handle control frame
@@ -218,7 +246,7 @@ STATIC mp_uint_t websocket_read(mp_obj_t self_in, void *buf, mp_uint_t size, int
218246
STATIC mp_uint_t websocket_write(mp_obj_t self_in, const void *buf, mp_uint_t size, int *errcode) {
219247
mp_obj_websocket_t *self = MP_OBJ_TO_PTR(self_in);
220248
assert(size < 0x10000);
221-
byte header[4] = {0x80 | (self->opts & FRAME_OPCODE_MASK)};
249+
byte header[8] = {0x80 | (self->opts & FRAME_OPCODE_MASK)};
222250
int hdr_sz;
223251
if (size < 126) {
224252
header[1] = size;
@@ -229,6 +257,34 @@ STATIC mp_uint_t websocket_write(mp_obj_t self_in, const void *buf, mp_uint_t si
229257
header[3] = size & 0xff;
230258
hdr_sz = 4;
231259
}
260+
if (self->do_write_masking != NO_WRITE_MASKING) {
261+
hdr_sz += 4;
262+
header[1] |= 0x80;
263+
if (self->do_write_masking == NORMAL_WRITE_MASKING) {
264+
265+
// RFC6455 Section 5.3 states that the masking key must be derived
266+
// from a strong source of entropy. The "urandom" module doesn't
267+
// qualify in this regard, but there isn't any cross-platform
268+
// alternative. Fortunately, the purpose of masking is not
269+
// cryptographically motivated. The "urandom" module should be
270+
// seeded though, otherwise upon restart, the same sequence of
271+
// masks will always be used. A seed could be derived from a
272+
// network resource, a network interface's characteristics or
273+
// statistics, or a platform specific resource. Examples of using
274+
// a platform specific resource include reading an ESP8266's
275+
// 32-bit Random Number Generator register, or reading consecutive
276+
// values from a floating analog pin.
277+
mp_obj_t dest[3];
278+
mp_load_method(mp_module_get(MP_QSTR_urandom), MP_QSTR_getrandbits, dest);
279+
dest[2] = mp_obj_new_int(32);
280+
unsigned int randbits = MP_OBJ_SMALL_INT_VALUE(mp_call_method_n_kw(1, 0, dest));
281+
for (int i = 0; i < 4; ++i) {
282+
header[hdr_sz - 4 + i] = (randbits >> ((i ^ 3) << 3)) & 0xff;
283+
}
284+
} else if (self->do_write_masking == DEBUG_WRITE_MASKING) {
285+
memcpy(&header[hdr_sz - 4], self->debug_write_mask, 4);
286+
}
287+
}
232288

233289
mp_obj_t dest[3];
234290
if (self->opts & BLOCKING_WRITE) {
@@ -239,7 +295,15 @@ STATIC mp_uint_t websocket_write(mp_obj_t self_in, const void *buf, mp_uint_t si
239295

240296
mp_uint_t out_sz = mp_stream_write_exactly(self->sock, header, hdr_sz, errcode);
241297
if (*errcode == 0) {
242-
out_sz = mp_stream_write_exactly(self->sock, buf, size, errcode);
298+
if (self->do_write_masking == NO_WRITE_MASKING) {
299+
out_sz = mp_stream_write_exactly(self->sock, buf, size, errcode);
300+
} else {
301+
byte masked_buf[size];
302+
for (mp_uint_t i = 0; i < size; ++i) {
303+
masked_buf[i] = ((byte*)buf)[i] ^ header[hdr_sz - 4 + (i & 3)];
304+
}
305+
out_sz = mp_stream_write_exactly(self->sock, masked_buf, size, errcode);
306+
}
243307
}
244308

245309
if (self->opts & BLOCKING_WRITE) {

‎tests/extmod/websocket_basic.py

Lines changed: 13 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,20 +1,26 @@
11
try:
2+
import urandom
23
import uio
34
import uerrno
45
import websocket
56
except ImportError:
67
print("SKIP")
78
raise SystemExit
89

10+
# When writing a websocket frame as a client, the masking key used is obtained
11+
# from the "urandom" module. Seeding the random number generator with a
12+
# constant guarantees that the masking keys used are deterministic.
13+
urandom.seed(0x875513b0)
14+
915
# put raw data in the stream and do a websocket read
1016
def ws_read(msg, sz):
1117
ws = websocket.websocket(uio.BytesIO(msg))
1218
return ws.read(sz)
1319

1420
# do a websocket write and then return the raw data from the stream
15-
def ws_write(msg, sz):
21+
def ws_write(msg, sz, **kwargs):
1622
s = uio.BytesIO()
17-
ws = websocket.websocket(s)
23+
ws = websocket.websocket(s, **kwargs)
1824
ws.write(msg)
1925
s.seek(0)
2026
return s.read(sz)
@@ -31,9 +37,13 @@ def ws_write(msg, sz):
3137
print(ws_read(b'\x81~\x00\x80' + b'ping' * 32, 128))
3238
print(ws_write(b"pong" * 32, 132))
3339

34-
# mask (returned data will be 'mask' ^ 'mask')
40+
# read mask (returned data will be 'mask' ^ 'mask')
3541
print(ws_read(b"\x81\x84maskmask", 4))
3642

43+
# write mask
44+
print(ws_write(b"pong", 10, debug_mask=b"\x01\x00\x01\x00"))
45+
print(ws_write(b"pong", 10, is_client=True))
46+
3747
# close control frame
3848
s = uio.BytesIO(b'\x88\x00') # FRAME_CLOSE
3949
ws = websocket.websocket(s)

‎tests/extmod/websocket_basic.py.exp

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,8 @@ b'\x81\x04pong'
44
b'pingpingpingpingpingpingpingpingpingpingpingpingpingpingpingpingpingpingpingpingpingpingpingpingpingpingpingpingpingpingpingping'
55
b'\x81~\x00\x80pongpongpongpongpongpongpongpongpongpongpongpongpongpongpongpongpongpongp 5764 ongpongpongpongpongpongpongpongpongpongpongpongpongpong'
66
b'\x00\x00\x00\x00'
7+
b'\x81\x84\x01\x00\x01\x00qoog'
8+
b'\x81\x84\x90\x03\xbf\xee\xe0l\xd1\x89'
79
b''
810
b'\x81\x02\x88\x00'
911
b'ping'

0 commit comments

Comments
 (0)
0