niom-turn/tests/udp_turn.rs

use std::net::SocketAddr;
use std::sync::Arc;

use niom_turn::alloc::AllocationManager;
use niom_turn::auth::InMemoryStore;
use niom_turn::server::udp_reader_loop;
use niom_turn::stun::parse_message;
use tokio::net::UdpSocket;

use crate::support::stun_builders::{
    build_allocate_request, build_create_permission_request, build_refresh_request,
    build_send_request, new_transaction_id, parse,
};
use crate::support::{default_test_credentials, init_tracing, test_auth_manager};

mod support;

const SERVER_ADDR: &str = "127.0.0.1:0";

#[tokio::test]
async fn allocate_requires_auth_then_succeeds() {
    init_tracing();
    let (server, client_addr) = start_udp_server().await;
    let (username, password) = default_test_credentials();
    let auth = test_auth_manager(username, password);
    let allocs = AllocationManager::new();
    let server_arc = Arc::new(server);

    let server_clone = server_arc.clone();
    let auth_clone = auth.clone();
    let alloc_clone = allocs.clone();
    tokio::spawn(async move {
        let _ = udp_reader_loop(server_clone, auth_clone, alloc_clone).await;
    });

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

    // initial unauthenticated allocate should trigger 401 with nonce
    let req = build_allocate_request(None, None, None, None, None);
    client
        .send_to(&req, client_addr)
        .await
        .expect("send unauth allocate");
    let mut buf = [0u8; 1500];
    let (len, _) = client.recv_from(&mut buf).await.expect("recv challenge");
    let resp = parse_message(&buf[..len]).expect("parse 401");
    assert_eq!(resp.header.msg_type & 0x0110, 0x0110);

    let nonce = resp
        .attributes
        .iter()
        .find(|a| a.typ == niom_turn::constants::ATTR_NONCE)
        .expect("nonce attr")
        .value
        .clone();
    let nonce_str = String::from_utf8(nonce).expect("nonce utf8");
    let key = niom_turn::auth::compute_a1_md5(username, auth.realm(), password);

    let req = build_allocate_request(
        Some(username),
        Some(auth.realm()),
        Some(&nonce_str),
        Some(&key),
        Some(600),
    );
    client
        .send_to(&req, client_addr)
        .await
        .expect("send auth allocate");
    let (len, _) = client.recv_from(&mut buf).await.expect("recv success");
    let resp = parse(&buf[..len]);
    assert_eq!(resp.header.msg_type & 0x0110, 0x0100);
}

#[tokio::test]
async fn refresh_zero_lifetime_releases_allocation() {
    init_tracing();
    let (server, client_addr) = start_udp_server().await;
    let (username, password) = default_test_credentials();
    let auth = test_auth_manager(username, password);
    let allocs = AllocationManager::new();
    let server_arc = Arc::new(server);

    let server_clone = server_arc.clone();
    let auth_clone = auth.clone();
    let alloc_clone = allocs.clone();
    tokio::spawn(async move {
        let _ = udp_reader_loop(server_clone, auth_clone, alloc_clone).await;
    });

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

    let nonce =
        perform_authenticated_allocate(&client, client_addr, &auth, username, password, &allocs)
            .await;

    let key = niom_turn::auth::compute_a1_md5(username, auth.realm(), password);
    let trans_id = new_transaction_id();
    let refresh = build_refresh_request(trans_id, username, auth.realm(), &nonce, &key, 0);
    client
        .send_to(&refresh, client_addr)
        .await
        .expect("send refresh");
    let mut buf = [0u8; 1500];
    let (len, _) = client.recv_from(&mut buf).await.expect("recv refresh resp");
    let resp = parse(&buf[..len]);
    assert_eq!(resp.header.msg_type & 0x0110, 0x0100);
    let lifetime = resp
        .attributes
        .iter()
        .find(|a| a.typ == niom_turn::constants::ATTR_LIFETIME)
        .expect("lifetime attr");
    assert_eq!(
        u32::from_be_bytes([
            lifetime.value[0],
            lifetime.value[1],
            lifetime.value[2],
            lifetime.value[3]
        ]),
        0
    );

    assert!(allocs
        .get_allocation(&client.local_addr().unwrap())
        .is_none());
}

