wlan-ap/feeds/ucentral/ratelimit/files/usr/bin/ratelimit

#!/usr/bin/env ucode
'use strict';

import { basename, popen } from 'fs';
import * as ubus from 'ubus';
import * as uloop from 'uloop';

let defaults = {};
let devices = {};

function cmd(command, ignore_error) {
//	if (ignore_error)
//		command += "> /dev/null 2>&1";
	warn(`> ${command}\n`);
	let rc = system(command);
	return ignore_error || rc == 0;
}

function qdisc_add_leaf(iface, id, opts) {
	opts ??= "";
	return cmd(`tc class replace dev ${iface} parent 1:1 classid 1:${id} htb rate 1mbit ${opts} burst 2k prio 1`) &&
	       cmd(`tc qdisc replace dev ${iface} parent 1:${id} handle ${id}: fq_codel flows 128 limit 800 quantum 300 noecn`);
}

function qdisc_del_leaf(iface, id) {
	cmd(`tc class del dev ${iface} parent 1:1 classid 1:${id}`, true);
}

function qdisc_add(iface) {
	return cmd(`tc qdisc add dev ${iface} root handle 1: htb default 2`) &&
	       cmd(`tc class add dev ${iface} parent 1: classid 1:1 htb rate 1000mbit burst 6k`) &&
	       qdisc_add_leaf(iface, 2, "ceil 1000mbit");
}

function qdisc_del(iface) {
	cmd(`tc qdisc del dev ${iface} root`, true);
}

function ifb_dev(iface) {
	return "i-" + iface;
}

function ifb_add(iface, ifbdev) {
	return cmd(`ip link add ${ifbdev} type ifb`) &&
	       cmd(`ip link set ${ifbdev} up`) &&
	       cmd(`tc qdisc add dev ${iface} clsact`, true) &&
	       cmd(`tc filter add dev ${iface} ingress protocol all prio 512 matchall action mirred egress redirect dev ${ifbdev}`);
}

function ifb_del(iface, ifbdev) {
	cmd(`tc filter del dev ${iface} ingress protocol all prio 512`);
	cmd(`ip link set ${ifbdev} down`, true);
	cmd(`ip link del ${ifbdev}`, true);
}

function macfilter_add(iface, id, type, mac) {
	return cmd(`tc filter add dev ${iface} protocol all parent 1: prio 1 handle 800::${id} u32 match ether ${type} ${mac} flowid 1:${id}`);
}

function macfilter_del(iface, id) {
	cmd(`tc filter del dev ${iface} protocol all parent 1: prio 1 handle 800::${id} u32`, true);
}

function linux_client_del(device, client) {
	printf('-> linux_client_del\n');
	let ifbdev = ifb_dev(device.name);
	let id = client.id + 3;

	macfilter_del(device.name, id);
	qdisc_del_leaf(device.name, id);
	macfilter_del(ifbdev, id);
	qdisc_del_leaf(ifbdev, id);
}

function linux_client_set(device, client) {
	printf('-> linux_client_set\n');
	let ifbdev = ifb_dev(device.name);
	let id = client.id + 3;

	linux_client_del(device, client);

	let ret = qdisc_add_leaf(device.name, id, `ceil ${client.data.rate_egress}`) &&
		  macfilter_add(device.name, id, "dst", client.address) &&
		  qdisc_add_leaf(ifbdev, id, `ceil ${client.data.rate_ingress}`) &&
		  macfilter_add(ifbdev, id, "src", client.address);

	if (!ret)
		linux_client_del(device, client);

	return ret;
}


let ops = {
	device: {
		add: function(name) {
			printf('-> device.add\n');
			let ifbdev = ifb_dev(name);

			qdisc_del(name);
			ifb_del(name, ifbdev);

			let ret = qdisc_add(name) &&
				  ifb_add(name, ifbdev) &&
				  qdisc_add(ifbdev);

			if (!ret) {
				qdisc_del(name);
				ifb_del(name, ifbdev);
			}

			return ret;
		},
		remove: function(name) {
			printf('-> device.remove\n');
			let ifbdev = ifb_dev(name);
			qdisc_del(name);
			ifb_del(name, ifbdev);
		}
	},
	client: {
		set: function(device, client) {
			printf('-> client.set\n');
			return linux_client_set(device, client);
		},
		remove: function(device, client) {
			printf('-> client.remove\n');
			linux_client_del(device, client);
		}
	}
};

function get_device(devices, name) {
	printf('-> get_device\n');
	let device = devices[name];

	if (device)
		return device;

	if (!ops.device.add(name))
		return null;

	device = {
		name: name,
		clients: {},
		client_order: [],
	};

	devices[name] = device;

	return device;
}

function del_device(name) {
	printf('-> del_device\n');
	if (!devices[name])
		return;
	ops.device.remove(name);
	delete devices[name];
}

