From cb1dcacd4b4f200648de059d44d25e2d454b0389 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E7=A0=B4=E5=A8=83=E9=85=B1?= <mmgac001@gmail.com>
Date: Wed, 26 Jul 2017 12:13:26 +0800
Subject: [PATCH] refine recv size try more data_size_list

---
 shadowsocks/obfsplugin/auth_chain.py | 13 ++++++++++++-
 shadowsocks/tcprelay.py              |  6 +++---
 2 files changed, 15 insertions(+), 4 deletions(-)

diff --git a/shadowsocks/obfsplugin/auth_chain.py b/shadowsocks/obfsplugin/auth_chain.py
index b39624f..3949008 100644
--- a/shadowsocks/obfsplugin/auth_chain.py
+++ b/shadowsocks/obfsplugin/auth_chain.py
@@ -640,16 +640,22 @@ class auth_chain_b(auth_chain_a):
         self.salt = b"auth_chain_b"
         self.no_compatible_method = 'auth_chain_b'
         self.data_size_list = []
+        self.data_size_list2 = []
 
     def init_data_size(self, key):
         if self.data_size_list:
             self.data_size_list = []
+            self.data_size_list2 = []
         random = xorshift128plus()
         random.init_from_bin(key)
-        list_len = random.next() % 16 + 4;
+        list_len = random.next() % 16 + 4
         for i in range(0, list_len):
             self.data_size_list.append((int)(random.next() % 1440))
         self.data_size_list.sort()
+        list_len = random.next() % 32 + 32
+        for i in range(0, list_len):
+            self.data_size_list2.append((int)(random.next() % 1440))
+        self.data_size_list2.sort()
 
     def set_server_info(self, server_info):
         self.server_info = server_info
@@ -669,6 +675,11 @@ class auth_chain_b(auth_chain_a):
         if final_pos < len(self.data_size_list):
             return self.data_size_list[final_pos] - buf_size - self.server_info.overhead
 
+        pos = bisect.bisect_left(self.data_size_list2, buf_size + self.server_info.overhead)
+        final_pos = pos + random.next() % (len(self.data_size_list2) + 1 - pos)
+        if final_pos < len(self.data_size_list2):
+            return self.data_size_list2[final_pos] - buf_size - self.server_info.overhead
+
         if buf_size > 1300:
             return random.next() % 31
         if buf_size > 900:
diff --git a/shadowsocks/tcprelay.py b/shadowsocks/tcprelay.py
index 78b3257..595e2be 100644
--- a/shadowsocks/tcprelay.py
+++ b/shadowsocks/tcprelay.py
@@ -816,15 +816,15 @@ class TCPRelayHandler(object):
         if self._overhead == 0:
             return recv_buffer_size
         buffer_size = len(sock.recv(recv_buffer_size, socket.MSG_PEEK))
+        frame_size = self._tcp_mss - self._overhead
         if up:
             buffer_size = min(buffer_size, self._recv_u_max_size)
-            self._recv_u_max_size = min(self._recv_u_max_size + self._tcp_mss - self._overhead, BUF_SIZE)
+            self._recv_u_max_size = min(self._recv_u_max_size + frame_size, BUF_SIZE)
         else:
             buffer_size = min(buffer_size, self._recv_d_max_size)
-            self._recv_d_max_size = min(self._recv_d_max_size + self._tcp_mss - self._overhead, BUF_SIZE)
+            self._recv_d_max_size = min(self._recv_d_max_size + frame_size, BUF_SIZE)
         if buffer_size == recv_buffer_size:
             return buffer_size
-        frame_size = self._tcp_mss - self._overhead
         if buffer_size > frame_size:
             buffer_size = int(buffer_size / frame_size) * frame_size
         return buffer_size