Files
remoteingress/rust/crates/remoteingress-core/src/stun.rs

265 lines
9.0 KiB
Rust

use std::net::Ipv4Addr;
use tokio::net::UdpSocket;
use tokio::time::{timeout, Duration};
const STUN_SERVER: &str = "stun.cloudflare.com:3478";
const STUN_TIMEOUT: Duration = Duration::from_secs(3);
// STUN constants
const STUN_BINDING_REQUEST: u16 = 0x0001;
const STUN_MAGIC_COOKIE: u32 = 0x2112A442;
const ATTR_XOR_MAPPED_ADDRESS: u16 = 0x0020;
const ATTR_MAPPED_ADDRESS: u16 = 0x0001;
/// Discover our public IP via STUN Binding Request (RFC 5389).
/// Returns `None` on timeout or parse failure.
pub async fn discover_public_ip() -> Option<String> {
discover_public_ip_from(STUN_SERVER).await
}
pub async fn discover_public_ip_from(server: &str) -> Option<String> {
let result = timeout(STUN_TIMEOUT, async {
let socket = UdpSocket::bind("0.0.0.0:0").await.ok()?;
socket.connect(server).await.ok()?;
// Build STUN Binding Request (20 bytes)
let mut request = [0u8; 20];
// Message Type: Binding Request (0x0001)
request[0..2].copy_from_slice(&STUN_BINDING_REQUEST.to_be_bytes());
// Message Length: 0 (no attributes)
request[2..4].copy_from_slice(&0u16.to_be_bytes());
// Magic Cookie
request[4..8].copy_from_slice(&STUN_MAGIC_COOKIE.to_be_bytes());
// Transaction ID: 12 random bytes
let txn_id: [u8; 12] = rand_bytes();
request[8..20].copy_from_slice(&txn_id);
socket.send(&request).await.ok()?;
let mut buf = [0u8; 512];
let n = socket.recv(&mut buf).await.ok()?;
if n < 20 {
return None;
}
parse_stun_response(&buf[..n], &txn_id)
})
.await;
match result {
Ok(ip) => ip,
Err(_) => None, // timeout
}
}
fn parse_stun_response(data: &[u8], _txn_id: &[u8; 12]) -> Option<String> {
if data.len() < 20 {
return None;
}
// Verify it's a Binding Response (0x0101)
let msg_type = u16::from_be_bytes([data[0], data[1]]);
if msg_type != 0x0101 {
return None;
}
let msg_len = u16::from_be_bytes([data[2], data[3]]) as usize;
let magic = u32::from_be_bytes([data[4], data[5], data[6], data[7]]);
// Parse attributes
let attrs = &data[20..std::cmp::min(20 + msg_len, data.len())];
let mut offset = 0;
while offset + 4 <= attrs.len() {
let attr_type = u16::from_be_bytes([attrs[offset], attrs[offset + 1]]);
let attr_len = u16::from_be_bytes([attrs[offset + 2], attrs[offset + 3]]) as usize;
offset += 4;
if offset + attr_len > attrs.len() {
break;
}
let attr_data = &attrs[offset..offset + attr_len];
match attr_type {
ATTR_XOR_MAPPED_ADDRESS if attr_data.len() >= 8 => {
let family = attr_data[1];
if family == 0x01 {
// IPv4
let port_xored = u16::from_be_bytes([attr_data[2], attr_data[3]]);
let _port = port_xored ^ (STUN_MAGIC_COOKIE >> 16) as u16;
let ip_xored = u32::from_be_bytes([
attr_data[4],
attr_data[5],
attr_data[6],
attr_data[7],
]);
let ip = ip_xored ^ magic;
return Some(Ipv4Addr::from(ip).to_string());
}
}
ATTR_MAPPED_ADDRESS if attr_data.len() >= 8 => {
let family = attr_data[1];
if family == 0x01 {
// IPv4 (non-XOR fallback)
let ip = u32::from_be_bytes([
attr_data[4],
attr_data[5],
attr_data[6],
attr_data[7],
]);
return Some(Ipv4Addr::from(ip).to_string());
}
}
_ => {}
}
// Pad to 4-byte boundary
offset += (attr_len + 3) & !3;
}
None
}
#[cfg(test)]
mod tests {
use super::*;
/// Build a synthetic STUN Binding Response with given attributes.
fn build_stun_response(attrs: &[(u16, &[u8])]) -> Vec<u8> {
let mut attrs_bytes = Vec::new();
for &(attr_type, attr_data) in attrs {
attrs_bytes.extend_from_slice(&attr_type.to_be_bytes());
attrs_bytes.extend_from_slice(&(attr_data.len() as u16).to_be_bytes());
attrs_bytes.extend_from_slice(attr_data);
// Pad to 4-byte boundary
let pad = (4 - (attr_data.len() % 4)) % 4;
attrs_bytes.extend(std::iter::repeat(0u8).take(pad));
}
let mut response = Vec::new();
// msg_type = 0x0101 (Binding Response)
response.extend_from_slice(&0x0101u16.to_be_bytes());
// message length
response.extend_from_slice(&(attrs_bytes.len() as u16).to_be_bytes());
// magic cookie
response.extend_from_slice(&STUN_MAGIC_COOKIE.to_be_bytes());
// transaction ID (12 bytes)
response.extend_from_slice(&[0u8; 12]);
// attributes
response.extend_from_slice(&attrs_bytes);
response
}
#[test]
fn test_xor_mapped_address_ipv4() {
// IP 203.0.113.1 = 0xCB007101, XOR'd with magic 0x2112A442 = 0xEA12D543
let attr_data: [u8; 8] = [
0x00, 0x01, // reserved + family (IPv4)
0x11, 0x2B, // port XOR'd with 0x2112 (port 0x3039 = 12345)
0xEA, 0x12, 0xD5, 0x43, // IP XOR'd
];
let data = build_stun_response(&[(ATTR_XOR_MAPPED_ADDRESS, &attr_data)]);
let txn_id = [0u8; 12];
let result = parse_stun_response(&data, &txn_id);
assert_eq!(result, Some("203.0.113.1".to_string()));
}
#[test]
fn test_mapped_address_fallback_ipv4() {
// IP 192.168.1.1 = 0xC0A80101 (no XOR)
let attr_data: [u8; 8] = [
0x00, 0x01, // reserved + family (IPv4)
0x00, 0x50, // port 80
0xC0, 0xA8, 0x01, 0x01, // IP
];
let data = build_stun_response(&[(ATTR_MAPPED_ADDRESS, &attr_data)]);
let txn_id = [0u8; 12];
let result = parse_stun_response(&data, &txn_id);
assert_eq!(result, Some("192.168.1.1".to_string()));
}
#[test]
fn test_response_too_short() {
let data = vec![0u8; 19]; // < 20 bytes
let txn_id = [0u8; 12];
assert_eq!(parse_stun_response(&data, &txn_id), None);
}
#[test]
fn test_wrong_msg_type() {
// Build with correct helper then overwrite msg_type to 0x0001 (Binding Request)
let mut data = build_stun_response(&[]);
data[0] = 0x00;
data[1] = 0x01;
let txn_id = [0u8; 12];
assert_eq!(parse_stun_response(&data, &txn_id), None);
}
#[test]
fn test_no_mapped_address_attributes() {
// Valid response with no attributes
let data = build_stun_response(&[]);
let txn_id = [0u8; 12];
assert_eq!(parse_stun_response(&data, &txn_id), None);
}
#[test]
fn test_xor_preferred_over_mapped() {
// XOR gives 203.0.113.1, MAPPED gives 192.168.1.1
let xor_data: [u8; 8] = [
0x00, 0x01,
0x11, 0x2B,
0xEA, 0x12, 0xD5, 0x43,
];
let mapped_data: [u8; 8] = [
0x00, 0x01,
0x00, 0x50,
0xC0, 0xA8, 0x01, 0x01,
];
// XOR listed first — should be preferred
let data = build_stun_response(&[
(ATTR_XOR_MAPPED_ADDRESS, &xor_data),
(ATTR_MAPPED_ADDRESS, &mapped_data),
]);
let txn_id = [0u8; 12];
let result = parse_stun_response(&data, &txn_id);
assert_eq!(result, Some("203.0.113.1".to_string()));
}
#[test]
fn test_truncated_attribute_data() {
// Attribute claims 8 bytes but only 4 are present
let mut data = build_stun_response(&[]);
// Manually append a truncated XOR_MAPPED_ADDRESS attribute
let attr_type = ATTR_XOR_MAPPED_ADDRESS.to_be_bytes();
let attr_len = 8u16.to_be_bytes(); // claims 8 bytes
let truncated = [0x00, 0x01, 0x11, 0x2B]; // only 4 bytes
// Update message length
let new_msg_len = (attr_type.len() + attr_len.len() + truncated.len()) as u16;
data[2..4].copy_from_slice(&new_msg_len.to_be_bytes());
data.extend_from_slice(&attr_type);
data.extend_from_slice(&attr_len);
data.extend_from_slice(&truncated);
let txn_id = [0u8; 12];
// Should return None, not panic
assert_eq!(parse_stun_response(&data, &txn_id), None);
}
}
/// Generate 12 random bytes for transaction ID.
fn rand_bytes() -> [u8; 12] {
let mut bytes = [0u8; 12];
// Use a simple approach: mix timestamp + counter
let now = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default();
let nanos = now.as_nanos();
bytes[0..8].copy_from_slice(&(nanos as u64).to_le_bytes());
// Fill remaining with process-id based data
let pid = std::process::id();
bytes[8..12].copy_from_slice(&pid.to_le_bytes());
bytes
}