function get_free_idx(list) {
	printf('-> get_free_idx\n');
	for (let i = 0; i < length(list); i++)
		if (list[i] == null)
			return i;

	return length(list);
}

function del_client(device, address) {
	printf('-> del_client\n');

	let client = device.clients[address];

	if (!client)
		return false;

	delete device.clients[address];
	device.client_order[client.id] = null;

	ops.client.remove(device, client);
	return true;
}

function get_client(device, address) {
	printf('-> get_client\n');
	let client = device.clients[address];

	if (client)
		return client;

	let i = get_free_idx(device.client_order);
	client = {};
	client.address = address;
	client.id = i;
	client.data = {};

	device.clients[address] = client;
	device.client_order[i] = client;

	return client;
}

function set_client(device, client, data) {
	printf('-> set_client\n');
	let update = false;

	for (let key in data) {
		if (client.data[key] != data[key])
			update = true;

		client.data[key] = data[key];
	}

	if (update && !ops.client.set(device, client)) {
		del_client(device, client.address);
		return false;
	}

	return true;
}

function run_service() {
	let uctx = ubus.connect();

	uctx.publish("ratelimit", {
		flush: {
			call: function(req) {
				printf('-> flush\n');
				defaults = {};
			},
			args: {

			}
		},
		defaults_set: {
			call: function(req) {
				printf('-> defaults_set\n');
				let r_i = req.args.rate_ingress ?? req.args.rate;
				let r_e = req.args.rate_egress ?? req.args.rate;
				let name = req.args.name;

				if (!name || !r_i || !r_e)
					return ubus.STATUS_INVALID_ARGUMENT;

				defaults[name] = [ r_e, r_i ];

				return 0;
			},
			args: {
				name:"",
				rate:"",
				rate_ingress:"",
				rate_egress:"",
			}
		},
		client_set: {
			call: function(req) {
				printf('-> client_set\n');
				let r_i = req.args.rate_ingress ?? req.args.rate;
				let r_e = req.args.rate_egress ?? req.args.rate;

				if (req.args.defaults && defaults[req.args.defaults]) {
					let def = defaults[req.args.defaults];

					r_e ??= def[0];
					r_i ??= def[1];
				}

				if (!req.args.device || !req.args.address || !r_i || !r_e)
					return ubus.STATUS_INVALID_ARGUMENT;

				let device = get_device(devices, req.args.device);
				if (!device)
					return ubus.STATUS_INVALID_ARGUMENT;

				let client = get_client(device, req.args.address);
				if (!client)
					return ubus.STATUS_INVALID_ARGUMENT;

				let data = {
					rate_ingress: r_i,
					rate_egress: r_e
				};

				if (!set_client(device, client, data))
					return ubus.STATUS_UNKNOWN_ERROR;

				return 0;
			},
			args: {
				device:"",
				defaults:"",
				address:"",
				rate:"",
				rate_ingress:"",
				rate_egress:"",
			}
		},
		client_delete: {
			call: function(req) {
				printf('-> client_delete\n');
				if (!req.args.address)
					return ubus.STATUS_INVALID_ARGUMENT;

				if (req.args.device) {
					let device = devices[req.args.device];
					if (!device)
						return ubus.STATUS_NOT_FOUND;

					if (!del_client(device, req.args.address))
						return ubus.STATUS_NOT_FOUND;
				} else {
					for (let dev in devices) {
						let device = devices[dev];

						del_client(device, req.args.address);
					}
				}

				return 0;
			},
			args: {
				device:"",
				address:"",
			}
		},
		device_delete: {
			call: function(req) {
				printf('-> device_delete\n');
				let name = req.args.device;

				if (!name)
					return ubus.STATUS_INVALID_ARGUMENT;

				if (!devices[name])
					return ubus.STATUS_NOT_FOUND;

				del_device(name);

				return 0;
			},
			args: {
				device:"",
			}
		},
		reload: {
			call: function(req) {
				printf('-> reload\n');
				let list = uctx.list();
				for (let obj in list) {
					if (!wildcard(obj, 'hostapd.wlan*') && !wildcard(obj, 'hostapd.phy*'))
						continue;
					let iface = split(obj, '.')[1];
					let device = get_device(devices, req.args.device);
					if (!device)
						continue;
					let status = uctx.call(obj, 'get_status');
					if (!status?.ssid)
						continue;
					if (!defaults[status?.ssid])
						continue;
					let data = {
						rate_ingress: defaults[status?.ssid][0],
						rate_egress: defaults[status?.ssid][1]
					};
					for (let k, client in device.clients)
						set_client(device, client, data);
				}
				return 0;
			},
			args: {
			}
		},
		dump: {
			call: function(req) {
				return { devices, defaults };
			},
			args: {}
		}
	});

	try {
		uloop.run();
	} catch (e) {
		warn(`Error: ${e}\n${e.stacktrace[0].context}`);
	}

	for (let dev in devices) {
		del_device(dev);
	}
}

uloop.init();
run_service();
uloop.done();