niom-turn/tests/tls_turn.rs

use std::sync::Arc;

use niom_turn::alloc::AllocationManager;
use niom_turn::stun::parse_message;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::{TcpListener, UdpSocket};
use tokio_rustls::{rustls::ServerConfig, TlsAcceptor};

use crate::support::stun_builders::{build_allocate_request, build_refresh_request};
use crate::support::{default_test_credentials, init_tracing_with, test_auth_manager};

mod support;

#[tokio::test]
async fn tls_allocate_refresh_flow() {
    init_tracing_with("warn,niom_turn=info");

    let udp = UdpSocket::bind("127.0.0.1:0").await.expect("udp bind");
    let udp_arc = Arc::new(udp);
    let (username, password) = default_test_credentials();
    let auth = test_auth_manager(username, password);
    let allocs = AllocationManager::new();

    let (cert, key) = support::tls::generate_self_signed_cert();
    let mut cfg = ServerConfig::builder()
        .with_safe_defaults()
        .with_no_client_auth()
        .with_single_cert(vec![cert.clone()], key)
        .expect("server config");
    cfg.alpn_protocols.push(b"turn".to_vec());
    let acceptor = TlsAcceptor::from(Arc::new(cfg));

    let tcp_listener = TcpListener::bind("127.0.0.1:0").await.expect("tcp bind");
    let tcp_addr = tcp_listener.local_addr().expect("tcp addr");
    let udp_clone = udp_arc.clone();
    let auth_clone = auth.clone();
    let alloc_clone = allocs.clone();

    tokio::spawn(async move {
        loop {
            let (stream, peer) = match tcp_listener.accept().await {
                Ok(conn) => conn,
                Err(_) => break,
            };
            let acceptor = acceptor.clone();
            let udp_clone = udp_clone.clone();
            let auth_clone = auth_clone.clone();
            let alloc_clone = alloc_clone.clone();
            tokio::spawn(async move {
                match acceptor.accept(stream).await {
                    Ok(mut tls_stream) => {
                        match niom_turn::tls::handle_tls_connection(
                            &mut tls_stream,
                            peer,
                            udp_clone,
                            auth_clone,
                            alloc_clone,
                        )
                        .await
                        {
                            Ok(_) => {}
                            Err(e) => {
                                tracing::error!("tls connection error: {:?}", e);
                            }
                        }
                    }
                    Err(e) => {
                        tracing::error!("tls accept failed: {:?}", e);
                    }
                }
            });
        }
    });

    // Build client config trusting generated cert
    let mut root_store = tokio_rustls::rustls::RootCertStore::empty();
    root_store.add(&cert).expect("add root");
    let client_config = tokio_rustls::rustls::ClientConfig::builder()
        .with_safe_defaults()
        .with_root_certificates(root_store)
        .with_no_client_auth();
    let connector = tokio_rustls::TlsConnector::from(Arc::new(client_config));

    let tcp_stream = tokio::net::TcpStream::connect(tcp_addr)
        .await
        .expect("tcp connect");
    let domain = tokio_rustls::rustls::ServerName::try_from("localhost").unwrap();
    let mut tls_stream = connector
        .connect(domain, tcp_stream)
        .await
        .expect("tls connect");
    tracing::info!("client connected");

    let allocate = build_allocate_request(None, None, None, None, None);
    tls_stream
        .write_all(&allocate)
        .await
        .expect("write allocate");
    tracing::info!("sent unauthenticated allocate request");

    let mut buf = vec![0u8; 1500];
    let n = tls_stream.read(&mut buf).await.expect("read challenge");
    tracing::info!(bytes = n, "received nonce challenge");
    let resp = parse_message(&buf[..n]).expect("parse 401");
    let nonce_attr = resp
        .attributes
        .iter()
        .find(|a| a.typ == niom_turn::constants::ATTR_NONCE)
        .expect("nonce attr");
    let nonce = String::from_utf8(nonce_attr.value.clone()).expect("nonce str");

    let key = niom_turn::auth::compute_a1_md5(username, auth.realm(), password);
    let allocate = build_allocate_request(
        Some(username),
        Some(auth.realm()),
        Some(&nonce),
        Some(&key),
        Some(600),
    );
    tls_stream
        .write_all(&allocate)
        .await
        .expect("write auth allocate");
    tracing::info!("sent authenticated allocate request");
    let n = tls_stream.read(&mut buf).await.expect("read success");
    tracing::info!(bytes = n, "received allocate success");
    let resp = parse_message(&buf[..n]).expect("parse alloc success");
    assert_eq!(resp.header.msg_type & 0x0110, 0x0100);

    let refresh = build_refresh_request(
        resp.header.transaction_id,
        username,
        auth.realm(),
        &nonce,
        &key,
        0,
    );
    tls_stream.write_all(&refresh).await.expect("write refresh");
    tracing::info!("sent refresh request");
    let n = tls_stream.read(&mut buf).await.expect("read refresh resp");
    tracing::info!(bytes = n, "received refresh response");
    let resp = parse_message(&buf[..n]).expect("parse refresh resp");
    let lifetime = resp
        .attributes
        .iter()
        .find(|a| a.typ == niom_turn::constants::ATTR_LIFETIME)
        .expect("lifetime attr");
    let secs = u32::from_be_bytes([
        lifetime.value[0],
        lifetime.value[1],
        lifetime.value[2],
        lifetime.value[3],
    ]);
    assert_eq!(secs, 0);
}