evan/protomask
1
Fork 0
Code

Add experimental profiling support

Browse Source
This commit is contained in:
Evan Pratten 2023-07-19 10:36:53 -04:00
parent c6db02c80a
commit edd83fad99
16 changed files with 100 additions and 22 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
Cargo.toml
View File

@ -12,10 +12,13 @@ license = "GPL-3.0"
keywords = []
categories = []
[features]
default = []
profiler = ["profiling/profile-with-puffin", "puffin_http", "puffin"]
# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
[dependencies]
protomask-tun = { path = "protomask-tun", version = "0.1.0" }
tokio = { version = "1.29.1", features = [
"macros",
"rt-multi-thread",
@ -25,6 +28,9 @@ tokio = { version = "1.29.1", features = [
clap = { version = "4.3.11", features = ["derive"] }
serde = { version = "1.0.171", features = ["derive"] }
ipnet = { version = "2.8.0", features = ["serde"] }
puffin_http = { version = "^0.13.0", optional = true }
puffin = { version = "0.12.1", optional = true }
profiling = "1.0.8"
toml = "0.7.6"
log = "0.4.19"
fern = "0.6.2"
@ -36,7 +42,7 @@ bimap = "0.6.3"
pnet_packet = "0.33.0"
rtnetlink = "0.13.0"
futures = "0.3.28"
protomask-tun = { path = "protomask-tun", version = "0.1.0" }
cfg-if = "1.0.0"
[[bin]]
name = "protomask"

11
src/cli/main.rs
View File

@ -1,5 +1,6 @@
//! This is the entrypoint for `protomask` from the command line.
use cfg_if::cfg_if;
use clap::Parser;
use config::Config;
use logging::enable_logger;
@ -26,6 +27,16 @@ pub async fn main() {
std::process::exit(1);
}
// Enable the profiler if enabled by build flags
cfg_if! {
if #[cfg(feature = "profiler")] {
let puffin_listen_addr = format!("[::]:{}", puffin_http::DEFAULT_PORT);
log::info!("Puffin HTTP server listening on: {}", puffin_listen_addr);
let _puffin_server = puffin_http::Server::new(&puffin_listen_addr).unwrap();
puffin::set_scopes_on(true);
}
}
// Create the NAT64 instance
let mut nat64 = Nat64::new(
config.nat64_prefix,

17
src/nat/error.rs Normal file
View File

@ -0,0 +1,17 @@
use super::table;
#[derive(Debug, thiserror::Error)]
pub enum Nat64Error {
#[error(transparent)]
TableError(#[from] table::TableError),
#[error(transparent)]
TunError(#[from] protomask_tun::Error),
#[error(transparent)]
IoError(#[from] std::io::Error),
#[error(transparent)]
PacketHandlingError(#[from] crate::packet::error::PacketError),
#[error(transparent)]
PacketReceiveError(#[from] tokio::sync::broadcast::error::RecvError),
#[error(transparent)]
PacketSendError(#[from] tokio::sync::mpsc::error::SendError<Vec<u8>>),
}

48
src/nat/mod.rs
View File

@ -4,6 +4,7 @@ use crate::packet::{
};
use self::{
error::Nat64Error,
table::Nat64Table,
utils::{embed_address, extract_address},
};
@ -15,27 +16,10 @@ use std::{
};
use tokio::sync::{broadcast, mpsc};
mod error;
mod table;
mod utils;
#[derive(Debug, thiserror::Error)]
pub enum Nat64Error {
#[error(transparent)]
TableError(#[from] table::TableError),
#[error(transparent)]
TunError(#[from] protomask_tun::Error),
#[error(transparent)]
IoError(#[from] std::io::Error),
// #[error(transparent)]
// XlatError(#[from] xlat::PacketTranslationError),
#[error(transparent)]
PacketHandlingError(#[from] crate::packet::error::PacketError),
#[error(transparent)]
PacketReceiveError(#[from] broadcast::error::RecvError),
#[error(transparent)]
PacketSendError(#[from] mpsc::error::SendError<Vec<u8>>),
}
pub struct Nat64 {
table: Nat64Table,
interface: TunDevice,
@ -81,6 +65,9 @@ impl Nat64 {
// Process packets in a loop
loop {
// Start a new profiler frame
profiling::finish_frame!();
// Try to read a packet
match rx.recv().await {
Ok(packet) => {
@ -90,6 +77,8 @@ impl Nat64 {
// Separate logic is needed for handling IPv4 vs IPv6 packets, so a check must be done here
match packet[0] >> 4 {
4 => {
profiling::scope!("Ipv4 to IPv6");
// Parse the packet
let packet: Ipv4Packet<Vec<u8>> = packet.try_into()?;
@ -113,14 +102,37 @@ impl Nat64 {
});
}
6 => {
profiling::scope!("Ipv6 to IPv4");
// Parse the packet
let packet: Ipv6Packet<Vec<u8>> = packet.try_into()?;
// Drop packets "coming from" the NAT64 prefix
if self.ipv6_nat_prefix.contains(&packet.source_address) {
log::warn!(
"Dropping packet \"from\" NAT64 prefix: {} -> {}",
packet.source_address,
packet.destination_address
);
continue;
}
// Get the new source and dest addresses
let new_source =
self.table.get_or_assign_ipv4(packet.source_address)?;
let new_destination = extract_address(packet.destination_address);
// Drop packets destined for private IPv4 addresses
if new_destination.is_private() {
log::warn!(
"Dropping packet destined for private IPv4 address: {} -> {} ({})",
packet.source_address,
packet.destination_address,
new_destination
);
continue;
}
// Spawn a task to process the packet
tokio::spawn(async move {
let output =

7
src/nat/table.rs
View File

@ -47,6 +47,7 @@ impl Nat64Table {
}
/// Make a reservation for an IP address pair for eternity
#[profiling::function]
pub fn add_infinite_reservation(
&mut self,
ipv6: Ipv6Addr,
@ -68,6 +69,7 @@ impl Nat64Table {
}
/// Check if a given address exists in the table
#[profiling::function]
pub fn contains(&self, address: &IpAddr) -> bool {
match address {
IpAddr::V4(ipv4) => self.reservations.contains_right(ipv4),
@ -76,6 +78,7 @@ impl Nat64Table {
}
/// Get or assign an IPv4 address for the given IPv6 address
#[profiling::function]
pub fn get_or_assign_ipv4(&mut self, ipv6: Ipv6Addr) -> Result<Ipv4Addr, TableError> {
// Prune old reservations
self.prune();
@ -110,6 +113,7 @@ impl Nat64Table {
}
/// Try to find an IPv6 address for the given IPv4 address
#[profiling::function]
pub fn get_reverse(&mut self, ipv4: Ipv4Addr) -> Result<Ipv6Addr, TableError> {
// Prune old reservations
self.prune();
@ -129,11 +133,13 @@ impl Nat64Table {
}
/// Check if an address is within the IPv4 pool
#[profiling::function]
pub fn is_address_within_pool(&self, address: &Ipv4Addr) -> bool {
self.ipv4_pool.iter().any(|net| net.contains(address))
}
/// Calculate the translated version of any address
#[profiling::function]
pub fn calculate_xlat_addr(
&mut self,
input: &IpAddr,
@ -186,6 +192,7 @@ impl Nat64Table {
impl Nat64Table {
/// Prune old reservations
#[profiling::function]
pub fn prune(&mut self) {
let now = Instant::now();

3
src/packet/protocols/icmp.rs
View File

@ -29,6 +29,7 @@ where
{
type Error = PacketError;
#[profiling::function]
fn try_from(bytes: Vec<u8>) -> Result<Self, Self::Error> {
// Parse the packet
let packet =
@ -43,11 +44,11 @@ where
}
}
impl<T> Into<Vec<u8>> for IcmpPacket<T>
where
T: Into<Vec<u8>>,
{
#[profiling::function]
fn into(self) -> Vec<u8> {
// Convert the payload into raw bytes
let payload: Vec<u8> = self.payload.into();

3
src/packet/protocols/icmpv6.rs
View File

@ -42,6 +42,7 @@ where
T: From<Vec<u8>>,
{
/// Construct a new ICMPv6 packet from raw bytes
#[profiling::function]
pub fn new_from_bytes(
bytes: &[u8],
source_address: Ipv6Addr,
@ -64,6 +65,7 @@ where
impl Icmpv6Packet<RawBytes> {
/// Construct a new ICMPv6 packet with a raw payload from raw bytes
#[profiling::function]
pub fn new_from_bytes_raw_payload(
bytes: &[u8],
source_address: Ipv6Addr,
@ -88,6 +90,7 @@ impl<T> Into<Vec<u8>> for Icmpv6Packet<T>
where
T: Into<Vec<u8>>,
{
#[profiling::function]
fn into(self) -> Vec<u8> {
// Convert the payload into raw bytes
let payload: Vec<u8> = self.payload.into();

3
src/packet/protocols/ipv4.rs
View File

@ -53,6 +53,7 @@ impl<T> Ipv4Packet<T> {
}
}
#[profiling::function]
fn options_length_words(&self) -> u8 {
self.options
.iter()
@ -68,6 +69,7 @@ where
{
type Error = PacketError;
#[profiling::function]
fn try_from(bytes: Vec<u8>) -> Result<Self, Self::Error> {
// Parse the packet
let packet =
@ -94,6 +96,7 @@ impl<T> Into<Vec<u8>> for Ipv4Packet<T>
where
T: Into<Vec<u8>> + Clone,
{
#[profiling::function]
fn into(self) -> Vec<u8> {
// Convert the payload into raw bytes
let payload: Vec<u8> = self.payload.clone().into();

2
src/packet/protocols/ipv6.rs
View File

@ -44,6 +44,7 @@ where
{
type Error = PacketError;
#[profiling::function]
fn try_from(bytes: Vec<u8>) -> Result<Self, Self::Error> {
// Parse the packet
let packet =
@ -66,6 +67,7 @@ impl<T> Into<Vec<u8>> for Ipv6Packet<T>
where
T: Into<Vec<u8>>,
{
#[profiling::function]
fn into(self) -> Vec<u8> {
// Convert the payload into raw bytes
let payload: Vec<u8> = self.payload.into();

3
src/packet/protocols/tcp.rs
View File

@ -113,6 +113,7 @@ where
T: From<Vec<u8>>,
{
/// Construct a new TCP packet from bytes
#[profiling::function]
pub fn new_from_bytes(
bytes: &[u8],
source_address: IpAddr,
@ -147,6 +148,7 @@ where
impl TcpPacket<RawBytes> {
/// Construct a new TCP packet with a raw payload from bytes
#[profiling::function]
pub fn new_from_bytes_raw_payload(
bytes: &[u8],
source_address: IpAddr,
@ -183,6 +185,7 @@ impl<T> Into<Vec<u8>> for TcpPacket<T>
where
T: Into<Vec<u8>>,
{
#[profiling::function]
fn into(self) -> Vec<u8> {
// Get the options length in words
let options_length = self.options_length();

3
src/packet/protocols/udp.rs
View File

@ -84,6 +84,7 @@ where
T: From<Vec<u8>>,
{
/// Construct a new UDP packet from bytes
#[profiling::function]
pub fn new_from_bytes(
bytes: &[u8],
source_address: IpAddr,
@ -112,6 +113,7 @@ where
impl UdpPacket<RawBytes> {
/// Construct a new UDP packet with a raw payload from bytes
#[profiling::function]
pub fn new_from_bytes_raw_payload(
bytes: &[u8],
source_address: IpAddr,
@ -142,6 +144,7 @@ impl<T> Into<Vec<u8>> for UdpPacket<T>
where
T: Into<Vec<u8>>,
{
#[profiling::function]
fn into(self) -> Vec<u8> {
// Convert the payload into raw bytes
let payload: Vec<u8> = self.payload.into();

4
src/packet/xlat/icmp/mod.rs
View File

@ -12,6 +12,7 @@ use super::ip::{translate_ipv4_to_ipv6, translate_ipv6_to_ipv4};
mod type_code;
/// Translates an ICMP packet to an ICMPv6 packet
#[profiling::function]
pub fn translate_icmp_to_icmpv6(
input: IcmpPacket<RawBytes>,
new_source: Ipv6Addr,
@ -55,6 +56,7 @@ pub fn translate_icmp_to_icmpv6(
}
/// Translates an ICMPv6 packet to an ICMP packet
#[profiling::function]
pub fn translate_icmpv6_to_icmp(
input: Icmpv6Packet<RawBytes>,
new_source: Ipv4Addr,
@ -83,7 +85,7 @@ pub fn translate_icmpv6_to_icmp(
buffer.extend_from_slice(&inner_payload);
buffer
})
},
}
_ => input.payload,
};

2
src/packet/xlat/icmp/type_code.rs
View File

@ -8,6 +8,7 @@ use pnet_packet::{
use crate::packet::error::PacketError;
/// Best effort translation from an ICMP type and code to an ICMPv6 type and code
#[profiling::function]
pub fn translate_type_and_code_4_to_6(
icmp_type: IcmpType,
icmp_code: IcmpCode,
@ -55,6 +56,7 @@ pub fn translate_type_and_code_4_to_6(
}
/// Best effort translation from an ICMPv6 type and code to an ICMP type and code
#[profiling::function]
pub fn translate_type_and_code_6_to_4(
icmp_type: Icmpv6Type,
icmp_code: Icmpv6Code,

2
src/packet/xlat/ip.rs
View File

@ -17,6 +17,7 @@ use super::{
};
/// Translates an IPv4 packet to an IPv6 packet
#[profiling::function]
pub fn translate_ipv4_to_ipv6(
input: Ipv4Packet<Vec<u8>>,
new_source: Ipv6Addr,
@ -69,6 +70,7 @@ pub fn translate_ipv4_to_ipv6(
}
/// Translates an IPv6 packet to an IPv4 packet
#[profiling::function]
pub fn translate_ipv6_to_ipv4(
input: Ipv6Packet<Vec<u8>>,
new_source: Ipv4Addr,

2
src/packet/xlat/tcp.rs
View File

@ -6,6 +6,7 @@ use crate::packet::{
};
/// Translates an IPv4 TCP packet to an IPv6 TCP packet
#[profiling::function]
pub fn translate_tcp4_to_tcp6(
input: TcpPacket<RawBytes>,
new_source_addr: Ipv6Addr,
@ -26,6 +27,7 @@ pub fn translate_tcp4_to_tcp6(
}
/// Translates an IPv6 TCP packet to an IPv4 TCP packet
#[profiling::function]
pub fn translate_tcp6_to_tcp4(
input: TcpPacket<RawBytes>,
new_source_addr: Ipv4Addr,

2
src/packet/xlat/udp.rs
View File

@ -6,6 +6,7 @@ use crate::packet::{
};
/// Translates an IPv4 UDP packet to an IPv6 UDP packet
#[profiling::function]
pub fn translate_udp4_to_udp6(
input: UdpPacket<RawBytes>,
new_source_addr: Ipv6Addr,
@ -20,6 +21,7 @@ pub fn translate_udp4_to_udp6(
}
/// Translates an IPv6 UDP packet to an IPv4 UDP packet
#[profiling::function]
pub fn translate_udp6_to_udp4(
input: UdpPacket<RawBytes>,
new_source_addr: Ipv4Addr,