waylume_server/lib/wireguard/traffic_control.dart

import 'dart:io';

class TrafficControlService {
  static Future<void> setSpeedLimit(String peerIP, int bytesPerSecond) async {
    final mark = _getMarkForIP(peerIP);
    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)');

    try {
      // Ensure HTB qdisc exists on wg0
      print('Setting up HTB qdisc on wg0...');
      try {
        await _runTcCommand(['qdisc', 'add', 'dev', 'wg0', 'root', 'handle', '1:', 'htb', 'default', '30']);
      } catch (e) {
        if (e.toString().contains('File exists') || e.toString().contains('Exclusivity flag')) {
          print('HTB qdisc already exists, continuing...');
        } else {
          rethrow;
        }
      }
      
      try {
        await _runTcCommand(['class', 'add', 'dev', 'wg0', 'parent', '1:', 'classid', '1:1', 'htb', 'rate', '1000mbit']);
      } catch (e) {
        if (e.toString().contains('File exists')) {
          print('HTB root class already exists, continuing...');
        } else {
          rethrow;
        }
      }
      
      // Create separate classes for upload and download
      final uploadMark = mark;
      final downloadMark = mark + 1000; // Offset to avoid conflicts
      
      print('Running iptables MARK commands for $peerIP...');
      // Clean existing rules for this peer first
      try {
        await _runIptablesCommand(['-D', 'FORWARD', '-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()]);
      } catch (e) { /* Rule doesn't exist, ignore */ }
      
      // Mark upload traffic (FROM peer) with uploadMark - use POSTROUTING for upload
      await _runIptablesCommand(['-t', 'mangle', '-I', 'POSTROUTING', '-s', peerIP, '-j', 'MARK', '--set-mark', uploadMark.toString()]);
      // 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...');
      
      // Upload class (traffic FROM peer)
      try {
        await _runTcCommand(['class', 'add', 'dev', 'wg0', 'parent', '1:1', 'classid', '1:$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', 'wg0', 'classid', '1:$uploadMark', 'htb', 'rate', '${kbps}kbit', 'ceil', '${kbps}kbit']);
        } else {
          rethrow;
        }
      }
      
      // Download class (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) {
        if (e.toString().contains('File exists')) {
          print('Download class exists, updating...');
          await _runTcCommand(['class', 'change', 'dev', 'wg0', 'classid', '1:$downloadMark', 'htb', 'rate', '${kbps}kbit', 'ceil', '${kbps}kbit']);
        } else {
          rethrow;
        }
      }
      
      print('Running tc filter add commands...');
      // Upload filter
      try {
        await _runTcCommand(['filter', 'add', 'dev', 'wg0', 'protocol', 'ip', 'parent', '1:', 'prio', '1', 'handle', uploadMark.toString(), 'fw', 'flowid', '1:$uploadMark']);
      } catch (e) {
        if (!e.toString().contains('File exists')) rethrow;
      }
      
      // Download filter
      try {
        await _runTcCommand(['filter', 'add', 'dev', 'wg0', 'protocol', 'ip', 'parent', '1:', 'prio', '1', 'handle', downloadMark.toString(), 'fw', 'flowid', '1:$downloadMark']);
      } catch (e) {
        if (!e.toString().contains('File exists')) rethrow;
      }
      
      print('Speed limit set successfully for $peerIP');
    } catch (e) {
      print('ERROR setting speed limit for $peerIP: $e');
      // Don't rethrow if it's just because qdisc already exists
      if (!e.toString().contains('File exists')) {
        rethrow;
      }
    }
  }

  static Future<void> setDataCap(String peerIP, int quotaBytes) async {
    print('Setting data cap for peer $peerIP to $quotaBytes bytes');
    
    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']);
    await _runIptablesCommand(['-A', 'FORWARD', '-d', peerIP, '-j', 'DROP']);
  }

  static int _getMarkForIP(String ip) {
    // Use last octet of IP + 10 to ensure small, valid class IDs (10-255)
    final lastOctet = int.parse(ip.split('.').last);
    return lastOctet + 10;
  }

  static Future<void> _runIptablesCommand(List<String> args) async {
    final command = 'iptables ${args.join(' ')}';
    print('Executing: $command');
    final result = await Process.run('iptables', args);
    print('iptables exit code: ${result.exitCode}');
    if (result.stdout.toString().isNotEmpty) {
      print('iptables stdout: ${result.stdout}');
    }
    if (result.stderr.toString().isNotEmpty) {
      print('iptables stderr: ${result.stderr}');
    }
    if (result.exitCode != 0) {
      throw Exception('iptables command failed: ${result.stderr}');
    }
  }

  static Future<void> _runTcCommand(List<String> args) async {
    final command = 'tc ${args.join(' ')}';
    print('Executing: $command');
    final result = await Process.run('tc', args);
    print('tc exit code: ${result.exitCode}');
    if (result.stdout.toString().isNotEmpty) {
      print('tc stdout: ${result.stdout}');
    }
    if (result.stderr.toString().isNotEmpty) {
      print('tc stderr: ${result.stderr}');
    }
    if (result.exitCode != 0) {
      throw Exception('tc command failed: ${result.stderr}');
    }
  }
}