feat: support snappy

Author: sazima

config_c.json

@@ -1,7 +1,8 @@
 {
   "server": {
     "url": "ws://192.168.9.224:18888/websocket_path",
-    "password": "helloworld"
+    "password": "helloworld",
+    "compress": true
   },
   "client_name": "windows10_sql",
   "log_file": "/var/log/nt/nt.log"

entity/client_config_entity.py

@@ -13,6 +13,7 @@ class _ServerEntity(TypedDict):
     https: bool
     password: str
     path: str
+    compress: bool
 
 
 class ClientConfigEntity(TypedDict):

run_client.py

@@ -10,13 +10,18 @@
 from optparse import OptionParser
 from threading import Thread
 from typing import List, Set, Dict
+try:
+    import snappy
+    has_snappy = True
+except ModuleNotFoundError:
+    has_snappy = False
 
 from tornado import ioloop
 
 sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
 
 from common.speed_limit import SpeedLimiter
-from common.websocket import WebSocketException
+from common.websocket import WebSocketException, ABNF, WebSocketConnectionClosedException
 
 from client.clear_nonce_task import ClearNonceTask
 from client.heart_beat_task import HeatBeatTask
@@ -113,12 +118,23 @@ def __init__(self, ws: websocket.WebSocketApp, tcp_forward_client, heart_beat_ta
         self.ws.on_message = self.on_message
         self.ws.on_close = self.on_close
         self.ws.on_open = self.on_open
+        self.ws.send = self.send
         self.forward_client: TcpForwardClient = tcp_forward_client
         self.heart_beat_task = heart_beat_task
         self.config_data: ClientConfigEntity = config_data
+        self.compress_support: bool = config_data['server']['compress']
+
+    def send(self, data, opcode=ABNF.OPCODE_TEXT):
+        if opcode == ABNF.OPCODE_BINARY and self.compress_support:
+            data = snappy.snappy.compress(data)
+        if not self.ws.sock or self.ws.sock.send(data, opcode) == 0:
+            raise WebSocketConnectionClosedException(
+                "Connection is already closed.")
 
     def on_message(self, ws, message: bytes):
         try:
+            if self.compress_support:
+                message = snappy.snappy.uncompress(message)
             message_data: MessageEntity = NatSerialization.loads(message, ContextUtils.get_password())
             start_time = time.time()
             time_ = message_data['type_']
@@ -222,7 +238,17 @@ def main():
         else:
             url += 'ws://'
         url += f"{server_config['host']}:{str(server_config['port'])}{server_config['path']}"
+    config_data['server'].setdefault('compress', False)
+    compress_support = config_data['server']['compress']
+    assert isinstance(compress_support, bool)
+    if compress_support and not has_snappy:
+        raise Exception('snappy is not installed')
     LoggerFactory.get_logger().info(f'start open {url}')
+    if compress_support:
+        if '?' in url:  # 补充 compress_support 参数
+            url += '&c=' + json.dumps(compress_support)
+        else:
+            url += '?c=' + json.dumps(compress_support)
     ws = websocket.WebSocketApp(url)
     forward_client = TcpForwardClient(ws)
     heart_beat_task = HeatBeatTask(ws)

server/websocket_handler.py

@@ -1,15 +1,18 @@
 import asyncio
 import json
 import logging
-import socket
 import time
 import traceback
 from asyncio import Lock
-from collections import defaultdict
-from threading import Thread
-from typing import List, Dict, Set, Tuple
+from json import JSONDecodeError
+from typing import List, Dict, Set
+
+try:
+    import snappy
+    has_snappy = True
+except ModuleNotFoundError:
+    has_snappy = False
 
-from tornado.ioloop import IOLoop
 from tornado.websocket import WebSocketHandler
 
 from common.nat_serialization import NatSerialization
@@ -33,19 +36,32 @@ class MyWebSocketaHandler(WebSocketHandler):
     names: Set[str]
     recv_time: float = None
 
+    compress_support: bool = False  # 是否支持snappy压缩
+
     # handler_to_recv_time: Dict['MyWebSocketaHandler', float] = {}
     client_name_to_handler: Dict[str, 'MyWebSocketaHandler'] = {}
     lock = Lock()
 
     def open(self, *args: str, **kwargs: str):
         self.client_name = None
         self.version = None
-        LoggerFactory.get_logger().info('new open websocket')
+        try:
+            self.compress_support = json.loads(self.get_argument('c', 'false'))
+        except JSONDecodeError:
+            self.compress_support = False
+        if self.compress_support and not has_snappy:
+            msg = 'python-snappy is not installed on the server'
+            LoggerFactory.get_logger().info(msg)
+            self.close(reason=msg)
+        LoggerFactory.get_logger().info(f'new open websocket, compress_support: {self.compress_support}')
 
     async def write_message(self, message, binary=False):
         start_time = time.time()
         try:
-            await (super(MyWebSocketaHandler, self).write_message(bytes(message), binary))
+            byte_message = bytes(message)
+            if self.compress_support:
+                byte_message = snappy.snappy.compress(byte_message)
+            await (super(MyWebSocketaHandler, self).write_message(byte_message, binary))
             if LoggerFactory.get_logger().isEnabledFor(logging.DEBUG):
                 LoggerFactory.get_logger().debug(f'write message cost time {time.time() - start_time}, len: {len(message)}')
             return
@@ -55,6 +71,8 @@ async def write_message(self, message, binary=False):
             raise
 
     def on_message(self, m_bytes):
+        if self.compress_support:
+            m_bytes = snappy.snappy.uncompress(m_bytes)
         asyncio.ensure_future(self.on_message_async(m_bytes))
 
     async def on_message_async(self, message):

setup.py

@@ -27,5 +27,8 @@
     nt_server = {package_name}.run_server:main
     """,
     packages=l,
-    install_requires=['tornado',  'typing_extensions']
+    install_requires=['tornado',  'typing_extensions'],
+    extras_require={
+        "snappy": ["python-snappy"],
+    },
 )