Rate limiting WebSocket messages for Agents

Problem

Protecting WebSocket-based AI services from Overload caused by high message rates, temporary spikes via burst control, resource waste from duplicate or repeated messages, aggressive/malicious agents with temporary penalties, and lack of visibility via structured JSON logging.

Solution

This iRule protects WebSocket endpoints from aggressive or misbehaving AI agents by enforcing message rate limits, burst controls, and duplicate suppression. Each client IP is allowed up to 40 messages per 10 seconds (rate_limit / rate_window) with a maximum of 20 messages per second (burst_limit). Duplicate messages within 5 seconds (dup_ttl) are dropped, and any client exceeding limits is temporarily penalized for 60 seconds (penalty_time) and disconnected. All violations are logged in JSON format to an HSL pool, including timestamp, client IP, event type, message content, and count.

Impact

For organizations running AI at scale, this is a huge game changer that safeguards availability, performance, and security across potentially thousands of clients simultaneously.

Code

when RULE_INIT {

    # HSL pool for JSON logging
    set static::hsl_pool "syslog_pool"

    # Sliding window rate limit: 40 messages per 10 seconds
    set static::rate_limit 40
    set static::rate_window 10

    # Burst protection: 20 messages per second
    set static::burst_limit 20

    # Duplicate message suppression TTL (seconds)
    set static::dup_ttl 5

    # Penalty/quarantine duration (seconds)
    set static::penalty_time 60

}

# -----------------------------
# Detect WebSocket Upgrade
# -----------------------------
when HTTP_REQUEST {
    if {[string tolower [HTTP::header "Upgrade"]] eq "websocket"} {
        # Nothing required here, IP can be grabbed from client_addr in other events
    }
}

# -----------------------------
# Inspect WebSocket Frames
# -----------------------------
when WS_CLIENT_DATA {
    set payload [WS::payload]
}

when WS_CLIENT_FRAME {

    set ip [IP::client_addr]

    # Open HSL
    set hsl [HSL::open -proto UDP -pool $static::hsl_pool]

    # -----------------------------
    # Check penalty/quarantine
    # -----------------------------
    if {[table lookup "ws_penalty:$ip"] ne ""} {

        # Log JSON event
        set ts [clock format [clock seconds] -gmt 1 -format "%Y-%m-%dT%H:%M:%SZ"]
        set logmsg [string map {\" \\\" \n "" \r ""} $payload]
        set json "{\
\"timestamp\":\"$ts\",\
\"client_ip\":\"$ip\",\
\"event\":\"penalty_block\",\
\"message\":\"$logmsg\",\
\"count\":\"0\"\
}"
        HSL::send $hsl $json

        # Mark violation for disconnect
        table set "ws_violation:$ip" 1 2
        return
    }

    # -----------------------------
    # Sliding window rate counter
    # -----------------------------
    set rate_key "ws_rate:$ip"
    set rate [table incr $rate_key]
    if {$rate == 1} {
        table timeout $rate_key $static::rate_window
    }

    # -----------------------------
    # Burst detection
    # -----------------------------
    set burst_key "ws_burst:$ip"
    set burst [table incr $burst_key]
    if {$burst == 1} {
        table timeout $burst_key 1
    }

    # -----------------------------
    # Duplicate message detection
    # -----------------------------
    set hash [crc32 $payload]
    set dup_key "ws_dup:$ip:$hash"

    if {[table lookup $dup_key] ne ""} {

        # Log duplicate message
        set ts [clock format [clock seconds] -gmt 1 -format "%Y-%m-%dT%H:%M:%SZ"]
        set logmsg [string map {\" \\\" \n "" \r ""} $payload]
        set json "{\
\"timestamp\":\"$ts\",\
\"client_ip\":\"$ip\",\
\"event\":\"duplicate_message\",\
\"message\":\"$logmsg\",\
\"count\":\"$rate\"\
}"
        HSL::send $hsl $json

        WS::frame drop
        return
    }

    # Store this message hash for duplicate detection
    table set $dup_key 1 $static::dup_ttl

    # -----------------------------
    # Rate violation check
    # -----------------------------
    if {$rate > $static::rate_limit || $burst > $static::burst_limit} {

        # Log rate limit exceeded
        set ts [clock format [clock seconds] -gmt 1 -format "%Y-%m-%dT%H:%M:%SZ"]
        set logmsg [string map {\" \\\" \n "" \r ""} $payload]
        set json "{\
\"timestamp\":\"$ts\",\
\"client_ip\":\"$ip\",\
\"event\":\"rate_limit_exceeded\",\
\"message\":\"$logmsg\",\
\"count\":\"$rate\"\
}"
        HSL::send $hsl $json

        # Apply penalty/quarantine
        table set "ws_penalty:$ip" 1 $static::penalty_time

        # Mark violation for disconnect in FRAME_DONE
        table set "ws_violation:$ip" 1 2
        return
    }

}

# -----------------------------
# Disconnect violating clients in valid event
# -----------------------------
when WS_CLIENT_FRAME_DONE {

    set ip [IP::client_addr]

    if {[table lookup "ws_violation:$ip"] eq "1"} {
        WS::disconnect 1000 "Violation occurred"
        table delete "ws_violation:$ip"
    }

}