niom-turn/tests/support/stun_builders.rs

#![allow(dead_code)]

use bytes::BytesMut;
use niom_turn::constants::*;
use niom_turn::stun::{compute_message_integrity, parse_message};
use uuid::Uuid;

/// Build a basic STUN header for TURN requests with a freshly generated transaction id.
pub fn new_transaction_id() -> [u8; 12] {
    let uuid = Uuid::new_v4();
    let mut trans = [0u8; 12];
    trans.copy_from_slice(&uuid.as_bytes()[..12]);
    trans
}

/// Construct a TURN Allocate request optionally including lifetime and auth attributes.
pub fn build_allocate_request(
    username: Option<&str>,
    realm: Option<&str>,
    nonce: Option<&str>,
    key: Option<&[u8]>,
    lifetime: Option<u32>,
) -> Vec<u8> {
    build_authenticated_request(
        METHOD_ALLOCATE,
        username,
        realm,
        nonce,
        key,
        lifetime,
        None,
        None,
    )
}

/// Construct a TURN Refresh request using MESSAGE-INTEGRITY.
pub fn build_refresh_request(
    trans: [u8; 12],
    username: &str,
    realm: &str,
    nonce: &str,
    key: &[u8],
    lifetime: u32,
) -> Vec<u8> {
    build_request_with_body(
        METHOD_REFRESH,
        Some(username),
        Some(realm),
        Some(nonce),
        Some(key),
        Some(lifetime),
        None,
        None,
        Some(trans),
    )
}

/// Build a CreatePermission request for the specified peer address.
pub fn build_create_permission_request(
    username: &str,
    realm: &str,
    nonce: &str,
    key: &[u8],
    peer: &std::net::SocketAddr,
) -> Vec<u8> {
    build_request_with_body(
        METHOD_CREATE_PERMISSION,
        Some(username),
        Some(realm),
        Some(nonce),
        Some(key),
        None,
        Some(peer),
        None,
        None,
    )
}

/// Build a Send indication to forward payload to peer.
pub fn build_send_request(
    username: &str,
    realm: &str,
    nonce: &str,
    key: &[u8],
    peer: &std::net::SocketAddr,
    payload: &[u8],
) -> Vec<u8> {
    build_request_with_body(
        METHOD_SEND,
        Some(username),
        Some(realm),
        Some(nonce),
        Some(key),
        None,
        Some(peer),
        Some(payload),
        None,
    )
}

fn build_authenticated_request(
    method: u16,
    username: Option<&str>,
    realm: Option<&str>,
    nonce: Option<&str>,
    key: Option<&[u8]>,
    lifetime: Option<u32>,
    peer: Option<&std::net::SocketAddr>,
    payload: Option<&[u8]>,
) -> Vec<u8> {
    build_request_with_body(
        method, username, realm, nonce, key, lifetime, peer, payload, None,
    )
}

fn build_request_with_body(
    method: u16,
    username: Option<&str>,
    realm: Option<&str>,
    nonce: Option<&str>,
    key: Option<&[u8]>,
    lifetime: Option<u32>,
    peer: Option<&std::net::SocketAddr>,
    payload: Option<&[u8]>,
    override_trans: Option<[u8; 12]>,
) -> Vec<u8> {
    let mut buf = BytesMut::new();
    buf.extend_from_slice(&method.to_be_bytes());
    buf.extend_from_slice(&0u16.to_be_bytes());
    buf.extend_from_slice(&MAGIC_COOKIE_BYTES);
    let trans = override_trans.unwrap_or_else(new_transaction_id);
    buf.extend_from_slice(&trans);

    if let Some(username) = username {
        push_string_attr(&mut buf, ATTR_USERNAME, username);
    }
    if let Some(realm) = realm {
        push_string_attr(&mut buf, ATTR_REALM, realm);
    }
    if let Some(nonce) = nonce {
        push_string_attr(&mut buf, ATTR_NONCE, nonce);
    }
    if let Some(lifetime) = lifetime {
        push_u32_attr(&mut buf, ATTR_LIFETIME, lifetime);
    }
    if let Some(peer) = peer {
        let encoded = niom_turn::stun::encode_xor_peer_address(peer, &trans);
        push_bytes_attr(&mut buf, ATTR_XOR_PEER_ADDRESS, &encoded);
    }
    if let Some(data) = payload {
        push_bytes_attr(&mut buf, ATTR_DATA, data);
    }

    if let Some(key) = key {
        append_message_integrity(&mut buf, key);
    }

    // update length
    let total_len = (buf.len() - 20) as u16;
    let len_bytes = total_len.to_be_bytes();
    buf[2] = len_bytes[0];
    buf[3] = len_bytes[1];

    buf.to_vec()
}

fn push_string_attr(buf: &mut BytesMut, typ: u16, value: &str) {
    push_bytes_attr(buf, typ, value.as_bytes());
}

fn push_u32_attr(buf: &mut BytesMut, typ: u16, value: u32) {
    push_bytes_attr(buf, typ, &value.to_be_bytes());
}

fn push_bytes_attr(buf: &mut BytesMut, typ: u16, data: &[u8]) {
    buf.extend_from_slice(&typ.to_be_bytes());
    buf.extend_from_slice(&(data.len() as u16).to_be_bytes());
    buf.extend_from_slice(data);
    while (buf.len() % 4) != 0 {
        buf.extend_from_slice(&[0]);
    }
}

fn append_message_integrity(buf: &mut BytesMut, key: &[u8]) {
    // position before adding MESSAGE-INTEGRITY attribute
    let attribute_start = buf.len();

    // append attribute header and placeholder value
    buf.extend_from_slice(&ATTR_MESSAGE_INTEGRITY.to_be_bytes());
    buf.extend_from_slice(&(20u16.to_be_bytes()));
    let value_start = buf.len();
    buf.extend_from_slice(&[0u8; 20]);

    // update message length to include the attribute (spec requires this before HMAC)
    let total_len = (buf.len() - 20) as u16;
    let len_bytes = total_len.to_be_bytes();
    buf[2] = len_bytes[0];
    buf[3] = len_bytes[1];

    // compute the HMAC over all bytes preceding the attribute (RFC 5389 §15.4)
    let signed = compute_message_integrity(key, &buf[..attribute_start]);

    // write the computed MAC into the placeholder we appended above
    buf[value_start..value_start + 20].copy_from_slice(&signed[..20]);
}

/// Ensure builders produce parseable STUN messages.
pub fn parse(buf: &[u8]) -> niom_turn::models::stun::StunMessage {
    parse_message(buf).expect("valid stun message")
}