set local network shaper for test ton-node-rust

set_shaper.sh

#!/usr/bin/env bash
# shapertest.sh — cross-platform network shaper for protocol testing
# Usage: ./shapertest.sh <set|reset|status> <node_count> <delay_ms> <loss_percent>
# Supports: Linux (tc/netem), FreeBSD (ipfw/dummynet), macOS (dnctl/pf)
# Ports per node: TCP 3000+N, UDP 3000+N, UDP 4000+N

set -euo pipefail

# ─── Argument parsing ────────────────────────────────────────────────────────

usage() {
    echo "Usage: $0 <set|reset|status> [node_count] [delay_ms] [loss_percent]"
    echo ""
    echo "  set    <node_count> <delay_ms> <loss_percent>  — apply shaping rules"
    echo "  reset                                          — remove shaping rules"
    echo "  status                                         — show rules and counters"
    echo ""
    echo "  node_count   — number of nodes (ports 3001..300N, 4001..400N)"
    echo "  delay_ms     — one-way delay in milliseconds (e.g. 25)"
    echo "  loss_percent — packet loss 0-100 (e.g. 10)"
    echo ""
    echo "Examples:"
    echo "  $0 set 6 25 10"
    echo "  $0 reset"
    echo "  $0 status"
    exit 1
}

[[ $# -lt 1 ]] && usage
ACTION="$1"

case "$ACTION" in
    set)
        [[ $# -lt 4 ]] && { echo "Error: 'set' requires node_count, delay_ms, loss_percent"; usage; }
        NODE_COUNT="$2"
        DELAY_MS="$3"
        LOSS_PCT="$4"
        ;;
    reset|status)
        NODE_COUNT=0
        DELAY_MS=0
        LOSS_PCT=0
        ;;
    *)
        echo "Error: action must be 'set', 'reset', or 'status'"
        usage
        ;;
esac

# ─── OS detection ────────────────────────────────────────────────────────────

detect_os() {
    case "$(uname -s)" in
        Linux)   echo "linux" ;;
        FreeBSD) echo "freebsd" ;;
        Darwin)  echo "macos" ;;
        *)
            echo "Error: unsupported OS: $(uname -s)" >&2
            exit 1
            ;;
    esac
}

OS=$(detect_os)
echo "==> Detected OS: $OS"

# ─── Port list builder ───────────────────────────────────────────────────────

build_port_lists() {
    TCP_PORTS=()
    UDP_PORTS=()
    for (( i=1; i<=NODE_COUNT; i++ )); do
        TCP_PORTS+=("$((3000 + i))")
        UDP_PORTS+=("$((3000 + i))")
        UDP_PORTS+=("$((4000 + i))")
    done
}

# ─── Linux ───────────────────────────────────────────────────────────────────

linux_set() {
    local iface="lo"

    echo "==> [Linux] Clearing existing qdiscs on $iface"
    tc qdisc del dev "$iface" root 2>/dev/null || true

    echo "==> [Linux] Adding HTB root qdisc"
    tc qdisc add dev "$iface" root handle 1: htb default 99
    tc class add dev "$iface" parent 1: classid 1:10 htb rate 1gbit
    tc class add dev "$iface" parent 1: classid 1:99 htb rate 1gbit

    echo "==> [Linux] Adding netem: delay=${DELAY_MS}ms loss=${LOSS_PCT}%"
    tc qdisc add dev "$iface" parent 1:10 handle 10: \
        netem delay "${DELAY_MS}ms" loss "${LOSS_PCT}%"

    echo "==> [Linux] Adding port filters"
    for port in "${TCP_PORTS[@]}"; do
        for dir in dport sport; do
            tc filter add dev "$iface" parent 1: protocol ip prio 1 u32 \
                match ip protocol 6 0xff \
                match ip "$dir" "$port" 0xffff \
                flowid 1:10
        done
    done

    for port in "${UDP_PORTS[@]}"; do
        for dir in dport sport; do
            tc filter add dev "$iface" parent 1: protocol ip prio 1 u32 \
                match ip protocol 17 0xff \
                match ip "$dir" "$port" 0xffff \
                flowid 1:10
        done
    done

    echo ""
    linux_status
}

linux_reset() {
    echo "==> [Linux] Removing qdiscs on lo"
    tc qdisc del dev lo root 2>/dev/null \
        && echo "    Removed." \
        || echo "    Nothing to remove."
}

