From 82f8fef28aa300ea3ad3e09a4742da5108b98e68 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E7=A0=B4=E5=A8=83=E9=85=B1?= <mmgac001@gmail.com>
Date: Sun, 4 Jun 2017 18:34:12 +0800
Subject: [PATCH] allow set speed limit in runtime

---
 db_transfer.py          |  2 +-
 mujson_mgr.py           |  6 +++++-
 server_pool.py          |  3 +--
 shadowsocks/tcprelay.py | 31 +++++++++++++++++++++++++------
 shadowsocks/udprelay.py | 15 ++++++++-------
 5 files changed, 40 insertions(+), 17 deletions(-)

diff --git a/db_transfer.py b/db_transfer.py
index 2f10870..e758381 100644
--- a/db_transfer.py
+++ b/db_transfer.py
@@ -123,7 +123,7 @@ class TransferBase(object):
 				continue
 
 			if allow:
-				allow_users[port] = passwd
+				allow_users[port] = cfg
 				if 'protocol' in cfg and 'protocol_param' in cfg and common.to_str(cfg['protocol']) in obfs.mu_protocol():
 					if '#' in common.to_str(cfg['protocol_param']):
 						mu_servers[port] = passwd
diff --git a/mujson_mgr.py b/mujson_mgr.py
index 39e6b6e..47e91a6 100644
--- a/mujson_mgr.py
+++ b/mujson_mgr.py
@@ -239,7 +239,7 @@ General options:
 
 
 def main():
-	shortopts = 'adeclu:i:p:k:O:o:G:g:m:t:f:h'
+	shortopts = 'adeclu:i:p:k:O:o:G:g:m:t:f:hs:S:'
 	longopts = ['help']
 	action = None
 	user = {}
@@ -302,6 +302,10 @@ def main():
 				user['obfs_param'] = value
 			elif key == '-G':
 				user['protocol_param'] = value
+			elif key == '-s':
+				user['speed_limit_per_con'] = int(value)
+			elif key == '-S':
+				user['speed_limit_per_user'] = int(value)
 			elif key == '-m':
 				if value in fast_set_method:
 					user['method'] = fast_set_method[value]
diff --git a/server_pool.py b/server_pool.py
index 2e6a85b..d159817 100644
--- a/server_pool.py
+++ b/server_pool.py
@@ -28,6 +28,7 @@ import time
 from shadowsocks import shell, eventloop, tcprelay, udprelay, asyncdns, common
 import threading
 import sys
+import traceback
 from socket import *
 from configloader import load_config, get_config
 
@@ -80,12 +81,10 @@ class ServerPool(object):
 			loop.run()
 		except (KeyboardInterrupt, IOError, OSError) as e:
 			logging.error(e)
-			import traceback
 			traceback.print_exc()
 			os.exit(0)
 		except Exception as e:
 			logging.error(e)
-			import traceback
 			traceback.print_exc()
 
 	def server_is_run(self, port):
diff --git a/shadowsocks/tcprelay.py b/shadowsocks/tcprelay.py
index 3322b52..7aae7e4 100644
--- a/shadowsocks/tcprelay.py
+++ b/shadowsocks/tcprelay.py
@@ -273,6 +273,11 @@ class TCPRelayHandler(object):
     def _update_user(self, user):
         self._user = user
         self._user_id = struct.unpack('<I', user)[0]
+        if self._user in self._server.server_users_cfg:
+            cfg = self._server.server_users_cfg[self._user]
+            speed = cfg.get('speed_limit_per_con', 0)
+            self.speed_tester_u.update_limit(speed)
+            self.speed_tester_d.update_limit(speed)
 
     def _update_activity(self, data_len=0):
         # tell the TCP Relay we have activities recently
@@ -1168,6 +1173,7 @@ class TCPRelay(object):
         self.server_transfer_ul = 0
         self.server_transfer_dl = 0
         self.server_users = {}
+        self.server_users_cfg = {}
         self.server_user_transfer_ul = {}
         self.server_user_transfer_dl = {}
         self.mu = False
@@ -1258,9 +1264,9 @@ class TCPRelay(object):
                             self.del_user(uid)
                         else:
                             passwd = items[1]
-                            self.add_user(uid, passwd)
+                            self.add_user(uid, {'password':passwd})
 
-    def update_user(self, id, passwd):
+    def _update_user(self, id, passwd):
         uid = struct.pack('<I', id)
         self.add_user(uid, passwd)
 
@@ -1273,12 +1279,25 @@ class TCPRelay(object):
             uid = struct.pack('<I', id)
             self.add_user(uid, users[id])
 
-    def add_user(self, user, passwd): # user: binstr[4], passwd: str
-        self.server_users[user] = common.to_bytes(passwd)
+    def add_user(self, uid, cfg): # user: binstr[4], passwd: str
+        passwd = cfg['password']
+        self.server_users[uid] = common.to_bytes(passwd)
+        self.server_users_cfg[uid] = cfg
+        speed = cfg.get("speed_limit_per_user", 0)
+        if uid in self._speed_tester_u:
+            self._speed_tester_u[uid].update_limit(speed)
+        else:
+            self._speed_tester_u[uid] = SpeedTester(speed)
+        if uid in self._speed_tester_d:
+            self._speed_tester_d[uid].update_limit(speed)
+        else:
+            self._speed_tester_d[uid] = SpeedTester(speed)
 
-    def del_user(self, user):
+    def del_user(self, uid):
         if user in self.server_users:
-            del self.server_users[user]
+            del self.server_users[uid]
+        if user in self.server_users_cfg:
+            del self.server_users_cfg[uid]
 
     def add_transfer_u(self, user, transfer):
         if user is None:
diff --git a/shadowsocks/udprelay.py b/shadowsocks/udprelay.py
index a7bbf3c..b9606cd 100644
--- a/shadowsocks/udprelay.py
+++ b/shadowsocks/udprelay.py
@@ -249,9 +249,9 @@ class UDPRelay(object):
                             self.del_user(uid)
                         else:
                             passwd = items[1]
-                            self.add_user(uid, passwd)
+                            self.add_user(uid, {'password':passwd})
 
-    def update_user(self, id, passwd):
+    def _update_user(self, id, passwd):
         uid = struct.pack('<I', id)
         self.add_user(uid, passwd)
 
@@ -264,12 +264,13 @@ class UDPRelay(object):
             uid = struct.pack('<I', id)
             self.add_user(uid, users[id])
 
-    def add_user(self, user, passwd): # user: binstr[4], passwd: str
-        self.server_users[user] = common.to_bytes(passwd)
+    def add_user(self, uid, cfg): # user: binstr[4], passwd: str
+        passwd = cfg['password']
+        self.server_users[uid] = common.to_bytes(passwd)
 
-    def del_user(self, user):
-        if user in self.server_users:
-            del self.server_users[user]
+    def del_user(self, uid):
+        if uid in self.server_users:
+            del self.server_users[uid]
 
     def add_transfer_u(self, user, transfer):
         if user is None: