[Refactor] Link Play TCP data transmission

- Code refactor of Link Play TCP data transmission for better security and scalability

latest version/linkplay_server/config.py

@@ -4,19 +4,22 @@ class Config:
     '''
 
     '''
-        服务器地址、端口号、校验码
-        Server address, port and verification code
+        服务器地址、端口号、校验码、传输加密密钥
+        Server address, port, verification code, and encryption key
     '''
     HOST = '0.0.0.0'
     UDP_PORT = 10900
     TCP_PORT = 10901
     AUTHENTICATION = 'my_link_play_server'
+    TCP_SECRET_KEY = '1145141919810'
     '''
     --------------------------------------------------
     '''
 
     DEBUG = False
 
+    TCP_MAX_LENGTH = 0x0FFFFFFF
+
     TIME_LIMIT = 3600000
 
     COMMAND_INTERVAL = 1000000

latest version/linkplay_server/main.py

@@ -2,6 +2,7 @@
 import logging
 import socketserver
 import threading
+from json import dumps, loads
 
 from .aes import decrypt, encrypt
 from .config import Config
@@ -55,27 +56,52 @@ class UDP_handler(socketserver.BaseRequestHandler):
                           ciphertext, self.client_address)
 
 
+AUTH_LEN = len(Config.AUTHENTICATION)
+TCP_AES_KEY = Config.TCP_SECRET_KEY.encode('utf-8').ljust(16, b'\x00')[:16]
+
+
 class TCP_handler(socketserver.StreamRequestHandler):
+
     def handle(self):
         try:
-            self.data = self.rfile.readline().strip()
+            if self.rfile.read(AUTH_LEN).decode('utf-8') != Config.AUTHENTICATION:
+                self.wfile.write(b'No authentication')
+                logging.warning(
+                    f'TCP-{self.client_address[0]}-No authentication')
+                return None
+
+            cipher_len = int.from_bytes(self.rfile.read(8), byteorder='little')
+            if cipher_len > Config.TCP_MAX_LENGTH:
+                self.wfile.write(b'Body too long')
+                logging.warning(f'TCP-{self.client_address[0]}-Body too long')
+                return None
+
+            iv = self.rfile.read(12)
+            tag = self.rfile.read(12)
+            ciphertext = self.rfile.read(cipher_len)
+
+            self.data = decrypt(TCP_AES_KEY, b'', iv, ciphertext, tag)
             message = self.data.decode('utf-8')
+            data = loads(message)
         except Exception as e:
             logging.error(e)
             return None
 
         if Config.DEBUG:
             logging.info(f'TCP-From-{self.client_address[0]}-{message}')
-        data = message.split('|')
-        if data[0] != Config.AUTHENTICATION:
-            self.wfile.write(b'No authentication')
-            logging.warning(f'TCP-{self.client_address[0]}-No authentication')
-            return None
 
-        r = TCPRouter(data[1:]).handle()
-        if Config.DEBUG:
-            logging.info(f'TCP-To-{self.client_address[0]}-{r}')
-        self.wfile.write(r.encode('utf-8'))
+        r = TCPRouter(data).handle()
+        try:
+            r = dumps(r)
+            if Config.DEBUG:
+                logging.info(f'TCP-To-{self.client_address[0]}-{r}')
+            iv, ciphertext, tag = encrypt(TCP_AES_KEY, r.encode('utf-8'), b'')
+            r = len(ciphertext).to_bytes(8, byteorder='little') + \
+                iv + tag[:12] + ciphertext
+        except Exception as e:
+            logging.error(e)
+            return None
+        self.wfile.write(r)
 
 
 def link_play(ip: str = Config.HOST, udp_port: int = Config.UDP_PORT, tcp_port: int = Config.TCP_PORT):

latest version/linkplay_server/store.py

@@ -87,19 +87,21 @@ def memory_clean(now):
 class TCPRouter:
     clean_timer = 0
     router = {
-        '0': 'debug',
-        '1': 'create_room',
-        '2': 'join_room',
-        '3': 'update_room',
+        'debug',
+        'create_room',
+        'join_room',
+        'update_room',
     }
 
-    def __init__(self, data: list):
-        self.data = data  # data: list[str] = [command, ...]
+    def __init__(self, raw_data: 'dict | list'):
+        self.raw_data = raw_data  # data: dict {endpoint: str, data: dict}
+        self.data = raw_data['data']
+        self.endpoint = raw_data['endpoint']
 
-    def debug(self):
+    def debug(self) -> dict:
         if Config.DEBUG:
-            return eval(self.data[1])
-        return 'ok'
+            return {'result': eval(self.data['code'])}
+        return {'hello_world': 'ok'}
 
     @staticmethod
     def clean_check():
@@ -109,14 +111,17 @@ class TCPRouter:
             TCPRouter.clean_timer = now
             memory_clean(now)
 
-    def handle(self) -> str:
+    def handle(self) -> dict:
         self.clean_check()
-        if self.data[0] not in self.router:
+        if self.endpoint not in self.router:
             return None
-        r = getattr(self, self.router[self.data[0]])()
-        if isinstance(r, tuple):
-            return '|'.join(map(str, r))
-        return str(r)
+        r = getattr(self, self.endpoint)()
+        if isinstance(r, int):
+            return {'code': r}
+        return {
+            'code': 0,
+            'data': r
+        }
 
     @staticmethod
     def generate_player(name: str) -> Player:
@@ -144,12 +149,12 @@ class TCPRouter:
 
         return room
 
-    def create_room(self) -> tuple:
+    def create_room(self) -> dict:
         # 开房
         # data = ['1', name, song_unlock, ]
         # song_unlock: base64 str
-        name = self.data[1]
-        song_unlock = b64decode(self.data[2])
+        name = self.data['name']
+        song_unlock = b64decode(self.data['song_unlock'])
 
         key = urandom(16)
         with Store.lock:
@@ -172,33 +177,39 @@ class TCPRouter:
             }
 
         logging.info(f'TCP-Create room `{room.room_code}` by player `{name}`')
-        return (0, room.room_code, room.room_id, token, b64encode(key).decode('utf-8'), player.player_id)
+        return {
+            'room_code': room.room_code,
+            'room_id': room.room_id,
+            'token': token,
+            'key': b64encode(key).decode('utf-8'),
+            'player_id': player.player_id
+        }
 
-    def join_room(self) -> tuple:
+    def join_room(self) -> 'dict | int':
         # 入房
         # data = ['2', name, song_unlock, room_code]
         # song_unlock: base64 str
-        room_code = self.data[3].upper()
+        room_code = self.data['room_code'].upper()
         key = urandom(16)
-        name = self.data[1]
-        song_unlock = b64decode(self.data[2])
+        name = self.data['name']
+        song_unlock = b64decode(self.data['song_unlock'])
 
         with Store.lock:
             if room_code not in Store.room_code_dict:
                 # 房间号错误 / 房间不存在
-                return '1202'
+                return 1202
             room: Room = Store.room_code_dict[room_code]
 
             player_num = room.player_num
             if player_num == 4:
                 # 满人
-                return '1201'
+                return 1201
             if player_num == 0:
                 # 房间不存在
-                return '1202'
+                return 1202
             if room.state != 2:
                 # 无法加入
-                return '1205'
+                return 1205
 
             token = unique_random(Store.link_play_data)
 
@@ -219,16 +230,30 @@ class TCPRouter:
             }
 
         logging.info(f'TCP-Player `{name}` joins room `{room_code}`')
-        return (0, room_code, room.room_id, token, b64encode(key).decode('utf-8'), player.player_id, b64encode(room.song_unlock).decode('utf-8'))
+        return {
+            'room_code': room_code,
+            'room_id': room.room_id,
+            'token': token,
+            'key': b64encode(key).decode('utf-8'),
+            'player_id': player.player_id,
+            'song_unlock': b64encode(room.song_unlock).decode('utf-8')
+        }
 
-    def update_room(self) -> tuple:
+    def update_room(self) -> dict:
         # 房间信息更新
         # data = ['3', token]
-        token = int(self.data[1])
+        token = int(self.data['token'])
         with Store.lock:
             if token not in Store.link_play_data:
-                return '108'
+                return 108
             r = Store.link_play_data[token]
             room = r['room']
             logging.info(f'TCP-Room `{room.room_code}` info update')
-            return (0, room.room_code, room.room_id, b64encode(r['key']).decode('utf-8'), room.players[r['player_index']].player_id, b64encode(room.song_unlock).decode('utf-8'))
+            return {
+                'room_code': room.room_code,
+                'room_id': room.room_id,
+                'key': b64encode(r['key']).decode('utf-8'),
+                # changed from room.players[r['player_index']].player_id,
+                'player_id': r['player_id'],
+                'song_unlock': b64encode(room.song_unlock).decode('utf-8')
+            }