linux_status() {
    local iface="lo"

    # Check if our qdisc is present at all
    if ! tc qdisc show dev "$iface" | grep -q "htb"; then
        echo "==> [Linux] No shaping rules active on $iface"
        return
    fi

    echo "==> [Linux] Qdisc tree with counters:"
    echo ""

    # netem stats: shows delay/loss config + sent/dropped counts
    tc -s qdisc show dev "$iface" | awk '
        /qdisc netem/ {
            print "  [netem] " $0
            getline; print "  " $0   # config line (delay, loss)
            getline; print "  " $0   # Sent X bytes Y pkts
            getline; print "  " $0   # dropped
            print ""
        }
        /qdisc htb/ {
            print "  [htb]   " $0
            getline; print "  " $0
            print ""
        }
    '

    echo "==> [Linux] Class counters:"
    echo ""
    tc -s class show dev "$iface" | awk '
        /class htb 1:10/ {
            print "  [shaped traffic — class 1:10]"
            # Print next 6 lines (rate, sent, dropped, etc.)
            for (i=0; i<6; i++) { getline; print "    " $0 }
            print ""
        }
        /class htb 1:99/ {
            print "  [unmatched traffic — class 1:99]"
            for (i=0; i<6; i++) { getline; print "    " $0 }
            print ""
        }
    '

    echo "==> [Linux] Filter hit counts:"
    echo ""
    # u32 filters with -s show match counts per filter
    tc -s filter show dev "$iface" | awk '
        /filter u32/ { filter_line = $0 }
        /match ip (d|s)port/ {
            # Extract port number and direction from match line
            split($0, a, " ")
            for (i=1; i<=length(a); i++) {
                if (a[i] == "dport" || a[i] == "sport") {
                    dir = a[i]; port = a[i+1]
                }
            }
        }
        /Sent/ {
            # Format: Sent X bytes Y pkts
            printf "  port %-6s %-5s  %s\n", port, dir, $0
        }
    '
}

# ─── FreeBSD ─────────────────────────────────────────────────────────────────

freebsd_set() {
    local plr
    plr=$(echo "scale=4; $LOSS_PCT / 100" | bc)

    echo "==> [FreeBSD] Configuring ipfw pipe 100: delay=${DELAY_MS}ms plr=${plr}"
    ipfw pipe 100 config delay "${DELAY_MS}" plr "$plr" bw 1Gbit/s

    echo "==> [FreeBSD] Adding ipfw rules (base rule number 1000)"
    local rule=1000

    for port in "${TCP_PORTS[@]}"; do
        ipfw add "$rule" pipe 100 tcp from any to any dst-port "$port" via lo0
        (( rule++ ))
        ipfw add "$rule" pipe 100 tcp from any "$port" to any via lo0
        (( rule++ ))
    done

    for port in "${UDP_PORTS[@]}"; do
        ipfw add "$rule" pipe 100 udp from any to any dst-port "$port" via lo0
        (( rule++ ))
        ipfw add "$rule" pipe 100 udp from any "$port" to any via lo0
        (( rule++ ))
    done

    echo ""
    freebsd_status
}

freebsd_reset() {
    echo "==> [FreeBSD] Removing ipfw rules 1000-1999"
    local rules
    rules=$(ipfw list 2>/dev/null | awk '$1 >= 1000 && $1 < 2000 {print $1}')
    if [[ -n "$rules" ]]; then
        echo "$rules" | xargs -n1 ipfw delete 2>/dev/null || true
        echo "    Rules removed."
    else
        echo "    No rules in range 1000-1999."
    fi

    echo "==> [FreeBSD] Deleting pipe 100"
    ipfw pipe 100 delete 2>/dev/null \
        && echo "    Pipe removed." \
        || echo "    No pipe to remove."
}

freebsd_status() {
    local has_pipe has_rules

    has_pipe=$(ipfw pipe list 2>/dev/null | grep -c "^00100" || true)
    has_rules=$(ipfw list 2>/dev/null | awk '$1 >= 1000 && $1 < 2000' | wc -l | tr -d ' ')

    if [[ "$has_pipe" -eq 0 && "$has_rules" -eq 0 ]]; then
        echo "==> [FreeBSD] No shaping rules active"
        return
    fi

    if [[ "$has_pipe" -gt 0 ]]; then
        echo "==> [FreeBSD] Pipe 100 configuration and counters:"
        echo ""
        # pipe show outputs: config line + per-bucket flow stats
        ipfw pipe 100 show 2>/dev/null | awk '
            NR==1 { print "  Config : " $0; next }
            /pkts/ { print "  Flows  : " $0; next }
            /^[[:space:]]/ { print "  " $0 }
        '
        echo ""
    fi

    if [[ "$has_rules" -gt 0 ]]; then
        echo "==> [FreeBSD] Rules 1000-1999 with packet/byte counters:"
        echo ""
        # -c flag shows packet/byte counts per rule
        ipfw -ac list 2>/dev/null \
            | awk '$1 >= 1000 && $1 < 2000 {
                printf "  rule %-5s  pkts=%-8s bytes=%-10s  %s %s %s %s %s %s %s\n",
                    $1, $2, $3, $4, $5, $6, $7, $8, $9, $10
              }'
        echo ""
    fi
}

# ─── macOS ───────────────────────────────────────────────────────────────────

MACOS_ANCHOR="shapertest"
MACOS_ANCHOR_FILE="/etc/pf.anchors/${MACOS_ANCHOR}"

