From cb17c7b3b510570d565176a501e5adf106aac47c Mon Sep 17 00:00:00 2001
From: BreakWa11 <mmgac001@gmail.com>
Date: Mon, 2 Nov 2015 15:13:54 +0800
Subject: [PATCH] let setting & code simple add "protocol" in config.json

---
 config.json                             |  1 +
 shadowsocks/obfs.py                     |  4 ++++
 shadowsocks/obfsplugin/plain.py         |  4 ++++
 shadowsocks/obfsplugin/verify_simple.py | 31 +------------------------
 shadowsocks/shell.py                    |  1 +
 shadowsocks/tcprelay.py                 | 27 ++++++++++++++++-----
 6 files changed, 32 insertions(+), 36 deletions(-)

diff --git a/config.json b/config.json
index 2fbf33e..97cd259 100644
--- a/config.json
+++ b/config.json
@@ -7,6 +7,7 @@
     "password": "m",
     "timeout": 300,
     "method": "aes-256-cfb",
+    "protocol": "origin",
     "obfs": "http_simple_compatible",
     "obfs_param": "",
     "dns_ipv6": false,
diff --git a/shadowsocks/obfs.py b/shadowsocks/obfs.py
index 92de548..1752a56 100644
--- a/shadowsocks/obfs.py
+++ b/shadowsocks/obfs.py
@@ -83,3 +83,7 @@ class obfs(object):
     def server_post_decrypt(self, buf):
         return self.obfs.server_post_decrypt(buf)
 
+    def dispose(self):
+        self.obfs.dispose()
+        del self.obfs
+
diff --git a/shadowsocks/obfsplugin/plain.py b/shadowsocks/obfsplugin/plain.py
index cc602fa..5450e7a 100644
--- a/shadowsocks/obfsplugin/plain.py
+++ b/shadowsocks/obfsplugin/plain.py
@@ -27,6 +27,7 @@ def create_obfs(method):
 
 obfs_map = {
         'plain': (create_obfs,),
+        'origin': (create_obfs,),
 }
 
 class plain(object):
@@ -66,3 +67,6 @@ class plain(object):
     def server_post_decrypt(self, buf):
         return buf
 
+    def dispose(self):
+        pass
+
diff --git a/shadowsocks/obfsplugin/verify_simple.py b/shadowsocks/obfsplugin/verify_simple.py
index da63d52..8945d96 100644
--- a/shadowsocks/obfsplugin/verify_simple.py
+++ b/shadowsocks/obfsplugin/verify_simple.py
@@ -57,57 +57,29 @@ def match_begin(str1, str2):
 
 class obfs_verify_data(object):
     def __init__(self):
-        self.sub_obfs = None
+        pass
 
 class verify_base(plain.plain):
     def __init__(self, method):
         super(verify_base, self).__init__(method)
         self.method = method
-        self.sub_obfs = None
 
     def init_data(self):
         return obfs_verify_data()
 
     def set_server_info(self, server_info):
-        try:
-            if server_info.param:
-                sub_param = ''
-                param_list = server_info.param.split(',', 1)
-                if len(param_list) > 1:
-                    self.sub_obfs = shadowsocks.obfs.obfs(param_list[0])
-                    sub_param = param_list[1]
-                else:
-                    self.sub_obfs = shadowsocks.obfs.obfs(server_info.param)
-                if server_info.data.sub_obfs is None:
-                    server_info.data.sub_obfs = self.sub_obfs.init_data()
-                _server_info = shadowsocks.obfs.server_info(server_info.data.sub_obfs)
-                _server_info.host = server_info.host
-                _server_info.port = server_info.port
-                _server_info.tcp_mss = server_info.tcp_mss
-                _server_info.param = sub_param
-                self.sub_obfs.set_server_info(_server_info)
-        except Exception as e:
-            shadowsocks.shell.print_exception(e)
         self.server_info = server_info
 
     def client_encode(self, buf):
-        if self.sub_obfs is not None:
-            return self.sub_obfs.client_encode(buf)
         return buf
 
     def client_decode(self, buf):
-        if self.sub_obfs is not None:
-            return self.sub_obfs.client_decode(buf)
         return (buf, False)
 
     def server_encode(self, buf):
-        if self.sub_obfs is not None:
-            return self.sub_obfs.server_encode(buf)
         return buf
 
     def server_decode(self, buf):
-        if self.sub_obfs is not None:
-            return self.sub_obfs.server_decode(buf)
         return (buf, True, False)
 
     def get_head_size(self, buf, def_value):
@@ -347,7 +319,6 @@ class client_queue(object):
 
 class obfs_auth_data(object):
     def __init__(self):
-        self.sub_obfs = None
         self.client_id = {}
         self.startup_time = int(time.time() - 30) & 0xFFFFFFFF
         self.local_client_id = b''
