allow set speed limit in runtime

Author: breakwa11

db_transfer.py

@@ -123,7 +123,7 @@ def del_server_out_of_bound_safe(self, last_rows, rows):
 				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

mujson_mgr.py

@@ -239,7 +239,7 @@ def print_server_help():
 
 
 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]

server_pool.py

@@ -28,6 +28,7 @@
 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 @@ def _loop(loop, dns_resolver, mgr):
 			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):

shadowsocks/tcprelay.py

@@ -273,6 +273,11 @@ def _create_encryptor(self, config):
     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 @@ def __init__(self, config, dns_resolver, is_local, stat_callback=None, stat_coun
         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 @@ def _update_users(self, protocol_param, acl):
                             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 @@ def update_users(self, users):
             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:

shadowsocks/udprelay.py

@@ -249,9 +249,9 @@ def _update_users(self, protocol_param, acl):
                             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 @@ def update_users(self, users):
             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: