use core::convert::TryFrom;
use bytes::{Buf, BytesMut};
use smolsocket::port_from_bytes;
use crate::field::Field;
use super::{addr::field_port, field, Addr, Decoder, Encodable, Encoder, Error, HasAddr, Result};
#[derive(Debug, PartialEq, Clone)]
pub struct Packet<T: AsRef<[u8]>>(HasAddr<T>);
impl<T: AsRef<[u8]>> Packet<T> {
pub fn new_unchecked(buffer: T) -> Packet<T> {
Packet(HasAddr::new_unchecked(field::ATYP, buffer))
}
pub fn new_checked(buffer: T) -> Result<Packet<T>> {
let packet = Self::new_unchecked(buffer);
packet.check_header_len()?;
Ok(packet)
}
fn buffer_ref(&self) -> &[u8] {
self.0.buffer.as_ref()
}
pub fn check_header_len(&self) -> Result<()> {
self.0.check_addr_len()
}
#[inline]
pub fn header_len(&self) -> usize {
self.0.len_to_port()
}
#[inline]
fn field_data(&self) -> Field {
let start = self.0.field_port().end;
start..self.0.buffer.as_ref().len()
}
#[inline]
pub fn rsv(&self) -> u16 {
let data = self.buffer_ref();
let rsv_bytes = &data[field::UDP_RSV];
port_from_bytes(rsv_bytes[0], rsv_bytes[1])
}
#[inline]
pub fn frag(&self) -> u8 {
let data = self.buffer_ref();
data[field::UDP_FRAG]
}
#[inline]
pub fn atyp(&self) -> u8 {
self.0.atyp()
}
#[inline]
pub fn port(&self) -> u16 {
self.0.port()
}
}
impl<'a, T: AsRef<[u8]> + ?Sized> Packet<&'a T> {
#[inline]
pub fn addr(&self) -> &'a [u8] {
self.0.addr()
}
#[inline]
pub fn socks_addr(&self) -> &'a [u8] {
self.0.socks_addr()
}
#[inline]
pub fn data(&self) -> &'a [u8] {
let field = self.field_data();
let data = self.0.buffer.as_ref();
&data[field]
}
}
impl<T: AsRef<[u8]> + AsMut<[u8]>> Packet<T> {
fn buffer_mut(&mut self) -> &mut [u8] {
self.0.buffer.as_mut()
}
#[inline]
pub fn set_frag(&mut self, value: u8) {
let data = self.buffer_mut();
data[field::UDP_FRAG] = value;
}
#[inline]
pub fn set_atyp(&mut self, value: u8) {
self.0.set_atyp(value)
}
#[inline]
pub fn set_addr(&mut self, value: &[u8]) {
self.0.set_addr(value)
}
#[inline]
pub fn set_port(&mut self, value: u16) {
self.0.set_port(value)
}
#[inline]
pub fn set_socks_addr(&mut self, value: &[u8]) {
self.0.set_socks_addr(value)
}
#[inline]
pub fn addr_mut(&mut self) -> &mut [u8] {
self.0.addr_mut()
}
#[inline]
pub fn socks_addr_mut(&mut self) -> &mut [u8] {
self.0.socks_addr_mut()
}
#[inline]
pub fn data_mut(&mut self) -> &mut [u8] {
let field = self.field_data();
let data = self.buffer_mut();
&mut data[field]
}
}
impl<T: AsRef<[u8]>> AsRef<[u8]> for Packet<T> {
#[inline]
fn as_ref(&self) -> &[u8] {
self.buffer_ref()
}
}
#[derive(Debug, PartialEq, Eq, Clone)]
pub struct Repr {
pub frag: u8,
pub addr: Addr,
pub payload_len: usize,
}
impl Repr {
pub fn parse<T: AsRef<[u8]> + ?Sized>(packet: &Packet<&T>) -> Result<Repr> {
if packet.rsv() != 0 as u16 {
return Err(Error::Malformed);
}
let frag = packet.as_ref()[field::UDP_FRAG];
if frag > 127 {
return Err(Error::Malformed);
}
packet.check_header_len()?;
Ok(Repr {
frag,
addr: Addr::try_from(packet.socks_addr())?,
payload_len: packet.as_ref().len() - packet.header_len(),
})
}
pub fn buffer_len(&self) -> usize {
let addr_len = self.addr.addr_len();
field_port(field::ADDR_PORT.start, addr_len).end + self.payload_len
}
pub fn emit<T: AsRef<[u8]> + AsMut<[u8]>>(&self, packet: &mut Packet<T>) {
packet.set_frag(self.frag);
packet.set_socks_addr(&self.addr.to_vec());
}
}
impl Decoder<Repr> for Repr {
fn decode(src: &mut BytesMut) -> Result<Option<Self>> {
let pkt = Packet::new_unchecked(src.as_ref());
match Repr::parse(&pkt) {
Ok(repr) => {
src.advance(repr.buffer_len());
Ok(Some(repr))
}
Err(Error::Truncated) => Ok(None),
Err(err) => Err(err),
}
}
}
impl Encodable for Repr {
fn try_encode(&self, dst: &mut BytesMut) -> Result<()> {
if dst.len() < self.buffer_len() {
dst.resize(self.buffer_len(), 0);
}
let mut pkt = Packet::new_unchecked(dst);
self.emit(&mut pkt);
Ok(())
}
}
impl Encoder<Repr> for Repr {
fn encode(item: &Repr, dst: &mut BytesMut) -> Result<()> {
item.try_encode(dst)
}
}
#[cfg(test)]
mod tests {
use bytes::BytesMut;
#[cfg(any(feature = "proto-ipv4", feature = "proto-ipv6"))]
use smolsocket::SocketAddr;
#[cfg(feature = "proto-ipv4")]
use smoltcp::wire::Ipv4Address;
use crate::Atyp;
use super::*;
#[cfg(feature = "proto-ipv4")]
#[test]
fn test_udp_invalid_len() {
let mut truncated_bytes = vec![0x00 as u8; 4];
let mut truncated = Packet::new_unchecked(&mut truncated_bytes);
truncated.set_frag(0);
truncated.set_atyp(Atyp::V4 as u8);
assert_eq!(truncated.check_header_len(), Err(Error::Truncated));
let mut truncated_bytes_mut = BytesMut::new();
truncated_bytes_mut.extend(truncated_bytes);
assert_eq!(Repr::decode(&mut truncated_bytes_mut), Ok(None));
let mut truncated_bytes = vec![0x00 as u8; 5];
let mut truncated = Packet::new_unchecked(&mut truncated_bytes);
truncated.set_frag(0);
truncated.set_atyp(Atyp::V4 as u8);
assert_eq!(truncated.check_header_len(), Err(Error::Truncated));
assert_eq!(truncated.header_len(), 10);
let mut malformed_bytes = vec![0x00 as u8; truncated.header_len()];
let mut malformed = Packet::new_unchecked(&mut malformed_bytes);
malformed.set_frag(0);
malformed.set_atyp(0);
assert_eq!(malformed.check_header_len(), Err(Error::Malformed));
let mut malformed_bytes_mut = BytesMut::new();
malformed_bytes_mut.extend(malformed_bytes);
assert_eq!(
Repr::decode(&mut malformed_bytes_mut),
Err(Error::Malformed)
);
}
#[cfg(feature = "proto-ipv4")]
#[test]
fn test_data_len_0() {
let socket_addr = SocketAddr::new_ip4_port(127, 0, 0, 1, 80);
let addr = Addr::SocketAddr(socket_addr);
let repr = Repr {
frag: 0,
addr: addr.clone(),
payload_len: 0,
};
assert_eq!(repr.addr.total_len(), 7);
assert_eq!(repr.buffer_len(), 10);
let mut bytes = vec![0x00 as u8; repr.buffer_len()];
let mut pkt = Packet::new_unchecked(&mut bytes[..]);
assert_eq!(pkt.frag(), 0);
pkt.set_frag(1);
assert_eq!(pkt.frag(), 1);
assert_eq!(pkt.atyp(), 0);
pkt.set_atyp(Atyp::V4 as u8);
assert_eq!(pkt.atyp(), Atyp::V4 as u8);
assert_eq!(&pkt.addr_mut(), &Ipv4Address::new(0, 0, 0, 0).as_bytes());
pkt.set_addr(Ipv4Address::new(192, 168, 0, 1).as_bytes());
assert_eq!(
&pkt.addr_mut(),
&Ipv4Address::new(192, 168, 0, 1).as_bytes()
);
assert_eq!(
&Packet::new_unchecked(pkt.as_ref()).addr(),
&Ipv4Address::new(192, 168, 0, 1).as_bytes()
);
assert_eq!(pkt.port(), 0);
pkt.set_port(8080);
assert_eq!(pkt.port(), 8080);
repr.emit(&mut pkt);
assert_eq!(pkt.field_data(), 10..10);
assert_eq!(pkt.data_mut().len(), 0);
assert_eq!(Packet::new_checked(pkt.as_ref()).unwrap().data().len(), 0);
assert_eq!(pkt.atyp(), Atyp::V4 as u8);
assert_eq!(&pkt.addr_mut(), &Ipv4Address::new(127, 0, 0, 1).as_bytes());
assert_eq!(pkt.port(), 80);
assert_eq!(&pkt.socks_addr_mut(), &addr.to_vec().as_slice());
let parsed = Repr::parse(&Packet::new_checked(pkt.buffer_ref()).unwrap()).unwrap();
assert_eq!(parsed, repr);
let mut bytes_mut = BytesMut::new();
assert!(Repr::encode(&repr, &mut bytes_mut).is_ok());
let decoded = Repr::decode(&mut bytes_mut);
assert_eq!(decoded, Ok(Some(repr)));
}
#[cfg(feature = "proto-ipv4")]
#[test]
fn test_data_len_1() {
let socket_addr = SocketAddr::new_ip4_port(127, 0, 0, 1, 80);
let addr = Addr::SocketAddr(socket_addr);
let repr = Repr {
frag: 0,
addr: addr.clone(),
payload_len: 1,
};
assert_eq!(repr.addr.total_len(), 7);
assert_eq!(repr.buffer_len(), 11);
let mut bytes = vec![0x00 as u8; repr.buffer_len()];
let mut pkt = Packet::new_unchecked(&mut bytes[..]);
assert_eq!(pkt.frag(), 0);
pkt.set_frag(1);
assert_eq!(pkt.frag(), 1);
assert_eq!(pkt.atyp(), 0);
pkt.set_atyp(Atyp::V4 as u8);
assert_eq!(pkt.atyp(), Atyp::V4 as u8);
assert_eq!(&pkt.addr_mut(), &Ipv4Address::new(0, 0, 0, 0).as_bytes());
pkt.set_addr(Ipv4Address::new(192, 168, 0, 1).as_bytes());
assert_eq!(
&pkt.addr_mut(),
&Ipv4Address::new(192, 168, 0, 1).as_bytes()
);
assert_eq!(
&Packet::new_unchecked(pkt.as_ref()).addr(),
&Ipv4Address::new(192, 168, 0, 1).as_bytes()
);
assert_eq!(pkt.port(), 0);
pkt.set_port(8080);
assert_eq!(pkt.port(), 8080);
repr.emit(&mut pkt);
assert_eq!(pkt.field_data(), 10..11);
assert_eq!(pkt.data_mut().len(), 1);
assert_eq!(Packet::new_checked(pkt.as_ref()).unwrap().data().len(), 1);
assert_eq!(pkt.atyp(), Atyp::V4 as u8);
assert_eq!(&pkt.addr_mut(), &Ipv4Address::new(127, 0, 0, 1).as_bytes());
assert_eq!(pkt.port(), 80);
assert_eq!(&pkt.socks_addr_mut(), &addr.to_vec().as_slice());
let parsed = Repr::parse(&Packet::new_checked(pkt.buffer_ref()).unwrap()).unwrap();
assert_eq!(parsed, repr);
let mut bytes_mut = BytesMut::new();
assert!(Repr::encode(&repr, &mut bytes_mut).is_ok());
let decoded = Repr::decode(&mut bytes_mut);
assert_eq!(decoded, Ok(Some(repr)));
}
}