From 48e241d99f3c7b65b802e98c7d002b4ba3761baa Mon Sep 17 00:00:00 2001
From: ImBenji <benjamin.watt@imbenji.net>
Date: Tue, 5 Aug 2025 13:56:03 +0100
Subject: [PATCH] Enhance traffic control by dynamically detecting outgoing
 interface and updating iptables rules for upload and download traffic

---
 lib/wireguard/traffic_control.dart | 70 ++++++++++++++++++++++++------
 1 file changed, 57 insertions(+), 13 deletions(-)

diff --git a/lib/wireguard/traffic_control.dart b/lib/wireguard/traffic_control.dart
index f88fa46..13d3780 100644
--- a/lib/wireguard/traffic_control.dart
+++ b/lib/wireguard/traffic_control.dart
@@ -34,12 +34,12 @@ class TrafficControlService {
       final downloadMark = mark + 1000; // Offset to avoid conflicts
       
       print('Running iptables MARK commands for $peerIP...');
-      // Clean existing rules for this peer first
+      // Clean existing mangle rules for this peer first
       try {
-        await _runIptablesCommand(['-D', 'FORWARD', '-s', peerIP, '-j', 'MARK', '--set-mark', uploadMark.toString()]);
+        await _runIptablesCommand(['-t', 'mangle', '-D', 'POSTROUTING', '-s', peerIP, '-j', 'MARK', '--set-mark', uploadMark.toString()]);
       } catch (e) { /* Rule doesn't exist, ignore */ }
       try {
-        await _runIptablesCommand(['-D', 'FORWARD', '-d', peerIP, '-j', 'MARK', '--set-mark', downloadMark.toString()]);
+        await _runIptablesCommand(['-t', 'mangle', '-D', 'FORWARD', '-d', peerIP, '-j', 'MARK', '--set-mark', downloadMark.toString()]);
       } catch (e) { /* Rule doesn't exist, ignore */ }
       
       // Mark upload traffic (FROM peer) with uploadMark - use POSTROUTING for upload
@@ -47,21 +47,47 @@ class TrafficControlService {
       // Mark download traffic (TO peer) with downloadMark - use FORWARD for download
       await _runIptablesCommand(['-t', 'mangle', '-I', 'FORWARD', '-d', peerIP, '-j', 'MARK', '--set-mark', downloadMark.toString()]);
       
-      print('Running tc class add/change commands for upload and download...');
+      // Detect outgoing interface dynamically
+      print('Detecting outgoing interface...');
+      final outgoingInterface = await _detectOutgoingInterface();
+      print('Setting up upload limiting on outgoing interface: $outgoingInterface');
       
-      // Upload class (traffic FROM peer)
+      // Setup HTB on outgoing interface for upload limiting
       try {
-        await _runTcCommand(['class', 'add', 'dev', 'wg0', 'parent', '1:1', 'classid', '1:$uploadMark', 'htb', 'rate', '${kbps}kbit', 'ceil', '${kbps}kbit']);
+        await _runTcCommand(['qdisc', 'add', 'dev', outgoingInterface, 'root', 'handle', '2:', 'htb', 'default', '30']);
       } catch (e) {
-        if (e.toString().contains('File exists')) {
-          print('Upload class exists, updating...');
-          await _runTcCommand(['class', 'change', 'dev', 'wg0', 'classid', '1:$uploadMark', 'htb', 'rate', '${kbps}kbit', 'ceil', '${kbps}kbit']);
+        if (e.toString().contains('File exists') || e.toString().contains('Exclusivity flag')) {
+          print('HTB qdisc already exists on $outgoingInterface, continuing...');
         } else {
           rethrow;
         }
       }
       
-      // Download class (traffic TO peer)
+      try {
+        await _runTcCommand(['class', 'add', 'dev', outgoingInterface, 'parent', '2:', 'classid', '2:1', 'htb', 'rate', '1000mbit']);
+      } catch (e) {
+        if (e.toString().contains('File exists')) {
+          print('HTB root class already exists on $outgoingInterface, continuing...');
+        } else {
+          rethrow;
+        }
+      }
+      
+      print('Running tc class add/change commands for upload and download...');
+      
+      // Upload class on outgoing interface (traffic FROM peer going out)
+      try {
+        await _runTcCommand(['class', 'add', 'dev', outgoingInterface, 'parent', '2:1', 'classid', '2:$uploadMark', 'htb', 'rate', '${kbps}kbit', 'ceil', '${kbps}kbit']);
+      } catch (e) {
+        if (e.toString().contains('File exists')) {
+          print('Upload class exists, updating...');
+          await _runTcCommand(['class', 'change', 'dev', outgoingInterface, 'classid', '2:$uploadMark', 'htb', 'rate', '${kbps}kbit', 'ceil', '${kbps}kbit']);
+        } else {
+          rethrow;
+        }
+      }
+      
+      // Download class on wg0 (traffic TO peer)
       try {
         await _runTcCommand(['class', 'add', 'dev', 'wg0', 'parent', '1:1', 'classid', '1:$downloadMark', 'htb', 'rate', '${kbps}kbit', 'ceil', '${kbps}kbit']);
       } catch (e) {
@@ -74,14 +100,14 @@ class TrafficControlService {
       }
       
       print('Running tc filter add commands...');
-      // Upload filter
+      // Upload filter on outgoing interface
       try {
-        await _runTcCommand(['filter', 'add', 'dev', 'wg0', 'protocol', 'ip', 'parent', '1:', 'prio', '1', 'handle', uploadMark.toString(), 'fw', 'flowid', '1:$uploadMark']);
+        await _runTcCommand(['filter', 'add', 'dev', outgoingInterface, 'protocol', 'ip', 'parent', '2:', 'prio', '1', 'handle', uploadMark.toString(), 'fw', 'flowid', '2:$uploadMark']);
       } catch (e) {
         if (!e.toString().contains('File exists')) rethrow;
       }
       
-      // Download filter
+      // Download filter on wg0
       try {
         await _runTcCommand(['filter', 'add', 'dev', 'wg0', 'protocol', 'ip', 'parent', '1:', 'prio', '1', 'handle', downloadMark.toString(), 'fw', 'flowid', '1:$downloadMark']);
       } catch (e) {
@@ -113,6 +139,24 @@ class TrafficControlService {
     return lastOctet + 10;
   }
 
+  static Future<String> _detectOutgoingInterface() async {
+    try {
+      // Get default route interface
+      final result = await Process.run('ip', ['route', 'show', 'default']);
+      if (result.exitCode == 0) {
+        final output = result.stdout.toString();
+        final match = RegExp(r'dev (\w+)').firstMatch(output);
+        if (match != null) {
+          return match.group(1)!;
+        }
+      }
+    } catch (e) {
+      print('Failed to detect outgoing interface: $e');
+    }
+    // Fallback to eth0
+    return 'eth0';
+  }
+
   static Future<void> _runIptablesCommand(List<String> args) async {
     final command = 'iptables ${args.join(' ')}';
     print('Executing: $command');