macos_set() {
    local plr
    plr=$(echo "scale=4; $LOSS_PCT / 100" | bc)

    echo "==> [macOS] Configuring dnctl pipe 1: delay=${DELAY_MS}ms plr=${plr} bw=unlimited"
    # bw 0 = unlimited — avoids the silent 1bit/s fallback from unsupported 'G' prefix
    dnctl pipe 1 config delay "${DELAY_MS}" plr "$plr" bw 0

    echo "==> [macOS] Writing pf anchor file: $MACOS_ANCHOR_FILE"
    mkdir -p /etc/pf.anchors

    local tcp_set udp_set
    tcp_set="{ $(IFS=', '; echo "${TCP_PORTS[*]}") }"
    udp_set="{ $(IFS=', '; echo "${UDP_PORTS[*]}") }"

    cat > "$MACOS_ANCHOR_FILE" <<EOF
# shapertest anchor — auto-generated by shapertest.sh, safe to delete
dummynet in  quick on lo0 proto tcp from any port ${tcp_set} to any pipe 1
dummynet out quick on lo0 proto tcp from any to any port ${tcp_set} pipe 1
dummynet in  quick on lo0 proto udp from any port ${udp_set} to any pipe 1
dummynet out quick on lo0 proto udp from any to any port ${udp_set} pipe 1
EOF

    echo "==> [macOS] Registering dummynet-anchor + anchor in main pf ruleset"
    # Must pipe through /etc/pf.conf content + anchor declarations together.
    # 'dummynet-anchor' registers the anchor for dummynet, 'anchor' for filter rules.
    (
        cat /etc/pf.conf
        echo "dummynet-anchor \"${MACOS_ANCHOR}\""
        echo "anchor \"${MACOS_ANCHOR}\""
    ) | pfctl -f - 2>&1 | grep -v "^$" || true

    echo "==> [macOS] Loading rules into anchor"
    pfctl -a "$MACOS_ANCHOR" -f "$MACOS_ANCHOR_FILE"

    echo "==> [macOS] Enabling pf"
    pfctl -E 2>/dev/null || true

    echo ""
    macos_status
}

macos_reset() {
    echo "==> [macOS] Flushing anchor rules"
    pfctl -a "$MACOS_ANCHOR" -F all 2>/dev/null || true
    rm -f "$MACOS_ANCHOR_FILE"

    echo "==> [macOS] Deleting dnctl pipe 1"
    dnctl pipe 1 delete 2>/dev/null \
        && echo "    Pipe removed." \
        || echo "    No pipe to remove."

    echo "==> [macOS] Reloading pf without anchor"
    pfctl -f /etc/pf.conf 2>&1 | grep -v "^$" || true
}

macos_status() {
    local pipe_active=0 anchor_active=0

    dnctl pipe 1 show 2>/dev/null | grep -q "^00001" && pipe_active=1 || true
    pfctl -a "$MACOS_ANCHOR" -s rules 2>/dev/null | grep -q "dummynet" && anchor_active=1 || true

    if [[ "$pipe_active" -eq 0 && "$anchor_active" -eq 0 ]]; then
        echo "==> [macOS] No shaping rules active"
        return
    fi

    if [[ "$pipe_active" -eq 1 ]]; then
        echo "==> [macOS] dnctl pipe 1 configuration and counters:"
        echo ""
        dnctl pipe 1 show 2>/dev/null | awk '
            NR==1 { print "  Config : " $0; next }
            /pkts/  { print "  Flows  : " $0; next }
            /^[[:space:]]/ { print "  " $0 }
        '
        echo ""
    fi

    if [[ "$anchor_active" -eq 1 ]]; then
        echo "==> [macOS] pf anchor rules with counters:"
        echo ""
        pfctl -a "$MACOS_ANCHOR" -v -s rules 2>/dev/null | awk '
            /dummynet/ {
                rule = $0
                getline
                match($0, /Packets: ([0-9]+).*Bytes: ([0-9]+)/, m)
                printf "  pkts=%-8s bytes=%-12s  %s\n", m[1], m[2], rule
            }
        '
        echo ""
    fi
}

# ─── Dispatch ────────────────────────────────────────────────────────────────

if [[ "$ACTION" == "set" ]]; then
    if [[ "$NODE_COUNT" -lt 1 ]]; then
        echo "Error: node_count must be >= 1 for 'set'"
        usage
    fi
    build_port_lists
    echo "==> TCP ports: ${TCP_PORTS[*]}"
    echo "==> UDP ports: ${UDP_PORTS[*]}"
    echo ""
fi

case "$OS" in
    linux)
        case "$ACTION" in
            set)    linux_set    ;;
            reset)  linux_reset  ;;
            status) linux_status ;;
        esac
        ;;
    freebsd)
        case "$ACTION" in
            set)    freebsd_set    ;;
            reset)  freebsd_reset  ;;
            status) freebsd_status ;;
        esac
        ;;
    macos)
        case "$ACTION" in
            set)    macos_set    ;;
            reset)  macos_reset  ;;
            status) macos_status ;;
        esac
        ;;
esac

echo ""
echo "==> Done: $ACTION on $OS"