#[tokio::test]
async fn create_permission_and_send_relays_data() {
    init_tracing();
    let (server, client_addr) = start_udp_server().await;
    let (username, password) = default_test_credentials();
    let auth = test_auth_manager(username, password);
    let allocs = AllocationManager::new();
    let server_arc = Arc::new(server);

    let server_clone = server_arc.clone();
    let auth_clone = auth.clone();
    let alloc_clone = allocs.clone();
    tokio::spawn(async move {
        let _ = udp_reader_loop(server_clone, auth_clone, alloc_clone).await;
    });

    let client = UdpSocket::bind("127.0.0.1:0").await.expect("client bind");
    let nonce =
        perform_authenticated_allocate(&client, client_addr, &auth, username, password, &allocs)
            .await;

    let key = niom_turn::auth::compute_a1_md5(username, auth.realm(), password);
    let peer_sock = UdpSocket::bind("127.0.0.1:0").await.expect("peer bind");
    let relay_addr = allocs
        .get_allocation(&client.local_addr().unwrap())
        .expect("allocation exists")
        .relay_addr;

    let perm_req = build_create_permission_request(
        username,
        auth.realm(),
        &nonce,
        &key,
        &peer_sock.local_addr().unwrap(),
    );
    client
        .send_to(&perm_req, client_addr)
        .await
        .expect("send create permission");
    let mut buf = [0u8; 1500];
    client.recv_from(&mut buf).await.expect("permission resp");

    let payload = b"hello-turn";
    let send_req = build_send_request(
        username,
        auth.realm(),
        &nonce,
        &key,
        &peer_sock.local_addr().unwrap(),
        payload,
    );
    client
        .send_to(&send_req, client_addr)
        .await
        .expect("send indication");

    let mut peer_buf = [0u8; 1500];
    let (len, addr) = peer_sock.recv_from(&mut peer_buf).await.expect("peer recv");
    assert_eq!(len, payload.len());
    assert_eq!(&peer_buf[..len], payload);
    assert_eq!(addr.port(), relay_addr.port());
    assert!(addr.ip().is_loopback());
}

async fn start_udp_server() -> (UdpSocket, SocketAddr) {
    let server = UdpSocket::bind(SERVER_ADDR).await.expect("server bind");
    let addr = server.local_addr().expect("server addr");
    (server, addr)
}

async fn perform_authenticated_allocate(
    client: &UdpSocket,
    server_addr: SocketAddr,
    auth: &niom_turn::auth::AuthManager<InMemoryStore>,
    username: &str,
    password: &str,
    allocs: &AllocationManager,
) -> String {
    init_tracing();
    // trigger nonce challenge
    let req = build_allocate_request(None, None, None, None, None);
    client
        .send_to(&req, server_addr)
        .await
        .expect("send initial allocate");
    let mut buf = [0u8; 1500];
    let (len, _) = client.recv_from(&mut buf).await.expect("recv nonce");
    let resp = parse_message(&buf[..len]).expect("parse nonce resp");
    let nonce = resp
        .attributes
        .iter()
        .find(|a| a.typ == niom_turn::constants::ATTR_NONCE)
        .expect("nonce attr")
        .value
        .clone();
    let nonce = String::from_utf8(nonce).expect("nonce utf8");
    let key = niom_turn::auth::compute_a1_md5(username, auth.realm(), password);

    let req = build_allocate_request(
        Some(username),
        Some(auth.realm()),
        Some(&nonce),
        Some(&key),
        Some(600),
    );
    client
        .send_to(&req, server_addr)
        .await
        .expect("send auth allocate");
    client.recv_from(&mut buf).await.expect("recv success");

    assert!(allocs
        .get_allocation(&client.local_addr().unwrap())
        .is_some());
    nonce
}