#!/bin/sh
# conn-limits-gen: translate the UCI 'conn_limits' config into an nftables
# ruleset and apply it. One 'config port' section per L2 device. Coexists with
# fw3 (own tables only).
#
# Per-port port security, all three enforced even on bridged LAN ports:
#   - connection rate  (new conns/sec per source IP)
#   - session limit    (concurrent conns per source IP)
#   - inbound policer  (inbound kbit/s, drop-on-exceed)
#
# Inbound policer runs at the physical netdev ingress hook. Connection rate/session limit need conntrack (only
# available at inet prerouting, where a bridged LAN port reads as br-lan), so
# each port's ingress hook tags packets with a reserved mark nibble and the
# inet rules match that mark -> true per-physical-port through the bridge.
#
# Flow offloading: the inbound policer needs to see every packet at the
# ingress hook, so flow offloading (software AND hardware) must be disabled
# for it to enforce correctly - offloaded flows bypass the hook. Connection
# rate and session limits match 'ct state new' (always slow path) and work
# with offloading enabled.
#
# Usage: conn-limits-gen [apply|flush]

. /lib/functions.sh

NFT=/usr/sbin/nft
INET_TABLE=conn_limits
SIG_FILE=/var/run/conn_limits.sig   # signature of the last-applied ruleset

# Reserved skb-mark bits for per-port matching (top nibble). OR-set at ingress,
# matched at prerouting, cleared afterwards so routing/mwan3 marks are untouched.
MARK_MASK=0x0f000000
MARK_CLEAR=0xf0ffffff

RATE_UNIT=second
# inbound_rate is given in kbit/s. nft's limit rate uses byte units only, so it
# is converted to bytes/s: bytes/s = kbit/s * 1000 / 8 = kbit/s * 125.
KBIT_TO_BYTES=125

inet_pre=""
netdev_tables=""
netdev_devs=""
have_inet=0
port_idx=0

_sanitize() { printf '%s' "$1" | sed 's/[^a-zA-Z0-9]/_/g'; }

handle_port() {
	local cfg="$1"
	local enabled dev
	local rate limit inbound_rate

	config_get_bool enabled "$cfg" enabled 0
	[ "$enabled" = 1 ] || return 0

	# 'name' is the L2 device (netdev).
	config_get dev "$cfg" name
	[ -n "$dev" ] || return 0
	[ -e "/sys/class/net/${dev}" ] || return 0

	config_get rate "$cfg" rate
	config_get limit "$cfg" limit
	config_get inbound_rate "$cfg" inbound_rate

	# nothing to enforce?
	[ -n "$rate$limit$inbound_rate" ] || return 0

	local need_mark=0
	[ -n "$rate" ] || [ -n "$limit" ] && need_mark=1

	# stable per-port mark id (1..15) in the reserved nibble
	port_idx=$((port_idx + 1))
	[ "$port_idx" -gt 15 ] && need_mark=0     # only 15 mark ids available
	local markval; markval=$(printf '0x0%x000000' "$port_idx")

	local dev_id; dev_id=$(_sanitize "$dev")

	# --- netdev ingress table (Inbound policer + per-port mark tag) ---
	local nd="\tchain ingress {\n\t\ttype filter hook ingress device \"${dev}\" priority 0; policy accept;\n"
	if [ -n "$inbound_rate" ]; then
		local ib=$(( inbound_rate * KBIT_TO_BYTES ))   # kbit/s -> bytes/s
		nd="${nd}\t\tlimit rate over ${ib} bytes/second burst ${ib} bytes counter drop\n"
	fi
	[ "$need_mark" = 1 ] && \
		nd="${nd}\t\tmeta mark set meta mark or ${markval}\n"
	nd="${nd}\t}"
	netdev_tables="${netdev_tables}add table netdev conn_limits_${dev_id}\ndelete table netdev conn_limits_${dev_id}\ntable netdev conn_limits_${dev_id} {\n${nd}\n}\n"
	netdev_devs="${netdev_devs} ${dev_id}"

	# --- inet prerouting rules (Connection rate/session limit), matched by the port mark ---
	[ "$need_mark" = 1 ] || return 0
	local fam sa r
	for fam in 4 6; do
		[ "$fam" = 4 ] && sa="ip saddr" || sa="ip6 saddr"
		if [ -n "$rate" ]; then
			r="meta mark & ${MARK_MASK} == ${markval} ct state new meter r_${dev_id}${fam} { ${sa} limit rate over ${rate}/${RATE_UNIT} burst ${rate} packets } counter drop"
			inet_pre="${inet_pre}\t\t${r}\n"
			have_inet=1
		fi
		if [ -n "$limit" ]; then
			r="meta mark & ${MARK_MASK} == ${markval} ct state new meter c_${dev_id}${fam} { ${sa} ct count over ${limit} } counter drop"
			inet_pre="${inet_pre}\t\t${r}\n"
			have_inet=1
		fi
	done
}

