#![allow(clippy::cast_possible_truncation)]
use std::net::IpAddr;
use hickory_proto::op::Message;
use hickory_proto::serialize::binary::BinEncodable;
use pnet::packet::ethernet::{EtherTypes, MutableEthernetPacket};
use pnet::packet::ip::IpNextHeaderProtocols;
use pnet::packet::ipv4::MutableIpv4Packet;
use pnet::packet::ipv6::MutableIpv6Packet;
use pnet::packet::udp::MutableUdpPacket;
use super::buffer::BufferPool;
use super::capture::PacketInfo;
use crate::error::{NetworkError, Result};
const ETHERNET_HEADER_SIZE: usize = 14;
const IPV4_HEADER_SIZE: usize = 20;
const IPV6_HEADER_SIZE: usize = 40;
const UDP_HEADER_SIZE: usize = 8;
pub struct PacketBuilder {
buffer_pool: BufferPool,
}
impl PacketBuilder {
pub const fn new(buffer_pool: BufferPool) -> Self {
Self { buffer_pool }
}
pub fn build_response(&self, response: &Message, packet_info: &PacketInfo) -> Result<Vec<u8>> {
let dns_bytes = response
.to_bytes()
.map_err(|e| NetworkError::PacketConstruction(e.to_string()))?;
match (packet_info.source_ip, packet_info.dest_ip) {
(IpAddr::V4(src_ip), IpAddr::V4(dst_ip)) => {
self.build_ipv4_response(&dns_bytes, packet_info, src_ip, dst_ip)
}
(IpAddr::V6(src_ip), IpAddr::V6(dst_ip)) => {
self.build_ipv6_response(&dns_bytes, packet_info, src_ip, dst_ip)
}
_ => Err(NetworkError::PacketConstruction("mixed IPv4/IPv6 addresses".into()).into()),
}
}
fn build_ipv4_response(
&self,
dns_bytes: &[u8],
packet_info: &PacketInfo,
src_ip: std::net::Ipv4Addr,
dst_ip: std::net::Ipv4Addr,
) -> Result<Vec<u8>> {
let udp_len = UDP_HEADER_SIZE + dns_bytes.len();
let ipv4_len = IPV4_HEADER_SIZE + udp_len;
let total_len = ETHERNET_HEADER_SIZE + ipv4_len;
let mut buffer = self.buffer_pool.get_zeroed(total_len);
let packet_data = buffer.as_mut_slice();
let udp_start = ETHERNET_HEADER_SIZE + IPV4_HEADER_SIZE;
{
let mut udp =
MutableUdpPacket::new(&mut packet_data[udp_start..]).ok_or_else(|| {
NetworkError::PacketConstruction("UDP packet creation failed".into())
})?;
udp.set_source(packet_info.dest_port);
udp.set_destination(packet_info.source_port);
udp.set_length(udp_len as u16);
udp.set_payload(dns_bytes);
}
let ipv4_start = ETHERNET_HEADER_SIZE;
{
let mut ipv4 =
MutableIpv4Packet::new(&mut packet_data[ipv4_start..]).ok_or_else(|| {
NetworkError::PacketConstruction("IPv4 packet creation failed".into())
})?;
ipv4.set_version(4);
ipv4.set_header_length(5);
ipv4.set_total_length(ipv4_len as u16);
ipv4.set_ttl(64);
ipv4.set_next_level_protocol(IpNextHeaderProtocols::Udp);
ipv4.set_source(dst_ip);
ipv4.set_destination(src_ip);
ipv4.set_checksum(pnet::packet::ipv4::checksum(&ipv4.to_immutable()));
}
{
let udp_immutable = pnet::packet::udp::UdpPacket::new(&packet_data[udp_start..])
.ok_or_else(|| NetworkError::PacketConstruction("UDP checksum failed".into()))?;
let checksum = pnet::packet::udp::ipv4_checksum(&udp_immutable, &dst_ip, &src_ip);
let mut udp = MutableUdpPacket::new(&mut packet_data[udp_start..]).unwrap();
udp.set_checksum(checksum);
}
{
let mut ethernet = MutableEthernetPacket::new(packet_data).ok_or_else(|| {
NetworkError::PacketConstruction("Ethernet frame creation failed".into())
})?;
ethernet.set_destination(packet_info.source_mac);
ethernet.set_source(packet_info.dest_mac);
ethernet.set_ethertype(EtherTypes::Ipv4);
}
Ok(packet_data.to_vec())
}
fn build_ipv6_response(
&self,
dns_bytes: &[u8],
packet_info: &PacketInfo,
src_ip: std::net::Ipv6Addr,
dst_ip: std::net::Ipv6Addr,
) -> Result<Vec<u8>> {
let udp_len = UDP_HEADER_SIZE + dns_bytes.len();
let total_len = ETHERNET_HEADER_SIZE + IPV6_HEADER_SIZE + udp_len;
let mut buffer = self.buffer_pool.get_zeroed(total_len);
let packet_data = buffer.as_mut_slice();
let udp_start = ETHERNET_HEADER_SIZE + IPV6_HEADER_SIZE;
{
let mut udp =
MutableUdpPacket::new(&mut packet_data[udp_start..]).ok_or_else(|| {
NetworkError::PacketConstruction("UDP packet creation failed".into())
})?;
udp.set_source(packet_info.dest_port);
udp.set_destination(packet_info.source_port);
udp.set_length(udp_len as u16);
udp.set_payload(dns_bytes);
}
let ipv6_start = ETHERNET_HEADER_SIZE;
{
let mut ipv6 =
MutableIpv6Packet::new(&mut packet_data[ipv6_start..]).ok_or_else(|| {
NetworkError::PacketConstruction("IPv6 packet creation failed".into())
})?;
ipv6.set_version(6);
ipv6.set_payload_length(udp_len as u16);
ipv6.set_next_header(IpNextHeaderProtocols::Udp);
ipv6.set_hop_limit(64);
ipv6.set_source(dst_ip);
ipv6.set_destination(src_ip);
}
{
let udp_immutable = pnet::packet::udp::UdpPacket::new(&packet_data[udp_start..])
.ok_or_else(|| NetworkError::PacketConstruction("UDP checksum failed".into()))?;
let checksum = pnet::packet::udp::ipv6_checksum(&udp_immutable, &dst_ip, &src_ip);
let mut udp = MutableUdpPacket::new(&mut packet_data[udp_start..]).unwrap();
udp.set_checksum(checksum);
}
{
let mut ethernet = MutableEthernetPacket::new(packet_data).ok_or_else(|| {
NetworkError::PacketConstruction("Ethernet frame creation failed".into())
})?;
ethernet.set_destination(packet_info.source_mac);
ethernet.set_source(packet_info.dest_mac);
ethernet.set_ethertype(EtherTypes::Ipv6);
}
Ok(packet_data.to_vec())
}
}
#[cfg(test)]
mod tests {
use super::*;
use hickory_proto::op::{MessageType, OpCode, ResponseCode};
use pnet::packet::Packet;
use pnet::packet::ethernet::EthernetPacket;
use pnet::packet::ipv4::Ipv4Packet;
use pnet::packet::udp::UdpPacket;
use pnet::util::MacAddr;
use std::net::{Ipv4Addr, Ipv6Addr};
fn create_test_response() -> Message {
let mut msg = Message::new();
msg.set_id(1234)
.set_message_type(MessageType::Response)
.set_op_code(OpCode::Query)
.set_response_code(ResponseCode::NoError);
msg
}
fn create_ipv4_packet_info() -> PacketInfo {
PacketInfo {
source_mac: MacAddr::new(0x11, 0x22, 0x33, 0x44, 0x55, 0x66),
dest_mac: MacAddr::new(0xaa, 0xbb, 0xcc, 0xdd, 0xee, 0xff),
source_ip: IpAddr::V4(Ipv4Addr::new(192, 168, 1, 100)),
dest_ip: IpAddr::V4(Ipv4Addr::new(192, 168, 1, 1)),
source_port: 12345,
dest_port: 53,
}
}
fn create_ipv6_packet_info() -> PacketInfo {
PacketInfo {
source_mac: MacAddr::new(0x11, 0x22, 0x33, 0x44, 0x55, 0x66),
dest_mac: MacAddr::new(0xaa, 0xbb, 0xcc, 0xdd, 0xee, 0xff),
source_ip: IpAddr::V6(Ipv6Addr::new(0xfe80, 0, 0, 0, 0, 0, 0, 1)),
dest_ip: IpAddr::V6(Ipv6Addr::new(0xfe80, 0, 0, 0, 0, 0, 0, 2)),
source_port: 12345,
dest_port: 53,
}
}
#[test]
fn should_build_valid_ipv4_response_packet() {
let pool = BufferPool::new(4);
let builder = PacketBuilder::new(pool);
let response = create_test_response();
let packet_info = create_ipv4_packet_info();
let packet = builder.build_response(&response, &packet_info).unwrap();
let eth = EthernetPacket::new(&packet).unwrap();
assert_eq!(eth.get_destination(), packet_info.source_mac);
assert_eq!(eth.get_source(), packet_info.dest_mac);
assert_eq!(eth.get_ethertype(), EtherTypes::Ipv4);
let ipv4 = Ipv4Packet::new(eth.payload()).unwrap();
assert_eq!(ipv4.get_version(), 4);
assert_eq!(
ipv4.get_source(),
match packet_info.dest_ip {
IpAddr::V4(ip) => ip,
IpAddr::V6(_) => panic!(),
}
);
assert_eq!(
ipv4.get_destination(),
match packet_info.source_ip {
IpAddr::V4(ip) => ip,
IpAddr::V6(_) => panic!(),
}
);
let udp = UdpPacket::new(ipv4.payload()).unwrap();
assert_eq!(udp.get_source(), packet_info.dest_port);
assert_eq!(udp.get_destination(), packet_info.source_port);
}
#[test]
fn should_build_valid_ipv6_response_packet() {
let pool = BufferPool::new(4);
let builder = PacketBuilder::new(pool);
let response = create_test_response();
let packet_info = create_ipv6_packet_info();
let packet = builder.build_response(&response, &packet_info).unwrap();
let eth = EthernetPacket::new(&packet).unwrap();
assert_eq!(eth.get_ethertype(), EtherTypes::Ipv6);
assert_eq!(eth.get_destination(), packet_info.source_mac);
assert_eq!(eth.get_source(), packet_info.dest_mac);
}
#[test]
fn should_return_error_for_mixed_ip_versions() {
let pool = BufferPool::new(4);
let builder = PacketBuilder::new(pool);
let response = create_test_response();
let packet_info = PacketInfo {
source_mac: MacAddr::new(0x11, 0x22, 0x33, 0x44, 0x55, 0x66),
dest_mac: MacAddr::new(0xaa, 0xbb, 0xcc, 0xdd, 0xee, 0xff),
source_ip: IpAddr::V4(Ipv4Addr::new(192, 168, 1, 100)),
dest_ip: IpAddr::V6(Ipv6Addr::LOCALHOST),
source_port: 12345,
dest_port: 53,
};
let result = builder.build_response(&response, &packet_info);
assert!(result.is_err());
}
}