niom-turn/tests/rate_limit_udp.rs

use std::sync::Arc;

use niom_turn::alloc::AllocationManager;
use niom_turn::config::LimitsOptions;
use niom_turn::rate_limit::RateLimiters;
use tokio::net::UdpSocket;

use crate::support::stun_builders::{build_allocate_request, build_binding_request, parse};
use crate::support::{default_test_credentials, init_tracing, test_auth_manager};

mod support;

#[tokio::test]
async fn udp_binding_is_rate_limited_by_ip() {
    init_tracing();

    // Configure a very small burst to make the test deterministic.
    // We only limit Binding here to avoid affecting other integration tests.
    let mut limits = LimitsOptions::default();
    limits.binding_rps = Some(1);
    limits.binding_burst = Some(1);
    let rate_limiters = Arc::new(RateLimiters::from_limits(&limits));

    // Start a UDP server loop.
    let server = UdpSocket::bind("127.0.0.1:0").await.expect("server bind");
    let server_addr = server.local_addr().expect("server addr");

    let (username, password) = default_test_credentials();
    let auth = test_auth_manager(username, password);
    let allocs = AllocationManager::new();

    let server_arc = Arc::new(server);
    tokio::spawn({
        let reader = server_arc.clone();
        let auth = auth.clone();
        let allocs = allocs.clone();
        let rl = rate_limiters.clone();
        async move {
            let _ = niom_turn::server::udp_reader_loop_with_limits(reader, auth, allocs, rl).await;
        }
    });

    let client = UdpSocket::bind("127.0.0.1:0").await.expect("client bind");

    // Fire multiple Binding requests quickly; with burst=1 we should only get 1 success response.
    for _ in 0..3 {
        let req = build_binding_request();
        client
            .send_to(&req, server_addr)
            .await
            .expect("send binding");
    }

    let mut buf = [0u8; 1500];
    let mut responses = 0usize;

    // Collect responses for a short, bounded period.
    let deadline = tokio::time::Instant::now() + tokio::time::Duration::from_millis(150);
    loop {
        let now = tokio::time::Instant::now();
        if now >= deadline {
            break;
        }

        let remaining = deadline - now;
        match tokio::time::timeout(remaining, client.recv_from(&mut buf)).await {
            Ok(Ok((len, _from))) => {
                let resp = parse(&buf[..len]);
                // Success class
                assert_eq!(resp.header.msg_type & 0x0110, 0x0100);
                responses += 1;
            }
            Ok(Err(e)) => panic!("recv error: {e}"),
            Err(_) => break, // timeout
        }
    }

    assert_eq!(responses, 1, "expected exactly 1 Binding response under burst=1");

}

#[tokio::test]
async fn udp_unauth_challenge_is_rate_limited_by_ip() {
    init_tracing();

    // Configure a very small burst so only the first unauth challenge is answered.
    let mut limits = LimitsOptions::default();
    limits.unauth_rps = Some(1);
    limits.unauth_burst = Some(1);
    let rate_limiters = Arc::new(RateLimiters::from_limits(&limits));

    // Start a UDP server loop.
    let server = UdpSocket::bind("127.0.0.1:0").await.expect("server bind");
    let server_addr = server.local_addr().expect("server addr");

    let (username, password) = default_test_credentials();
    let auth = test_auth_manager(username, password);
    let allocs = AllocationManager::new();

    let server_arc = Arc::new(server);
    tokio::spawn({
        let reader = server_arc.clone();
        let auth = auth.clone();
        let allocs = allocs.clone();
        let rl = rate_limiters.clone();
        async move {
            let _ = niom_turn::server::udp_reader_loop_with_limits(reader, auth, allocs, rl).await;
        }
    });

    let client = UdpSocket::bind("127.0.0.1:0").await.expect("client bind");

    for _ in 0..3 {
        let req = build_allocate_request(None, None, None, None, None);
        client
            .send_to(&req, server_addr)
            .await
            .expect("send unauth allocate");
    }

    let mut buf = [0u8; 1500];
    let mut responses = 0usize;

    let deadline = tokio::time::Instant::now() + tokio::time::Duration::from_millis(150);
    loop {
        let now = tokio::time::Instant::now();
        if now >= deadline {
            break;
        }

        let remaining = deadline - now;
        match tokio::time::timeout(remaining, client.recv_from(&mut buf)).await {
            Ok(Ok((len, _from))) => {
                let resp = parse(&buf[..len]);
                assert_eq!(resp.header.msg_type & 0x0110, niom_turn::constants::CLASS_ERROR);
                resp.attributes
                    .iter()
                    .find(|a| a.typ == niom_turn::constants::ATTR_NONCE)
                    .expect("nonce attr");
                responses += 1;
            }
            Ok(Err(e)) => panic!("recv error: {e}"),
            Err(_) => break,
        }
    }

    assert_eq!(responses, 1, "expected exactly 1 unauth challenge under burst=1");
}