Files
waylume_server/lib/wireguard/traffic_control.dart

97 lines
3.8 KiB
Dart

import 'dart:io';
class TrafficControlService {
static Future<void> setSpeedLimit(String peerIP, int speedKbps) async {
final mark = _getMarkForIP(peerIP);
print('Setting speed limit for peer $peerIP to ${speedKbps}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;
}
}
print('Running iptables MARK commands for $peerIP...');
await _runIptablesCommand(['-I', 'FORWARD', '-s', peerIP, '-j', 'MARK', '--set-mark', mark.toString()]);
await _runIptablesCommand(['-I', 'FORWARD', '-d', peerIP, '-j', 'MARK', '--set-mark', mark.toString()]);
print('Running tc class add command...');
await _runTcCommand(['class', 'add', 'dev', 'wg0', 'parent', '1:1', 'classid', '1:$mark', 'htb', 'rate', '${speedKbps}kbit', 'ceil', '${speedKbps}kbit']);
print('Running tc filter add command...');
await _runTcCommand(['filter', 'add', 'dev', 'wg0', 'protocol', 'ip', 'parent', '1:', 'prio', '1', 'handle', mark.toString(), 'fw', 'flowid', '1:$mark']);
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 dataCapMB) async {
final quotaBytes = dataCapMB * 1024 * 1024;
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}');
}
}
}