From f61bbefa95ee43c485b44f0fc0e1a8af92e2a452 Mon Sep 17 00:00:00 2001
From: AkaneAkaza <AkaneAkaza@protonmail.com>
Date: Mon, 23 Oct 2017 14:18:58 +0800
Subject: [PATCH] Add cache flag

---
 shadowsocks/encrypt.py               | 19 ++++++++++++-------
 shadowsocks/obfsplugin/auth_chain.py |  7 ++-----
 shadowsocks/tcprelay.py              |  2 +-
 3 files changed, 15 insertions(+), 13 deletions(-)

diff --git a/shadowsocks/encrypt.py b/shadowsocks/encrypt.py
index 29a9e6d..a9fd02a 100644
--- a/shadowsocks/encrypt.py
+++ b/shadowsocks/encrypt.py
@@ -46,7 +46,7 @@ def try_cipher(key, method=None):
     Encryptor(key, method)
 
 
-def EVP_BytesToKey(password, key_len, iv_len):
+def EVP_BytesToKey(password, key_len, iv_len, cache):
     # equivalent to OpenSSL's EVP_BytesToKey() with count 1
     # so that we make the same key and iv as nodejs version
     cached_key = '%s-%d-%d' % (password, key_len, iv_len)
@@ -66,13 +66,14 @@ def EVP_BytesToKey(password, key_len, iv_len):
     ms = b''.join(m)
     key = ms[:key_len]
     iv = ms[key_len:key_len + iv_len]
-    cached_keys[cached_key] = (key, iv)
-    cached_keys.sweep()
+    if cache:
+        cached_keys[cached_key] = (key, iv)
+        cached_keys.sweep()
     return key, iv
 
 
 class Encryptor(object):
-    def __init__(self, key, method, iv = None):
+    def __init__(self, key, method, iv = None, cache = False):
         self.key = key
         self.method = method
         self.iv = None
@@ -81,6 +82,7 @@ class Encryptor(object):
         self.iv_buf = b''
         self.cipher_key = b''
         self.decipher = None
+        self.cache = cache
         method = method.lower()
         self._method_info = self.get_method_info(method)
         if self._method_info:
@@ -105,7 +107,7 @@ class Encryptor(object):
         password = common.to_bytes(password)
         m = self._method_info
         if m[0] > 0:
-            key, iv_ = EVP_BytesToKey(password, m[0], m[1])
+            key, iv_ = EVP_BytesToKey(password, m[0], m[1], self.cache)
         else:
             # key_length == 0 indicates we should use the key directly
             key, iv = password, b''
@@ -119,6 +121,9 @@ class Encryptor(object):
 
     def encrypt(self, buf):
         if len(buf) == 0:
+            if not self.iv_sent:
+                self.iv_sent = True
+                return self.cipher_iv
             return buf
         if self.iv_sent:
             return self.cipher.update(buf)
@@ -155,7 +160,7 @@ def encrypt_all(password, method, op, data):
     method = method.lower()
     (key_len, iv_len, m) = method_supported[method]
     if key_len > 0:
-        key, _ = EVP_BytesToKey(password, key_len, iv_len)
+        key, _ = EVP_BytesToKey(password, key_len, iv_len, True)
     else:
         key = password
     if op:
@@ -172,7 +177,7 @@ def encrypt_key(password, method):
     method = method.lower()
     (key_len, iv_len, m) = method_supported[method]
     if key_len > 0:
-        key, _ = EVP_BytesToKey(password, key_len, iv_len)
+        key, _ = EVP_BytesToKey(password, key_len, iv_len, True)
     else:
         key = password
     return key
diff --git a/shadowsocks/obfsplugin/auth_chain.py b/shadowsocks/obfsplugin/auth_chain.py
index 6619f4f..2860655 100644
--- a/shadowsocks/obfsplugin/auth_chain.py
+++ b/shadowsocks/obfsplugin/auth_chain.py
@@ -349,7 +349,6 @@ class auth_chain_a(auth_base):
     def pack_client_data(self, buf):
         buf = self.encryptor.encrypt(buf)
         data = self.rnd_data(len(buf), buf, self.last_client_hash, self.random_client)
-        data_len = len(data) + 8
         mac_key = self.user_key + struct.pack('<I', self.pack_id)
         length = len(buf) ^ struct.unpack('<H', self.last_client_hash[14:])[0]
         data = struct.pack('<H', length) + data
@@ -361,7 +360,6 @@ class auth_chain_a(auth_base):
     def pack_server_data(self, buf):
         buf = self.encryptor.encrypt(buf)
         data = self.rnd_data(len(buf), buf, self.last_server_hash, self.random_server)
-        data_len = len(data) + 8
         mac_key = self.user_key + struct.pack('<I', self.pack_id)
         length = len(buf) ^ struct.unpack('<H', self.last_server_hash[14:])[0]
         data = struct.pack('<H', length) + data
@@ -372,7 +370,6 @@ class auth_chain_a(auth_base):
 
     def pack_auth_data(self, auth_data, buf):
         data = auth_data
-        data_len = 12 + 4 + 16 + 4
         data = data + (struct.pack('<H', self.server_info.overhead) + struct.pack('<H', 0))
         mac_key = self.server_info.iv + self.server_info.key
 
@@ -565,7 +562,7 @@ class auth_chain_a(auth_base):
             if length >= 4096:
                 self.raw_trans = True
                 self.recv_buf = b''
-                if self.recv_id == 0:
+                if self.recv_id == 1:
                     logging.info(self.no_compatible_method + ': over size')
                     return (b'E' * 2048, False)
                 else:
@@ -581,7 +578,7 @@ class auth_chain_a(auth_base):
                 ))
                 self.raw_trans = True
                 self.recv_buf = b''
-                if self.recv_id == 0:
+                if self.recv_id == 1:
                     return (b'E' * 2048, False)
                 else:
                     raise Exception('server_post_decrype data uncorrect checksum')
diff --git a/shadowsocks/tcprelay.py b/shadowsocks/tcprelay.py
index ce84d84..e4b8fbe 100644
--- a/shadowsocks/tcprelay.py
+++ b/shadowsocks/tcprelay.py
@@ -266,7 +266,7 @@ class TCPRelayHandler(object):
     def _create_encryptor(self, config):
         try:
             self._encryptor = encrypt.Encryptor(config['password'],
-                                                config['method'])
+                                                config['method'], None, True)
             return True
         except Exception:
             self._stage = STAGE_DESTROYED