use crate::dissect::{Dissect, DissectResult};
use crate::Error;
#[derive(Debug, Clone)]
pub struct Udp<'a> {
data: &'a [u8],
}
impl<'a> Udp<'a> {
pub fn new(data: &'a [u8]) -> Result<Self, Error> {
if data.len() < 8 {
return Err(Error::truncated(8, data.len()));
}
Ok(Self { data })
}
pub fn src_port(&self) -> u16 {
u16::from_be_bytes([self.data[0], self.data[1]])
}
pub fn dst_port(&self) -> u16 {
u16::from_be_bytes([self.data[2], self.data[3]])
}
pub fn length(&self) -> u16 {
u16::from_be_bytes([self.data[4], self.data[5]])
}
pub fn checksum(&self) -> u16 {
u16::from_be_bytes([self.data[6], self.data[7]])
}
pub fn payload(&self) -> &'a [u8] {
if self.data.len() <= 8 {
&[]
} else {
&self.data[8..]
}
}
}
impl<'a> Dissect<'a> for Udp<'a> {
type Output = &'a [u8];
fn dissect(&self) -> DissectResult<Self::Output> {
Ok(self.payload())
}
fn data(&self) -> &'a [u8] {
self.data
}
fn header_len(&self) -> usize {
8
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_udp_new_valid() {
let data = vec![
0x00, 0x35, 0x00, 0x35, 0x00, 0x10, 0x00, 0x00, ];
let udp = Udp::new(&data).unwrap();
assert_eq!(udp.src_port(), 53);
assert_eq!(udp.dst_port(), 53);
}
#[test]
fn test_udp_new_truncated() {
let data = vec![0u8; 4];
let result = Udp::new(&data);
assert!(result.is_err());
}
#[test]
fn test_udp_payload() {
let mut data = vec![
0x00, 0x35, 0x01, 0xbb, 0x00, 0x10, 0x00, 0x00, ];
data.extend_from_slice(&[0xde, 0xad, 0xbe, 0xef]); let udp = Udp::new(&data).unwrap();
assert_eq!(udp.payload(), &[0xde, 0xad, 0xbe, 0xef]);
}
}