Files
ethereum-rpc-docker/limit-bandwidth.sh

321 lines
13 KiB
Bash
Executable File

#!/bin/bash
# General script to limit bandwidth for all public ports in a Docker Compose file
# Limits each port to specified bandwidth (default: 100 MBit/s)
#
# Usage:
# ./limit-bandwidth.sh <compose-file> [start|stop|status] [--limit BANDWIDTH]
# ./limit-bandwidth.sh rpc/gnosis/reth/gnosis-mainnet-reth-pruned-trace.yml start
# ./limit-bandwidth.sh rpc/gnosis/reth/gnosis-mainnet-reth-pruned-trace.yml start --limit 20mbit
#
# Environment variable:
# BANDWIDTH_LIMIT=20mbit ./limit-bandwidth.sh <compose-file> start
#
# For cronjob:
# */5 * * * * /path/to/rpc/limit-bandwidth.sh /path/to/compose.yml start
# */5 * * * * BANDWIDTH_LIMIT=20mbit /path/to/rpc/limit-bandwidth.sh /path/to/compose.yml start
set -uo pipefail
BASEPATH="$(cd "$(dirname "$0")" && pwd)"
# Bandwidth limit (can be overridden via BANDWIDTH_LIMIT env var or --limit parameter)
# Default: 100mbit
BANDWIDTH_LIMIT="${BANDWIDTH_LIMIT:-100mbit}"
BURST_MULTIPLIER="${BURST_MULTIPLIER:-0.1}" # Burst is 10% of limit by default
LATENCY="50ms" # Latency for shaping
# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color
# Check if running as root
if [ "$EUID" -ne 0 ]; then
echo -e "${RED}Error: This script must be run as root (use sudo)${NC}"
exit 1
fi
# Check if required tools are installed
if ! command -v tc >/dev/null 2>&1; then
echo -e "${RED}Error: 'tc' command not found${NC}"
echo "Please install iproute2:"
echo " sudo apt-get update && sudo apt-get install -y iproute2"
exit 1
fi
if ! command -v iptables >/dev/null 2>&1; then
echo -e "${RED}Error: 'iptables' command not found${NC}"
echo "Please install iptables:"
echo " sudo apt-get update && sudo apt-get install -y iptables"
exit 1
fi
# Parse arguments
if [ $# -lt 1 ]; then
echo "Usage: $0 <compose-file> [start|stop|status] [--limit BANDWIDTH]"
echo ""
echo "Options:"
echo " --limit BANDWIDTH Set bandwidth limit (e.g., 20mbit, 100mbit)"
echo " Can also use BANDWIDTH_LIMIT environment variable"
echo ""
echo "Examples:"
echo " $0 rpc/gnosis/reth/gnosis-mainnet-reth-pruned-trace.yml start"
echo " $0 rpc/gnosis/reth/gnosis-mainnet-reth-pruned-trace.yml start --limit 20mbit"
echo " BANDWIDTH_LIMIT=20mbit $0 rpc/gnosis/reth/gnosis-mainnet-reth-pruned-trace.yml start"
echo " $0 rpc/ethereum/geth/ethereum-mainnet-geth-pruned-pebble-path.yml status"
exit 1
fi
COMPOSE_FILE="$1"
ACTION=${2:-start}
# Parse optional --limit parameter
if [ "$ACTION" = "--limit" ] && [ $# -ge 3 ]; then
BANDWIDTH_LIMIT="$3"
ACTION=${4:-start}
elif [ $# -ge 3 ] && [ "$2" = "--limit" ]; then
BANDWIDTH_LIMIT="$3"
ACTION=${4:-start}
elif [ $# -ge 4 ] && [ "$3" = "--limit" ]; then
BANDWIDTH_LIMIT="$4"
elif [ "$ACTION" = "--limit" ]; then
echo -e "${RED}Error: --limit requires a bandwidth value${NC}"
echo "Example: $0 $COMPOSE_FILE start --limit 20mbit"
exit 1
fi
# Resolve compose file path
if [ ! -f "$COMPOSE_FILE" ]; then
# Try relative to BASEPATH
if [ -f "$BASEPATH/$COMPOSE_FILE" ]; then
COMPOSE_FILE="$BASEPATH/$COMPOSE_FILE"
else
echo -e "${RED}Error: Compose file not found: $1${NC}"
exit 1
fi
fi
COMPOSE_FILE=$(realpath "$COMPOSE_FILE")
COMPOSE_DIR=$(dirname "$COMPOSE_FILE")
# Calculate burst from limit (10% of limit)
# Extract number and unit from BANDWIDTH_LIMIT (e.g., "20mbit" -> "20" and "mbit")
if [[ "$BANDWIDTH_LIMIT" =~ ^([0-9]+)([a-zA-Z]+)$ ]]; then
LIMIT_NUM="${BASH_REMATCH[1]}"
LIMIT_UNIT="${BASH_REMATCH[2]}"
# Calculate 10% of limit (simple integer division for compatibility)
BURST_NUM=$((LIMIT_NUM / 10))
# Ensure minimum burst of 1
if [ "$BURST_NUM" -lt 1 ]; then
BURST_NUM=1
fi
BURST="${BURST_NUM}${LIMIT_UNIT}"
else
# Fallback: use 10% of limit as string manipulation
BURST="${BANDWIDTH_LIMIT}"
fi
echo -e "${BLUE}Compose file: ${COMPOSE_FILE}${NC}"
echo -e "${BLUE}Bandwidth limit: ${BANDWIDTH_LIMIT} per port${NC}"
# Function to extract ports from YAML
extract_ports() {
local file="$1"
# Extract ports in format "HOST:CONTAINER" or "HOST:CONTAINER/PROTO"
# Look for lines after "ports:" that contain port mappings
awk '
/^[[:space:]]*ports:[[:space:]]*$/ { in_ports=1; next }
/^[[:space:]]*-/ && in_ports {
# Match patterns like "13516:13516" or "13516:13516/udp"
if (match($0, /[0-9]+:[0-9]+/)) {
port = substr($0, RSTART, RLENGTH)
# Extract just the host port (first number)
if (match(port, /^[0-9]+/)) {
print substr(port, RSTART, RLENGTH)
}
}
}
/^[[:space:]]*[a-zA-Z]/ && in_ports { in_ports=0 }
' "$file" | sort -u
}
# Function to find Docker bridge for compose project
find_bridge() {
local compose_dir="$1"
local project_name=$(basename "$compose_dir" | tr '[:upper:]' '[:lower:]' | sed 's/[^a-z0-9]/-/g')
# Try to find the network by checking running containers from this compose file
local container_name=$(docker compose -f "$COMPOSE_FILE" ps --format json 2>/dev/null | \
jq -r '.[0].Name' 2>/dev/null | head -1)
if [ -n "$container_name" ] && [ "$container_name" != "null" ]; then
# Get network from container
local network=$(docker inspect "$container_name" --format '{{range $net, $conf := .NetworkSettings.Networks}}{{$net}}{{end}}' 2>/dev/null | head -1)
if [ -n "$network" ]; then
# Find bridge interface for this network
local bridge=$(docker network inspect "$network" --format '{{range .IPAM.Config}}{{.Gateway}}{{end}}' 2>/dev/null)
if [ -n "$bridge" ]; then
# Find interface with this gateway IP
local iface=$(ip route | grep "$bridge" | awk '{print $3}' | head -1)
if [ -n "$iface" ]; then
echo "$iface"
return 0
fi
fi
fi
fi
# Fallback: find Docker bridge interfaces
local bridges=$(ip link show | grep -E "^[0-9]+: (br-|docker0)" | cut -d: -f2 | tr -d ' ' | head -1)
if [ -n "$bridges" ]; then
echo "$bridges"
return 0
fi
# Last resort: use docker0
if ip link show docker0 >/dev/null 2>&1; then
echo "docker0"
return 0
fi
echo "eth0" # Final fallback
}
# Extract ports from compose file
PORTS_RAW=$(extract_ports "$COMPOSE_FILE")
if [ -z "$PORTS_RAW" ]; then
echo -e "${YELLOW}Warning: No public ports found in compose file${NC}"
exit 0
fi
# Convert to array (simple approach)
PORTS=()
while IFS= read -r line; do
[ -n "$line" ] && PORTS+=("$line")
done <<< "$PORTS_RAW"
if [ ${#PORTS[@]} -eq 0 ]; then
echo -e "${YELLOW}Warning: No public ports found in compose file${NC}"
exit 0
fi
echo -e "${GREEN}Found ${#PORTS[@]} public port(s): ${PORTS[*]}${NC}"
# Find the bridge interface
BRIDGE=$(find_bridge "$COMPOSE_DIR")
echo -e "${BLUE}Using network interface: ${BRIDGE}${NC}"
# Verify interface exists
if ! ip link show "$BRIDGE" >/dev/null 2>&1; then
echo -e "${RED}Error: Network interface '$BRIDGE' not found${NC}"
exit 1
fi
case "$ACTION" in
start)
echo -e "${YELLOW}Setting up bandwidth limiting...${NC}"
# Remove existing qdisc if any
tc qdisc del dev "$BRIDGE" root 2>/dev/null || true
# Create HTB (Hierarchical Token Bucket) qdisc
tc qdisc add dev "$BRIDGE" root handle 1: htb default 30
# Create root class with high bandwidth
tc class add dev "$BRIDGE" parent 1: classid 1:1 htb rate 1000mbit
# Create unlimited class for non-limited traffic
tc class add dev "$BRIDGE" parent 1:1 classid 1:30 htb rate 1000mbit
# Process each port
PORT_ID=10
for port in "${PORTS[@]}"; do
echo -e "${BLUE} Limiting port ${port} to ${BANDWIDTH_LIMIT}...${NC}"
# Create limited class for this port
tc class add dev "$BRIDGE" parent 1:1 classid 1:${PORT_ID} htb rate ${BANDWIDTH_LIMIT} burst ${BURST} ceil ${BANDWIDTH_LIMIT}
# Add filter to route marked packets to this class
tc filter add dev "$BRIDGE" parent 1: protocol ip prio ${PORT_ID} handle ${PORT_ID} fw flowid 1:${PORT_ID}
# Mark TCP traffic for this port
iptables -t mangle -C OUTPUT -p tcp --sport ${port} -j MARK --set-mark ${PORT_ID} 2>/dev/null || \
iptables -t mangle -A OUTPUT -p tcp --sport ${port} -j MARK --set-mark ${PORT_ID}
iptables -t mangle -C OUTPUT -p tcp --dport ${port} -j MARK --set-mark ${PORT_ID} 2>/dev/null || \
iptables -t mangle -A OUTPUT -p tcp --dport ${port} -j MARK --set-mark ${PORT_ID}
# Mark UDP traffic for this port
iptables -t mangle -C OUTPUT -p udp --sport ${port} -j MARK --set-mark ${PORT_ID} 2>/dev/null || \
iptables -t mangle -A OUTPUT -p udp --sport ${port} -j MARK --set-mark ${PORT_ID}
iptables -t mangle -C OUTPUT -p udp --dport ${port} -j MARK --set-mark ${PORT_ID} 2>/dev/null || \
iptables -t mangle -A OUTPUT -p udp --dport ${port} -j MARK --set-mark ${PORT_ID}
# Mark in FORWARD chain (container traffic)
iptables -t mangle -C FORWARD -p tcp --sport ${port} -j MARK --set-mark ${PORT_ID} 2>/dev/null || \
iptables -t mangle -A FORWARD -p tcp --sport ${port} -j MARK --set-mark ${PORT_ID}
iptables -t mangle -C FORWARD -p tcp --dport ${port} -j MARK --set-mark ${PORT_ID} 2>/dev/null || \
iptables -t mangle -A FORWARD -p tcp --dport ${port} -j MARK --set-mark ${PORT_ID}
iptables -t mangle -C FORWARD -p udp --sport ${port} -j MARK --set-mark ${PORT_ID} 2>/dev/null || \
iptables -t mangle -A FORWARD -p udp --sport ${port} -j MARK --set-mark ${PORT_ID}
iptables -t mangle -C FORWARD -p udp --dport ${port} -j MARK --set-mark ${PORT_ID} 2>/dev/null || \
iptables -t mangle -A FORWARD -p udp --dport ${port} -j MARK --set-mark ${PORT_ID}
PORT_ID=$((PORT_ID + 1))
done
echo -e "${GREEN}✓ Bandwidth limiting configured!${NC}"
echo -e "${GREEN}All ports are now limited to ${BANDWIDTH_LIMIT} each${NC}"
;;
stop)
echo -e "${YELLOW}Removing bandwidth limiting...${NC}"
# Remove iptables rules for each port
PORT_ID=10
for port in "${PORTS[@]}"; do
echo -e "${BLUE} Removing limits for port ${port}...${NC}"
iptables -t mangle -D OUTPUT -p tcp --sport ${port} -j MARK --set-mark ${PORT_ID} 2>/dev/null || true
iptables -t mangle -D OUTPUT -p tcp --dport ${port} -j MARK --set-mark ${PORT_ID} 2>/dev/null || true
iptables -t mangle -D OUTPUT -p udp --sport ${port} -j MARK --set-mark ${PORT_ID} 2>/dev/null || true
iptables -t mangle -D OUTPUT -p udp --dport ${port} -j MARK --set-mark ${PORT_ID} 2>/dev/null || true
iptables -t mangle -D FORWARD -p tcp --sport ${port} -j MARK --set-mark ${PORT_ID} 2>/dev/null || true
iptables -t mangle -D FORWARD -p tcp --dport ${port} -j MARK --set-mark ${PORT_ID} 2>/dev/null || true
iptables -t mangle -D FORWARD -p udp --sport ${port} -j MARK --set-mark ${PORT_ID} 2>/dev/null || true
iptables -t mangle -D FORWARD -p udp --dport ${port} -j MARK --set-mark ${PORT_ID} 2>/dev/null || true
PORT_ID=$((PORT_ID + 1))
done
# Remove tc qdisc
tc qdisc del dev "$BRIDGE" root 2>/dev/null || true
echo -e "${GREEN}✓ Bandwidth limiting removed${NC}"
;;
status)
echo -e "${YELLOW}Current bandwidth limiting status:${NC}"
echo ""
echo -e "${BLUE}TC Qdisc on ${BRIDGE}:${NC}"
tc qdisc show dev "$BRIDGE" 2>/dev/null || echo " No qdisc configured"
echo ""
echo -e "${BLUE}TC Classes:${NC}"
tc class show dev "$BRIDGE" 2>/dev/null || echo " No classes configured"
echo ""
echo -e "${BLUE}Iptables rules (OUTPUT):${NC}"
iptables -t mangle -L OUTPUT -n --line-numbers | grep -E "MARK|${PORTS[0]}" || echo " No rules found"
echo ""
echo -e "${BLUE}Iptables rules (FORWARD):${NC}"
iptables -t mangle -L FORWARD -n --line-numbers | grep -E "MARK|${PORTS[0]}" || echo " No rules found"
echo ""
echo -e "${BLUE}Monitored ports: ${PORTS[*]}${NC}"
;;
*)
echo -e "${RED}Error: Invalid action '$ACTION'${NC}"
echo "Usage: $0 <compose-file> [start|stop|status]"
exit 1
;;
esac