use num_enum::{IntoPrimitive, TryFromPrimitive};
use super::{Opcode, Version};
#[derive(Debug, PartialEq, Eq)]
pub enum Request {
ExternalAddress,
Mapping {
proto: MapProtocol,
local_port: u16,
external_port: u16,
lifetime_seconds: u32,
},
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, IntoPrimitive, TryFromPrimitive)]
#[repr(u8)]
pub enum MapProtocol {
Udp = 1,
Tcp = 2,
}
impl Request {
pub fn encode(&self) -> Vec<u8> {
match self {
Request::ExternalAddress => vec![
Version::NatPmp.into(),
Opcode::DetermineExternalAddress.into(),
],
Request::Mapping {
proto,
local_port,
external_port,
lifetime_seconds,
} => {
let opcode = match proto {
MapProtocol::Udp => Opcode::MapUdp,
MapProtocol::Tcp => Opcode::MapTcp,
};
let mut buf = vec![Version::NatPmp.into(), opcode.into()];
buf.push(0); buf.push(0); buf.extend_from_slice(&local_port.to_be_bytes());
buf.extend_from_slice(&external_port.to_be_bytes());
buf.extend_from_slice(&lifetime_seconds.to_be_bytes());
buf
}
}
}
#[cfg(test)]
fn random<R: rand::Rng>(opcode: super::Opcode, rng: &mut R) -> Self {
use rand::RngExt;
match opcode {
Opcode::DetermineExternalAddress => Request::ExternalAddress,
Opcode::MapUdp => Request::Mapping {
proto: MapProtocol::Udp,
local_port: rng.random(),
external_port: rng.random(),
lifetime_seconds: rng.random(),
},
Opcode::MapTcp => Request::Mapping {
proto: MapProtocol::Tcp,
local_port: rng.random(),
external_port: rng.random(),
lifetime_seconds: rng.random(),
},
}
}
#[cfg(test)]
fn decode_map(buf: &[u8], proto: MapProtocol) -> Request {
let local_port_bytes = buf[4..6].try_into().expect("slice has the right size");
let local_port = u16::from_be_bytes(local_port_bytes);
let external_port_bytes = buf[6..8].try_into().expect("slice has the right size");
let external_port = u16::from_be_bytes(external_port_bytes);
let lifetime_bytes: [u8; 4] = buf[8..12].try_into().unwrap();
let lifetime_seconds = u32::from_be_bytes(lifetime_bytes);
Request::Mapping {
proto,
local_port,
external_port,
lifetime_seconds,
}
}
#[cfg(test)]
#[track_caller]
fn decode(buf: &[u8]) -> Self {
let _version: Version = buf[0].try_into().unwrap();
let opcode: super::Opcode = buf[1].try_into().unwrap();
match opcode {
Opcode::DetermineExternalAddress => Request::ExternalAddress,
Opcode::MapUdp => Self::decode_map(buf, MapProtocol::Udp),
Opcode::MapTcp => Self::decode_map(buf, MapProtocol::Tcp),
}
}
}
#[cfg(test)]
mod tests {
use rand::SeedableRng;
use super::*;
#[test]
fn test_encode_decode_addr_request() {
let mut rng = rand_chacha::ChaCha8Rng::seed_from_u64(42);
let request = Request::random(super::Opcode::DetermineExternalAddress, &mut rng);
let encoded = request.encode();
assert_eq!(request, Request::decode(&encoded));
}
#[test]
fn test_encode_decode_map_request() {
let mut rng = rand_chacha::ChaCha8Rng::seed_from_u64(42);
let request = Request::random(super::Opcode::MapUdp, &mut rng);
let encoded = request.encode();
assert_eq!(request, Request::decode(&encoded));
}
}