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 "py/builtin.h"
34 #ifdef MICROPY_PY_WEBREPL_DELAY
35 #include "py/mphal.h"
36 #endif
37 #include "extmod/moduwebsocket.h"
38 
39 #if MICROPY_PY_WEBREPL
40 
41 #if 0 // print debugging info
42 #define DEBUG_printf DEBUG_printf
43 #else // don't print debugging info
44 #define DEBUG_printf(...) (void)0
45 #endif
46 
47 struct webrepl_file {
48     char sig[2];
49     char type;
50     char flags;
51     uint64_t offset;
52     uint32_t size;
53     uint16_t fname_len;
54     char fname[64];
55 } __attribute__((packed));
56 
57 enum { PUT_FILE = 1, GET_FILE, GET_VER };
58 enum { STATE_PASSWD, STATE_NORMAL };
59 
60 typedef struct _mp_obj_webrepl_t {
61     mp_obj_base_t base;
62     mp_obj_t sock;
63     byte state;
64     byte hdr_to_recv;
65     uint32_t data_to_recv;
66     struct webrepl_file hdr;
67     mp_obj_t cur_file;
68 } mp_obj_webrepl_t;
69 
70 STATIC const char passwd_prompt[] = "Password: ";
71 STATIC const char connected_prompt[] = "\r\nWebREPL connected\r\n>>> ";
72 STATIC const char denied_prompt[] = "\r\nAccess denied\r\n";
73 
74 STATIC char webrepl_passwd[10];
75 
write_webrepl(mp_obj_t websock,const void * buf,size_t len)76 STATIC void write_webrepl(mp_obj_t websock, const void *buf, size_t len) {
77     const mp_stream_p_t *sock_stream = mp_get_stream(websock);
78     int err;
79     int old_opts = sock_stream->ioctl(websock, MP_STREAM_SET_DATA_OPTS, FRAME_BIN, &err);
80     sock_stream->write(websock, buf, len, &err);
81     sock_stream->ioctl(websock, MP_STREAM_SET_DATA_OPTS, old_opts, &err);
82 }
83 
84 #define SSTR(s) s, sizeof(s) - 1
write_webrepl_str(mp_obj_t websock,const char * str,int sz)85 STATIC void write_webrepl_str(mp_obj_t websock, const char *str, int sz) {
86     int err;
87     const mp_stream_p_t *sock_stream = mp_get_stream(websock);
88     sock_stream->write(websock, str, sz, &err);
89 }
90 
write_webrepl_resp(mp_obj_t websock,uint16_t code)91 STATIC void write_webrepl_resp(mp_obj_t websock, uint16_t code) {
92     char buf[4] = {'W', 'B', code & 0xff, code >> 8};
93     write_webrepl(websock, buf, sizeof(buf));
94 }
95 
webrepl_make_new(const mp_obj_type_t * type,size_t n_args,size_t n_kw,const mp_obj_t * args)96 STATIC mp_obj_t webrepl_make_new(const mp_obj_type_t *type, size_t n_args, size_t n_kw, const mp_obj_t *args) {
97     mp_arg_check_num(n_args, n_kw, 1, 2, false);
98     mp_get_stream_raise(args[0], MP_STREAM_OP_READ | MP_STREAM_OP_WRITE | MP_STREAM_OP_IOCTL);
99     DEBUG_printf("sizeof(struct webrepl_file) = %lu\n", sizeof(struct webrepl_file));
100     mp_obj_webrepl_t *o = m_new_obj(mp_obj_webrepl_t);
101     o->base.type = type;
102     o->sock = args[0];
103     o->hdr_to_recv = sizeof(struct webrepl_file);
104     o->data_to_recv = 0;
105     o->state = STATE_PASSWD;
106     write_webrepl_str(args[0], SSTR(passwd_prompt));
107     return MP_OBJ_FROM_PTR(o);
108 }
109 
check_file_op_finished(mp_obj_webrepl_t * self)110 STATIC void check_file_op_finished(mp_obj_webrepl_t *self) {
111     if (self->data_to_recv == 0) {
112         mp_stream_close(self->cur_file);
113         self->hdr_to_recv = sizeof(struct webrepl_file);
114         DEBUG_printf("webrepl: Finished file operation %d\n", self->hdr.type);
115         write_webrepl_resp(self->sock, 0);
116     }
117 }
118 
write_file_chunk(mp_obj_webrepl_t * self)119 STATIC int write_file_chunk(mp_obj_webrepl_t *self) {
120     const mp_stream_p_t *file_stream = mp_get_stream(self->cur_file);
121     byte readbuf[2 + 256];
122     int err;
123     mp_uint_t out_sz = file_stream->read(self->cur_file, readbuf + 2, sizeof(readbuf) - 2, &err);
124     if (out_sz == MP_STREAM_ERROR) {
125         return out_sz;
126     }
127     readbuf[0] = out_sz;
128     readbuf[1] = out_sz >> 8;
129     DEBUG_printf("webrepl: Sending %d bytes of file\n", out_sz);
130     write_webrepl(self->sock, readbuf, 2 + out_sz);
131     return out_sz;
132 }
133 
handle_op(mp_obj_webrepl_t * self)134 STATIC void handle_op(mp_obj_webrepl_t *self) {
135 
136     // Handle operations not requiring opened file
137 
138     switch (self->hdr.type) {
139         case GET_VER: {
140             static const char ver[] = {MICROPY_VERSION_MAJOR, MICROPY_VERSION_MINOR, MICROPY_VERSION_MICRO};
141             write_webrepl(self->sock, ver, sizeof(ver));
142             self->hdr_to_recv = sizeof(struct webrepl_file);
143             return;
144         }
145     }
146 
147     // Handle operations requiring opened file
148 
149     mp_obj_t open_args[2] = {
150         mp_obj_new_str(self->hdr.fname, strlen(self->hdr.fname)),
151         MP_OBJ_NEW_QSTR(MP_QSTR_rb)
152     };
153 
154     if (self->hdr.type == PUT_FILE) {
155         open_args[1] = MP_OBJ_NEW_QSTR(MP_QSTR_wb);
156     }
157 
158     self->cur_file = mp_builtin_open(2, open_args, (mp_map_t *)&mp_const_empty_map);
159 
160     #if 0
161     struct mp_stream_seek_t seek = { .offset = self->hdr.offset, .whence = 0 };
162     int err;
163     mp_uint_t res = file_stream->ioctl(self->cur_file, MP_STREAM_SEEK, (uintptr_t)&seek, &err);
164     assert(res != MP_STREAM_ERROR);
165     #endif
166 
167     write_webrepl_resp(self->sock, 0);
168 
169     if (self->hdr.type == PUT_FILE) {
170         self->data_to_recv = self->hdr.size;
171         check_file_op_finished(self);
172     } else if (self->hdr.type == GET_FILE) {
173         self->data_to_recv = 1;
174     }
175 }
176 
177 STATIC mp_uint_t _webrepl_read(mp_obj_t self_in, void *buf, mp_uint_t size, int *errcode);
178 
webrepl_read(mp_obj_t self_in,void * buf,mp_uint_t size,int * errcode)179 STATIC mp_uint_t webrepl_read(mp_obj_t self_in, void *buf, mp_uint_t size, int *errcode) {
180     mp_uint_t out_sz;
181     do {
182         out_sz = _webrepl_read(self_in, buf, size, errcode);
183     } while (out_sz == -2);
184     return out_sz;
185 }
186 
_webrepl_read(mp_obj_t self_in,void * buf,mp_uint_t size,int * errcode)187 STATIC mp_uint_t _webrepl_read(mp_obj_t self_in, void *buf, mp_uint_t size, int *errcode) {
188     // We know that os.dupterm always calls with size = 1
189     assert(size == 1);
190     mp_obj_webrepl_t *self = MP_OBJ_TO_PTR(self_in);
191     const mp_stream_p_t *sock_stream = mp_get_stream(self->sock);
192     mp_uint_t out_sz = sock_stream->read(self->sock, buf, size, errcode);
193     // DEBUG_printf("webrepl: Read %d initial bytes from websocket\n", out_sz);
194     if (out_sz == 0 || out_sz == MP_STREAM_ERROR) {
195         return out_sz;
196     }
197 
198     if (self->state == STATE_PASSWD) {
199         char c = *(char *)buf;
200         if (c == '\r' || c == '\n') {
201             self->hdr.fname[self->data_to_recv] = 0;
202             DEBUG_printf("webrepl: entered password: %s\n", self->hdr.fname);
203 
204             if (strcmp(self->hdr.fname, webrepl_passwd) != 0) {
205                 write_webrepl_str(self->sock, SSTR(denied_prompt));
206                 return 0;
207             }
208 
209             self->state = STATE_NORMAL;
210             self->data_to_recv = 0;
211             write_webrepl_str(self->sock, SSTR(connected_prompt));
212         } else if (self->data_to_recv < 10) {
213             self->hdr.fname[self->data_to_recv++] = c;
214         }
215         return -2;
216     }
217 
218     // If last read data belonged to text record (== REPL)
219     int err;
220     if (sock_stream->ioctl(self->sock, MP_STREAM_GET_DATA_OPTS, 0, &err) == 1) {
221         return out_sz;
222     }
223 
224     DEBUG_printf("webrepl: received bin data, hdr_to_recv: %d, data_to_recv=%d\n", self->hdr_to_recv, self->data_to_recv);
225 
226     if (self->hdr_to_recv != 0) {
227         char *p = (char *)&self->hdr + sizeof(self->hdr) - self->hdr_to_recv;
228         *p++ = *(char *)buf;
229         if (--self->hdr_to_recv != 0) {
230             mp_uint_t hdr_sz = sock_stream->read(self->sock, p, self->hdr_to_recv, errcode);
231             if (hdr_sz == MP_STREAM_ERROR) {
232                 return hdr_sz;
233             }
234             self->hdr_to_recv -= hdr_sz;
235             if (self->hdr_to_recv != 0) {
236                 return -2;
237             }
238         }
239 
240         DEBUG_printf("webrepl: op: %d, file: %s, chunk @%x, sz=%d\n", self->hdr.type, self->hdr.fname, (uint32_t)self->hdr.offset, self->hdr.size);
241 
242         handle_op(self);
243 
244         return -2;
245     }
246 
247     if (self->data_to_recv != 0) {
248         // Ports that don't have much available stack can make this filebuf static
249         #if MICROPY_PY_WEBREPL_STATIC_FILEBUF
250         static
251         #endif
252         byte filebuf[512];
253         filebuf[0] = *(byte *)buf;
254         mp_uint_t buf_sz = 1;
255         if (--self->data_to_recv != 0) {
256             size_t to_read = MIN(sizeof(filebuf) - 1, self->data_to_recv);
257             mp_uint_t sz = sock_stream->read(self->sock, filebuf + 1, to_read, errcode);
258             if (sz == MP_STREAM_ERROR) {
259                 return sz;
260             }
261             self->data_to_recv -= sz;
262             buf_sz += sz;
263         }
264 
265         if (self->hdr.type == PUT_FILE) {
266             DEBUG_printf("webrepl: Writing %lu bytes to file\n", buf_sz);
267             int err;
268             mp_uint_t res = mp_stream_write_exactly(self->cur_file, filebuf, buf_sz, &err);
269             if (err != 0 || res != buf_sz) {
270                 assert(0);
271             }
272         } else if (self->hdr.type == GET_FILE) {
273             assert(buf_sz == 1);
274             assert(self->data_to_recv == 0);
275             assert(filebuf[0] == 0);
276             mp_uint_t out_sz = write_file_chunk(self);
277             if (out_sz != 0) {
278                 self->data_to_recv = 1;
279             }
280         }
281 
282         check_file_op_finished(self);
283 
284         #ifdef MICROPY_PY_WEBREPL_DELAY
285         // Some platforms may have broken drivers and easily gets
286         // overloaded with modest traffic WebREPL file transfers
287         // generate. The basic workaround is a crude rate control
288         // done in such way.
289         mp_hal_delay_ms(MICROPY_PY_WEBREPL_DELAY);
290         #endif
291     }
292 
293     return -2;
294 }
295 
webrepl_write(mp_obj_t self_in,const void * buf,mp_uint_t size,int * errcode)296 STATIC mp_uint_t webrepl_write(mp_obj_t self_in, const void *buf, mp_uint_t size, int *errcode) {
297     mp_obj_webrepl_t *self = MP_OBJ_TO_PTR(self_in);
298     if (self->state == STATE_PASSWD) {
299         // Don't forward output until passwd is entered
300         return size;
301     }
302     const mp_stream_p_t *stream_p = mp_get_stream(self->sock);
303     return stream_p->write(self->sock, buf, size, errcode);
304 }
305 
webrepl_ioctl(mp_obj_t o_in,mp_uint_t request,uintptr_t arg,int * errcode)306 STATIC mp_uint_t webrepl_ioctl(mp_obj_t o_in, mp_uint_t request, uintptr_t arg, int *errcode) {
307     mp_obj_webrepl_t *self = MP_OBJ_TO_PTR(o_in);
308     (void)arg;
309     switch (request) {
310         case MP_STREAM_CLOSE:
311             // TODO: This is a place to do cleanup
312             mp_stream_close(self->sock);
313             return 0;
314 
315         default:
316             *errcode = MP_EINVAL;
317             return MP_STREAM_ERROR;
318     }
319 }
320 
webrepl_set_password(mp_obj_t passwd_in)321 STATIC mp_obj_t webrepl_set_password(mp_obj_t passwd_in) {
322     size_t len;
323     const char *passwd = mp_obj_str_get_data(passwd_in, &len);
324     if (len > sizeof(webrepl_passwd) - 1) {
325         mp_raise_ValueError(NULL);
326     }
327     strcpy(webrepl_passwd, passwd);
328     return mp_const_none;
329 }
330 STATIC MP_DEFINE_CONST_FUN_OBJ_1(webrepl_set_password_obj, webrepl_set_password);
331 
332 STATIC const mp_rom_map_elem_t webrepl_locals_dict_table[] = {
333     { MP_ROM_QSTR(MP_QSTR_read), MP_ROM_PTR(&mp_stream_read_obj) },
334     { MP_ROM_QSTR(MP_QSTR_readinto), MP_ROM_PTR(&mp_stream_readinto_obj) },
335     { MP_ROM_QSTR(MP_QSTR_write), MP_ROM_PTR(&mp_stream_write_obj) },
336     { MP_ROM_QSTR(MP_QSTR_close), MP_ROM_PTR(&mp_stream_close_obj) },
337 };
338 STATIC MP_DEFINE_CONST_DICT(webrepl_locals_dict, webrepl_locals_dict_table);
339 
340 STATIC const mp_stream_p_t webrepl_stream_p = {
341     .read = webrepl_read,
342     .write = webrepl_write,
343     .ioctl = webrepl_ioctl,
344 };
345 
346 STATIC const mp_obj_type_t webrepl_type = {
347     { &mp_type_type },
348     .name = MP_QSTR__webrepl,
349     .make_new = webrepl_make_new,
350     .protocol = &webrepl_stream_p,
351     .locals_dict = (mp_obj_dict_t *)&webrepl_locals_dict,
352 };
353 
354 STATIC const mp_rom_map_elem_t webrepl_module_globals_table[] = {
355     { MP_ROM_QSTR(MP_QSTR___name__), MP_ROM_QSTR(MP_QSTR__webrepl) },
356     { MP_ROM_QSTR(MP_QSTR__webrepl), MP_ROM_PTR(&webrepl_type) },
357     { MP_ROM_QSTR(MP_QSTR_password), MP_ROM_PTR(&webrepl_set_password_obj) },
358 };
359 
360 STATIC MP_DEFINE_CONST_DICT(webrepl_module_globals, webrepl_module_globals_table);
361 
362 const mp_obj_module_t mp_module_webrepl = {
363     .base = { &mp_type_module },
364     .globals = (mp_obj_dict_t *)&webrepl_module_globals,
365 };
366 
367 #endif // MICROPY_PY_WEBREPL
368