# live netdev tables owned by us (one per port), as "conn_limits_<dev_id>"
live_netdev_tables() {
	"$NFT" list tables netdev 2>/dev/null | \
		awk '$1 == "table" && index($3, "conn_limits_") == 1 { print $3 }'
}

build_ruleset() {
	# drop tables of ports that are no longer configured/enabled, so that
	# their policers stop enforcing (same transaction, stays atomic)
	local t
	for t in $(live_netdev_tables); do
		case " ${netdev_devs} " in
			*" ${t#conn_limits_} "*) ;;
			*) echo "delete table netdev ${t}" ;;
		esac
	done
	echo "add table inet ${INET_TABLE}"
	echo "delete table inet ${INET_TABLE}"
	if [ "$have_inet" = 1 ]; then
		echo "table inet ${INET_TABLE} {"
		echo "	chain prerouting {"
		echo "		type filter hook prerouting priority filter; policy accept;"
		printf "%b" "$inet_pre"
		# clear our reserved mark bits so routing/mwan3 marks stay intact
		echo "		meta mark set meta mark & ${MARK_CLEAR}"
		echo "	}"
		echo "}"
	fi
	printf "%b" "$netdev_tables"
}

# true if everything the last apply created is still present in the kernel;
# guards the signature skip against external 'nft flush ruleset' & friends
ruleset_live() {
	if [ "$have_inet" = 1 ]; then
		"$NFT" list table inet ${INET_TABLE} >/dev/null 2>&1 || return 1
	fi
	local d
	for d in $netdev_devs; do
		"$NFT" list table netdev "conn_limits_${d}" >/dev/null 2>&1 || return 1
	done
	return 0
}

do_apply() {
	config_load conn_limits
	config_foreach handle_port port

	local tmp; tmp="$(mktemp)"
	build_ruleset > "$tmp"

	# Skip reapply when the ruleset is identical to the last applied one, so that
	# rebuild-on-interface-change (hotplug) does not wipe counters/meters/buckets.
	# Only if the rules are actually still loaded - the signature alone proves
	# what was generated, not that it is still in the kernel.
	local cur; cur="$(md5sum "$tmp" | cut -d' ' -f1)"
	if [ -f "$SIG_FILE" ] && [ "$(cat "$SIG_FILE")" = "$cur" ] && ruleset_live; then
		rm -f "$tmp"
		return 0
	fi

	if "$NFT" -c -f "$tmp" && "$NFT" -f "$tmp"; then
		echo "$cur" > "$SIG_FILE"
		rm -f "$tmp"
		return 0
	fi
	# on failure keep the previous ruleset; save the rejected one for inspection
	cp "$tmp" /tmp/conn_limits.failed.nft
	rm -f "$tmp"
	return 1
}

do_flush() {
	"$NFT" delete table inet ${INET_TABLE} 2>/dev/null
	# delete every live table of ours, not the configured set - catches ports
	# that were removed/disabled since the last apply
	local t
	for t in $(live_netdev_tables); do
		"$NFT" delete table netdev "$t" 2>/dev/null
	done
	rm -f "$SIG_FILE"          # force a real apply next time (tables are gone)
	return 0
}

case "${1:-apply}" in
	apply) do_apply ;;
	flush) do_flush ;;
	*) echo "usage: $0 [apply|flush]" >&2; exit 2 ;;
esac
