cross_socket/packet/
udp.rs

1use pnet::packet::Packet;
2use std::net::{IpAddr, SocketAddr};
3use crate::packet::ethernet::ETHERNET_HEADER_LEN;
4use crate::packet::ipv4::IPV4_HEADER_LEN;
5use crate::packet::ipv6::IPV6_HEADER_LEN;
6
7/// UDP Header Length
8pub const UDP_HEADER_LEN: usize = 8;
9/// UDP BASE Destination Port. Usually used for traceroute.
10pub const UDP_BASE_DST_PORT: u16 = 33435;
11
12/// UDP (IPv4) Minimum Packet Length
13pub const UDPV4_PACKET_LEN: usize = ETHERNET_HEADER_LEN + IPV4_HEADER_LEN + UDP_HEADER_LEN;
14/// UDP (IPv4) Minimum IP Packet Length
15pub const UDPV4_IP_PACKET_LEN: usize = IPV4_HEADER_LEN + UDP_HEADER_LEN;
16/// UDP (IPv6) Minimum Packet Length
17pub const UDPV6_PACKET_LEN: usize = ETHERNET_HEADER_LEN + IPV6_HEADER_LEN + UDP_HEADER_LEN;
18/// UDP (IPv6) Minimum IP Packet Length
19pub const UDPV6_IP_PACKET_LEN: usize = IPV6_HEADER_LEN + UDP_HEADER_LEN;
20
21/// Represents the UDP packet.
22#[derive(Clone, Debug, PartialEq)]
23pub struct UdpPacket {
24    /// Source Port
25    pub source: u16,
26    /// Destination Port
27    pub destination: u16,
28    /// Length
29    pub length: u16,
30    /// Checksum
31    pub checksum: u16,
32    /// Payload
33    pub payload: Vec<u8>,
34}
35
36impl UdpPacket {
37    pub(crate) fn from_pnet_packet(packet: &pnet::packet::udp::UdpPacket) -> UdpPacket {
38        UdpPacket {
39            source: packet.get_source(),
40            destination: packet.get_destination(),
41            length: packet.get_length(),
42            checksum: packet.get_checksum(),
43            payload: packet.payload().to_vec(),
44        }
45    }
46    /// Constructs a new UdpPacket from bytes
47    pub fn from_bytes(packet: &[u8]) -> UdpPacket {
48        let udp_packet = pnet::packet::udp::UdpPacket::new(packet).unwrap();
49        UdpPacket::from_pnet_packet(&udp_packet)
50    }
51}
52
53pub(crate) fn build_udp_packet(
54    udp_packet: &mut pnet::packet::udp::MutableUdpPacket,
55    src_ip: IpAddr,
56    src_port: u16,
57    dst_ip: IpAddr,
58    dst_port: u16,
59) {
60    udp_packet.set_length(8);
61    udp_packet.set_source(src_port);
62    udp_packet.set_destination(dst_port);
63    match src_ip {
64        IpAddr::V4(src_ip) => match dst_ip {
65            IpAddr::V4(dst_ip) => {
66                let checksum =
67                    pnet::packet::udp::ipv4_checksum(&udp_packet.to_immutable(), &src_ip, &dst_ip);
68                udp_packet.set_checksum(checksum);
69            }
70            IpAddr::V6(_) => {}
71        },
72        IpAddr::V6(src_ip) => match dst_ip {
73            IpAddr::V4(_) => {}
74            IpAddr::V6(dst_ip) => {
75                let checksum =
76                    pnet::packet::udp::ipv6_checksum(&udp_packet.to_immutable(), &src_ip, &dst_ip);
77                udp_packet.set_checksum(checksum);
78            }
79        },
80    }
81}
82
83/// UDP Packet Builder
84#[derive(Clone, Debug)]
85pub struct UdpPacketBuilder {
86    /// Source IP address
87    pub src_ip: IpAddr,
88    /// Source Port
89    pub src_port: u16,
90    /// Destination IP address
91    pub dst_ip: IpAddr,
92    /// Destination Port
93    pub dst_port: u16,
94    /// Payload
95    pub payload: Vec<u8>,
96}
97
98impl UdpPacketBuilder {
99    /// Constructs a new UdpPacketBuilder
100    pub fn new(src_addr: SocketAddr, dst_addr: SocketAddr) -> Self {
101        UdpPacketBuilder {
102            src_ip: src_addr.ip(),
103            src_port: src_addr.port(),
104            dst_ip: dst_addr.ip(),
105            dst_port: dst_addr.port(),
106            payload: Vec::new(),
107        }
108    }
109    /// Builds a new UdpPacket and return bytes
110    pub fn build(&self) -> Vec<u8> {
111        let mut buffer: Vec<u8> = vec![0; UDP_HEADER_LEN + self.payload.len()];
112        let mut udp_packet = pnet::packet::udp::MutableUdpPacket::new(&mut buffer).unwrap();
113        udp_packet.set_source(self.src_port);
114        udp_packet.set_destination(self.dst_port);
115        if self.payload.len() > 0 {
116            udp_packet.set_payload(&self.payload);
117        }
118        udp_packet.set_length(UDP_HEADER_LEN as u16 + self.payload.len() as u16);
119        match self.src_ip {
120            IpAddr::V4(src_ip) => match self.dst_ip {
121                IpAddr::V4(dst_ip) => {
122                    let checksum = pnet::packet::udp::ipv4_checksum(
123                        &udp_packet.to_immutable(),
124                        &src_ip,
125                        &dst_ip,
126                    );
127                    udp_packet.set_checksum(checksum);
128                }
129                IpAddr::V6(_) => {}
130            },
131            IpAddr::V6(src_ip) => match self.dst_ip {
132                IpAddr::V4(_) => {}
133                IpAddr::V6(dst_ip) => {
134                    let checksum = pnet::packet::udp::ipv6_checksum(
135                        &udp_packet.to_immutable(),
136                        &src_ip,
137                        &dst_ip,
138                    );
139                    udp_packet.set_checksum(checksum);
140                }
141            },
142        }
143        udp_packet.packet().to_vec()
144    }
145}