Enhance traffic control by dynamically detecting outgoing interface and updating iptables rules for upload and download traffic

lib/wireguard/traffic_control.dart

@@ -3,6 +3,14 @@ import 'dart:io';
 class TrafficControlService {
   static Future<void> setSpeedLimit(String peerIP, int bytesPerSecond) async {
     final mark = _getMarkForIP(peerIP);
+    
+    // Handle unlimited speed (-1)
+    if (bytesPerSecond == -1) {
+      print('Removing speed limit for peer $peerIP (unlimited)');
+      await _removeSpeedLimit(peerIP, mark);
+      return;
+    }
+    
     final kbps = (bytesPerSecond * 8 / 1000).round(); // Convert bytes/s to kbps
     print('Setting speed limit for peer $peerIP to ${bytesPerSecond} bytes/s (${kbps}kbps, mark: $mark)');
 
@@ -125,8 +133,18 @@ class TrafficControlService {
   }
 
   static Future<void> setDataCap(String peerIP, int quotaBytes) async {
+    // Handle unlimited data cap (-1)
+    if (quotaBytes == -1) {
+      print('Removing data cap for peer $peerIP (unlimited)');
+      await _removeDataCap(peerIP);
+      return;
+    }
+    
     print('Setting data cap for peer $peerIP to $quotaBytes bytes');
     
+    // Remove existing data cap rules first
+    await _removeDataCap(peerIP);
+    
     await _runIptablesCommand(['-I', 'FORWARD', '-s', peerIP, '-m', 'quota', '--quota', quotaBytes.toString(), '-j', 'ACCEPT']);
     await _runIptablesCommand(['-I', 'FORWARD', '-d', peerIP, '-m', 'quota', '--quota', quotaBytes.toString(), '-j', 'ACCEPT']);
     await _runIptablesCommand(['-A', 'FORWARD', '-s', peerIP, '-j', 'DROP']);
@@ -157,6 +175,74 @@ class TrafficControlService {
     return 'eth0';
   }
 
+  static Future<void> _removeSpeedLimit(String peerIP, int mark) async {
+    final uploadMark = mark;
+    final downloadMark = mark + 1000;
+    
+    print('Removing speed limit rules for $peerIP...');
+    
+    // Remove iptables mangle rules
+    try {
+      await _runIptablesCommand(['-t', 'mangle', '-D', 'POSTROUTING', '-s', peerIP, '-j', 'MARK', '--set-mark', uploadMark.toString()]);
+    } catch (e) { /* Rule doesn't exist, ignore */ }
+    try {
+      await _runIptablesCommand(['-t', 'mangle', '-D', 'FORWARD', '-d', peerIP, '-j', 'MARK', '--set-mark', downloadMark.toString()]);
+    } catch (e) { /* Rule doesn't exist, ignore */ }
+    
+    // Remove tc classes and filters
+    final outgoingInterface = await _detectOutgoingInterface();
+    
+    try {
+      await _runTcCommand(['filter', 'del', 'dev', outgoingInterface, 'protocol', 'ip', 'parent', '2:', 'prio', '1', 'handle', uploadMark.toString(), 'fw']);
+    } catch (e) { /* Filter doesn't exist, ignore */ }
+    try {
+      await _runTcCommand(['filter', 'del', 'dev', 'wg0', 'protocol', 'ip', 'parent', '1:', 'prio', '1', 'handle', downloadMark.toString(), 'fw']);
+    } catch (e) { /* Filter doesn't exist, ignore */ }
+    
+    try {
+      await _runTcCommand(['class', 'del', 'dev', outgoingInterface, 'classid', '2:$uploadMark']);
+    } catch (e) { /* Class doesn't exist, ignore */ }
+    try {
+      await _runTcCommand(['class', 'del', 'dev', 'wg0', 'classid', '1:$downloadMark']);
+    } catch (e) { /* Class doesn't exist, ignore */ }
+    
+    print('Speed limit removed for $peerIP');
+  }
+
+  static Future<void> _removeDataCap(String peerIP) async {
+    print('Removing data cap rules for $peerIP...');
+    
+    // Remove quota rules (these are harder to target specifically, so we try common patterns)
+    try {
+      await _runIptablesCommand(['-D', 'FORWARD', '-s', peerIP, '-j', 'DROP']);
+    } catch (e) { /* Rule doesn't exist, ignore */ }
+    try {
+      await _runIptablesCommand(['-D', 'FORWARD', '-d', peerIP, '-j', 'DROP']);
+    } catch (e) { /* Rule doesn't exist, ignore */ }
+    
+    // Remove quota rules by flushing and re-adding basic rules
+    // This is a bit aggressive but necessary since quota rules are hard to target specifically
+    try {
+      final result = await Process.run('iptables', ['-S', 'FORWARD']);
+      if (result.exitCode == 0) {
+        final lines = result.stdout.toString().split('\n');
+        for (final line in lines) {
+          if (line.contains(peerIP) && line.contains('quota')) {
+            // Extract the rule and convert -A to -D to delete it
+            final deleteRule = line.replaceFirst('-A', '-D');
+            try {
+              await Process.run('iptables', deleteRule.split(' ').skip(1).toList());
+            } catch (e) { /* Ignore individual failures */ }
+          }
+        }
+      }
+    } catch (e) {
+      print('Failed to remove specific quota rules: $e');
+    }
+    
+    print('Data cap removed for $peerIP');
+  }
+
   static Future<void> _runIptablesCommand(List<String> args) async {
     final command = 'iptables ${args.join(' ')}';
     print('Executing: $command');