Optimize code

Signed-off-by: Jianhui Zhao <jianhuizhao329@gmail.com>

server/rtty.py

@@ -14,9 +14,15 @@ from aiohttp import web, WSMsgType
 class Devices:
     devs = {}
     def add(self, ws, mac):
-        self.devs[mac] = {'ws': ws, 'active': 1}
+        self.devs[mac] = {'ws': ws, 'active': 3}
     def active(self, mac):
-        self.devs[mac]['active'] += 1
+        self.devs[mac]['active'] = 3
+    def flush(self):
+        for mac in list(self.devs):
+            self.devs[mac]['active'] -= 1
+            if self.devs[mac]['active'] == 0:
+                self.devs[mac]['ws'].close()
+                del self.devs[mac]
     def login(self, ws, mac):
         sid = md5((mac + str(random.uniform(1, 100))).encode('utf8')).hexdigest()
         dev = self.devs.get(mac)
@@ -27,7 +33,11 @@ class Devices:
         dev[sid] = ws
         dev['ws'].send_str(json.dumps({'type': 'login', 'mac': mac, 'sid': sid}))
         syslog.syslog('new logged to ' + mac)
-        return True
+        return sid
+    def logout(self, mac, sid):
+        dev = self.devs.get(mac)
+        del dev[sid]
+        dev['ws'].send_str(json.dumps({'type': 'logout', 'mac': mac, 'sid': sid}))
     def send_data2user(self, msg):
         mac = msg['mac']
         sid = msg['sid']
@@ -67,7 +77,8 @@ async def websocket_handler_browser(request):
     await ws.prepare(request)
 
     mac = request.query['mac']
-    if not devices.login(ws, mac):
+    sid = devices.login(ws, mac)
+    if not sid:
         ws.close()
         return ws
 
@@ -79,6 +90,7 @@ async def websocket_handler_browser(request):
                 devices.send_data2device(msg)
         elif msg.type == WSMsgType.ERROR:
             syslog.syslog('browser connection closed with exception %s' % ws.exception())
+    devices.logout(mac, sid)
     return ws
 
 async def handle_list(request):
@@ -105,4 +117,11 @@ app.router.add_get('/ws/browser', websocket_handler_browser)
 app.router.add_get('/', handle_root)
 app.router.add_static('/', path = document, name = 'static')
 
+def keepalive(loop):
+    devices.flush()
+    loop.call_later(5, keepalive, loop)
+
+loop = asyncio.get_event_loop()
+loop.call_later(5, keepalive, loop)
+
 web.run_app(app, port = port)

src/main.c

@@ -86,9 +86,19 @@ static void keepalive(struct uloop_timeout *utm)
     uloop_timeout_set(utm, KEEPALIVE_INTERVAL * 1000);
 }
 
+static void del_tty_session(struct tty_session *tty)
+{
+    list_del(&tty->node);
+    uloop_process_delete(&tty->up);
+    ustream_free(&tty->sfd.stream);
+    close(tty->pty);
+    kill(tty->pid, SIGTERM);
+    waitpid(tty->pid, NULL, 0);
+}
+
 static void pty_read_cb(struct ustream *s, int bytes)
 {
-    struct tty_session *ts = container_of(s, struct tty_session, sfd.stream);
+    struct tty_session *tty = container_of(s, struct tty_session, sfd.stream);
     char *str;
     int len;
 
@@ -97,7 +107,7 @@ static void pty_read_cb(struct ustream *s, int bytes)
     blobmsg_buf_init(&b);
     blobmsg_add_string(&b, "type", "data");
     blobmsg_add_string(&b, "mac", mac);
-    blobmsg_add_string(&b, "sid", ts->sid);
+    blobmsg_add_string(&b, "sid", tty->sid);
 
     b64_encode(str, len, buf, sizeof(buf));
     ustream_consume(s, len);
@@ -111,17 +121,19 @@ static void pty_read_cb(struct ustream *s, int bytes)
 
 static void pty_on_exit(struct uloop_process *p, int ret)
 {
-    struct tty_session *ts = container_of(p, struct tty_session, up);
+    struct tty_session *tty = container_of(p, struct tty_session, up);
     char *str;
 
     blobmsg_buf_init(&b);
     blobmsg_add_string(&b, "type", "logout");
     blobmsg_add_string(&b, "mac", mac);
-    blobmsg_add_string(&b, "sid", ts->sid);
+    blobmsg_add_string(&b, "sid", tty->sid);
 
     str = blobmsg_format_json(b.head, true);
     cl->send(cl, str, strlen(str), WEBSOCKET_OP_TEXT);
     free(str);
+
+    del_tty_session(tty);
 }
 
 static void new_tty_session(struct blob_attr **tb)
@@ -183,6 +195,15 @@ static void uwsc_onmessage(struct uwsc_client *cl, char *msg, uint64_t len, enum
     type = blobmsg_get_string(tb[RTTYD_TYPE]);
     if (!strcmp(type, "login")) {
         new_tty_session(tb);
+    } else if (!strcmp(type, "logout")) {
+        const char *sid = blobmsg_get_string(tb[RTTYD_SID]);
+        struct tty_session *tty, *tmp;
+
+        list_for_each_entry_safe(tty, tmp, &tty_sessions, node) {
+            if (!strcmp(tty->sid, sid)) {
+                del_tty_session(tty);
+            }
+        }
     } else if (!strcmp(type, "data")) {
         const char *sid = blobmsg_get_string(tb[RTTYD_SID]);
         const char *data = blobmsg_get_string(tb[RTTYD_DATA]);