Implemented traffic limits support.

broker/limits.py

@@ -0,0 +1,59 @@
+import logging
+import struct
+
+from . import protocol, traffic_control
+
+# Logger.
+logger = logging.getLogger("tunneldigger.limits")
+
+
+class LimitManager(object):
+    """
+    Tunnel traffic limit manager.
+    """
+
+    def __init__(self, tunnel, session_id):
+        """
+        Class constructor.
+
+        :param tunnel: Tunnel instance
+        :parma session_id: Session identifier
+        """
+
+        self.tunnel = tunnel
+        self.session_id = session_id
+
+    def configure(self, limit_message):
+        """
+        Configures a specific limit.
+
+        :param limit_message: Received limit control message
+        """
+
+        try:
+            limit_type, config_len = struct.unpack('!BB', limit_message[:2])
+        except ValueError:
+            logger.warning("Malformed limit configuration received on tunnel %d." % self.tunnel.tunnel_id)
+            return False
+
+        if limit_type == protocol.LIMIT_TYPE_BANDWIDTH_DOWN:
+            # Downstream (client-wise) limit setup.
+            try:
+                bandwidth = struct.unpack('!I', limit_message[2:2 + config_len])[0]
+            except ValueError:
+                logger.warning("Malformed bandwidth limit configuration received on tunnel %d." % self.tunnel.tunnel_id)
+                return False
+
+            logger.info("Setting downstream bandwidth limit to %d kbps on tunnel %d." % (bandwidth, self.tunnel.tunnel_id))
+
+            # Setup bandwidth limit using Linux traffic shaping.
+            try:
+                tc = traffic_control.TrafficControl(self.tunnel.get_session_name(self.session_id))
+                tc.reset()
+                tc.set_fixed_bandwidth(bandwidth)
+            except traffic_control.TrafficControlError:
+                logger.warning("Unable to configure traffic shaping for tunnel %d." % self.tunnel.tunnel_id)
+
+            return True
+        else:
+            return False

broker/network.py

@@ -139,6 +139,11 @@ def write_message(self, address, msg_type, msg_data=''):
         data += struct.pack('!BB', msg_type, len(msg_data))
         data += msg_data
 
+        # Pad the message to be at least 12 bytes long, as otherwise some firewalls
+        # may filter it when used over port 53.
+        if len(data) < 12:
+            data += '\x00' * (12 - len(data))
+
         self.write(address, data)
 
     def read(self, file_object):
@@ -162,7 +167,7 @@ def read(self, file_object):
             raise
         except:
             logger.error("Unhandled exception during message processing.")
-            logger.debug(traceback.format_exc())
+            logger.error(traceback.format_exc())
 
     def message(self, address, msg_type, msg_data, raw_length):
         """

broker/traffic_control.py

@@ -1,7 +1,5 @@
 import os
 
-TC = '/sbin/tc'
-
 
 class TrafficControlError(Exception):
     pass
@@ -17,20 +15,20 @@ def __init__(self, interface):
 
         self.interface = interface
 
-    def tc(self, command, ignore_fails = False):
+    def tc(self, command, ignore_fails=False):
         """
         Executes a traffic control command.
         """
 
-        if os.system('%s %s' % (TC, command)) != 0 and not ignore_fails:
+        if os.system('tc %s' % (command)) != 0 and not ignore_fails:
             raise TrafficControlError
 
     def reset(self):
         """
         Clears all existing traffic control rules.
         """
 
-        self.tc('qdisc del dev %s root handle 1: htb default 0' % self.interface, ignore_fails = True)
+        self.tc('qdisc del dev %s root handle 1: htb default 0' % self.interface, ignore_fails=True)
         self.tc('qdisc add dev %s root handle 1: htb default 1' % self.interface)
 
     def set_fixed_bandwidth(self, bandwidth):

broker/tunnel.py

@@ -8,7 +8,7 @@
 import struct
 import time
 
-from . import l2tp, protocol, network
+from . import l2tp, protocol, network, limits
 
 # Socket options.
 IP_MTU_DISCOVER = 10
@@ -358,6 +358,7 @@ def message(self, address, msg_type, msg_data, raw_length):
         elif msg_type == protocol.CONTROL_TYPE_PMTUD:
             # The other side is performing PMTU discovery.
             self.write_message(self.endpoint, protocol.CONTROL_TYPE_PMTUD_ACK, struct.pack('!H', raw_length))
+            return True
         elif msg_type == protocol.CONTROL_TYPE_PMTUD_ACK:
             # The other side is acknowledging a specific PMTU value.
             pmtu = struct.unpack('!H', msg_data)[0] + IPV4_HDR_OVERHEAD
@@ -368,19 +369,25 @@ def message(self, address, msg_type, msg_data, raw_length):
 
                 # Notify the other side of our measured MTU.
                 self.write_message(self.endpoint, protocol.CONTROL_TYPE_PMTU_NTFY, struct.pack('!H', self.measured_pmtu))
+
+            return True
         elif msg_type == protocol.CONTROL_TYPE_PMTU_NTFY:
             # The other side is notifying us about their tunnel MTU.
             remote_mtu = struct.unpack('!H', msg_data)[0]
 
             if remote_mtu != self.remote_tunnel_mtu:
                 self.remote_tunnel_mtu = remote_mtu
                 self.update_mtu()
+
+            return True
         elif msg_type & protocol.MASK_CONTROL_TYPE_RELIABLE:
             # Acknowledge reliable control messages.
             self.write_message(self.endpoint, protocol.CONTROL_TYPE_REL_ACK, msg_data[:2])
 
             if msg_type == protocol.CONTROL_TYPE_LIMIT:
-                # TODO: Client requests limit configuration.
+                # Client requests limit configuration.
+                limit_manager = limits.LimitManager(self, 1)
+                limit_manager.configure(msg_data[2:])
                 return True
 
         return False