Files
waylume_server/lib/wireguard/traffic_control.dart

277 lines
11 KiB
Dart

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)');
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 mangle rules for this peer first
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 */ }
// 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()]);
// Detect outgoing interface dynamically
print('Detecting outgoing interface...');
final outgoingInterface = await _detectOutgoingInterface();
print('Setting up upload limiting on outgoing interface: $outgoingInterface');
// Setup HTB on outgoing interface for upload limiting
try {
await _runTcCommand(['qdisc', 'add', 'dev', outgoingInterface, 'root', 'handle', '2:', 'htb', 'default', '30']);
} catch (e) {
if (e.toString().contains('File exists') || e.toString().contains('Exclusivity flag')) {
print('HTB qdisc already exists on $outgoingInterface, continuing...');
} else {
rethrow;
}
}
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) {
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 on outgoing interface
try {
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 on wg0
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 {
// 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']);
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<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> _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');
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}');
}
}
}