diff --git a/shadowsocks/shell.py b/shadowsocks/shell.py
index 66e38bb..38d2432 100644
--- a/shadowsocks/shell.py
+++ b/shadowsocks/shell.py
@@ -218,6 +218,7 @@ def get_config(is_local):
 
     config['password'] = to_bytes(config.get('password', b''))
     config['method'] = to_str(config.get('method', 'aes-256-cfb'))
+    config['protocol'] = to_str(config.get('protocol', 'origin'))
     config['obfs'] = to_str(config.get('obfs', 'plain'))
     config['obfs_param'] = to_str(config.get('obfs_param', ''))
     config['port_password'] = config.get('port_password', None)
diff --git a/shadowsocks/tcprelay.py b/shadowsocks/tcprelay.py
index e1ea954..395ecbf 100644
--- a/shadowsocks/tcprelay.py
+++ b/shadowsocks/tcprelay.py
@@ -122,6 +122,14 @@ class TCPRelayHandler(object):
         server_info.param = config['obfs_param']
         self._obfs.set_server_info(server_info)
 
+        self._protocol = obfs.obfs(config['protocol'])
+        server_info = obfs.server_info(server.protocol_data)
+        server_info.host = config['server']
+        server_info.port = server._listen_port
+        server_info.tcp_mss = 1440
+        server_info.param = ''
+        self._protocol.set_server_info(server_info)
+
         self._fastopen_connected = False
         self._data_to_write_to_local = []
         self._data_to_write_to_remote = []
@@ -330,7 +338,7 @@ class TCPRelayHandler(object):
 
     def _handle_stage_connecting(self, data):
         if self._is_local:
-            data = self._obfs.client_pre_encrypt(data)
+            data = self._protocol.client_pre_encrypt(data)
             data = self._encryptor.encrypt(data)
             data = self._obfs.client_encode(data)
         if data:
@@ -428,7 +436,7 @@ class TCPRelayHandler(object):
                     data = b'\x88' + struct.pack('>H', total_len) + chr(rnd_len) + (b' ' * (rnd_len - 1)) + data
                     crc = (0xffffffff - binascii.crc32(data)) & 0xffffffff
                     data += struct.pack('<I', crc)
-                data = self._obfs.client_pre_encrypt(data)
+                data = self._protocol.client_pre_encrypt(data)
                 data_to_send = self._encryptor.encrypt(data)
                 data_to_send = self._obfs.client_encode(data_to_send)
                 if data_to_send:
@@ -574,7 +582,7 @@ class TCPRelayHandler(object):
                 else:
                     data = obfs_decode[0]
                 try:
-                    data = self._obfs.server_post_decrypt(data)
+                    data = self._protocol.server_post_decrypt(data)
                 except Exception as e:
                     shell.print_exception(e)
                     self.destroy()
@@ -583,7 +591,7 @@ class TCPRelayHandler(object):
         self._server.server_transfer_ul += len(data)
         if self._stage == STAGE_STREAM:
             if self._is_local:
-                data = self._obfs.client_pre_encrypt(data)
+                data = self._protocol.client_pre_encrypt(data)
                 data = self._encryptor.encrypt(data)
                 data = self._obfs.client_encode(data)
             self._write_to_sock(data, self._remote_sock)
@@ -634,10 +642,10 @@ class TCPRelayHandler(object):
                 send_back = self._obfs.client_encode(b'')
                 self._write_to_sock(send_back, self._remote_sock)
             data = self._encryptor.decrypt(obfs_decode[0])
-            data = self._obfs.client_post_decrypt(data)
+            data = self._protocol.client_post_decrypt(data)
         else:
             if self._encrypt_correct:
-                data = self._obfs.server_pre_encrypt(data)
+                data = self._protocol.server_pre_encrypt(data)
                 data = self._encryptor.encrypt(data)
         try:
             self._write_to_sock(data, self._local_sock)
@@ -756,6 +764,12 @@ class TCPRelayHandler(object):
             del self._fd_to_handlers[self._local_sock.fileno()]
             self._local_sock.close()
             self._local_sock = None
+        if self._obfs:
+            self._obfs.dispose()
+            self._obfs = None
+        if self._protocol:
+            self._protocol.dispose()
+            self._protocol = None
         self._dns_resolver.remove_callback(self._handle_dns_resolved)
         self._server.remove_handler(self)
         self._server.add_connection(-1)
@@ -771,6 +785,7 @@ class TCPRelay(object):
         self.server_transfer_ul = 0
         self.server_transfer_dl = 0
         self.server_connections = 0
+        self.protocol_data = obfs.obfs(config['protocol']).init_data()
         self.obfs_data = obfs.obfs(config['obfs']).init_data()
 
         self._timeout = config['timeout']