evan/protomask
1
Fork 0
Code

Implement optional profiling

Browse Source
This commit is contained in:
Evan Pratten 2023-08-04 22:29:42 -04:00
parent 6c0e6b2341
commit eabfbfc8b6
13 changed files with 317 additions and 119 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

41
Cargo.toml
View File

@ -13,33 +13,6 @@ keywords = []
categories = []
exclude = ["/.github/", "/.vscode/"]
# # 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",
# # "process",
# "sync"
# ] }
# clap = { version = "4.3.11", features = ["derive"] }
# serde = { version = "1.0.171", features = ["derive"] }
# ipnet = { version = "2.8.0", features = ["serde"] }
# hyper = { version = "0.14.27", features = ["server", "http1", "tcp"] }
# owo-colors = { version = "3.5.0", features = ["supports-colors"] }
# toml = "0.7.6"
# log = "0.4.19"
# fern = "0.6.2"
# serde_path_to_error = "0.1.13"
# thiserror = "1.0.43"
# tun-tap = "0.1.3"
# bimap = "0.6.3"
# pnet_packet = "0.34.0"
# rtnetlink = "0.13.0"
# futures = "0.3.28"
# prometheus = "0.13.3"
# lazy_static = "1.4.0"
[workspace]
members = [
"libs/easy-tun",
@ -50,6 +23,16 @@ members = [
"libs/protomask-metrics",
]
[features]
default = []
profiler = [
"puffin",
"puffin_http",
"easy-tun/profile-puffin",
"fast-nat/profile-puffin",
"interproto/profile-puffin",
]
[[bin]]
name = "protomask"
path = "src/protomask.rs"
@ -76,12 +59,16 @@ tokio = { version = "1.29.1", features = ["macros", "rt-multi-thread"] }
owo-colors = { version = "3.5.0", features = ["supports-colors"] }
clap = { version = "4.3.11", features = ["derive"] }
ipnet = { version = "2.8.0", features = ["serde"] }
puffin_http = { version = "0.13.0", optional = true }
puffin = { version = "0.16.0", optional = true }
serde = { version = "^1.0", features = ["derive"] }
serde_json = "^1.0"
log = "0.4.19"
fern = "0.6.2"
nix = "0.26.2"
thiserror = "1.0.44"
cfg-if = "1.0.0"
profiling = "1.0.9"
[package.metadata.deb]
section = "network"

4
libs/easy-tun/Cargo.toml
View File

@ -12,6 +12,10 @@ license = "GPL-3.0"
keywords = []
categories = []
[features]
default = []
profile-puffin = ["profiling/profile-with-puffin"]
[dependencies]
log = "^0.4"
libc = "^0.2"

6
libs/easy-tun/src/tun.rs
View File

@ -96,6 +96,12 @@ impl Tun {
pub fn name(&self) -> &str {
&self.name
}
/// Get the underlying file descriptor
#[must_use]
pub fn fd(&self) -> &File {
&self.fd
}
}
impl AsRawFd for Tun {

4
libs/fast-nat/Cargo.toml
View File

@ -12,6 +12,10 @@ license = "GPL-3.0"
keywords = []
categories = []
[features]
default = []
profile-puffin = ["profiling/profile-with-puffin"]
[dependencies]
log = "^0.4"
rustc-hash = "1.1.0"

1
libs/interproto/Cargo.toml
View File

@ -15,6 +15,7 @@ categories = []
[features]
default = []
metrics = ["protomask-metrics"]
profile-puffin = ["profiling/profile-with-puffin"]
[dependencies]
protomask-metrics = { path = "../protomask-metrics", optional = true }

18
src/args/mod.rs
View File

@ -1,4 +1,22 @@
//! This module contains the definitions for each binary's CLI arguments and config file structure for the sake of readability.
use cfg_if::cfg_if;
pub mod protomask;
pub mod protomask_clat;
// Used to trick the build process into including a CLI argument based on a feature flag
cfg_if! {
if #[cfg(feature = "profiler")] {
#[derive(Debug, clap::Args)]
pub struct ProfilerArgs {
/// Expose the puffin HTTP server on this endpoint
#[clap(long)]
pub puffin_endpoint: Option<std::net::SocketAddr>,
}
} else {
#[derive(Debug, clap::Args)]
pub struct ProfilerArgs;
}
}

6
src/args/protomask.rs
View File

@ -7,6 +7,8 @@ use ipnet::{Ipv4Net, Ipv6Net};
use crate::common::rfc6052::parse_network_specific_prefix;
use super::ProfilerArgs;
#[derive(clap::Parser)]
#[clap(author, version, about="Fast and simple NAT64", long_about = None)]
pub struct Args {
@ -21,6 +23,9 @@ pub struct Args {
#[clap(short, long, default_value_t = ("nat%d").to_string())]
pub interface: String,
#[command(flatten)]
pub profiler_args: ProfilerArgs,
/// Enable verbose logging
#[clap(short, long)]
pub verbose: bool,
@ -102,3 +107,4 @@ impl From<StaticMap> for (Ipv4Addr, Ipv6Addr) {
(val.ipv4, val.ipv6)
}
}

4
src/args/protomask_clat.rs
View File

@ -1,5 +1,6 @@
//! Commandline arguments and config file definitions for `protomask-clat`
use super::ProfilerArgs;
use crate::common::rfc6052::parse_network_specific_prefix;
use ipnet::{Ipv4Net, Ipv6Net};
use std::{net::SocketAddr, path::PathBuf};
@ -18,6 +19,9 @@ pub struct Args {
#[clap(short, long, default_value_t = ("clat%d").to_string())]
pub interface: String,
#[command(flatten)]
pub profiler_args: ProfilerArgs,
/// Enable verbose logging
#[clap(short, long)]
pub verbose: bool,

1
src/common/mod.rs
View File

@ -4,3 +4,4 @@ pub mod logging;
pub mod packet_handler;
pub mod permissions;
pub mod rfc6052;
pub mod profiler;

156
src/common/packet_handler.rs
View File

@ -8,16 +8,9 @@ pub enum PacketHandlingError {
FastNatError(#[from] fast_nat::error::Error),
}
/// Handles checking the version number of an IP packet and calling the correct handler with needed data
pub fn handle_packet<Ipv4Handler, Ipv6Handler>(
packet: &[u8],
mut ipv4_handler: Ipv4Handler,
mut ipv6_handler: Ipv6Handler,
) -> Option<Vec<u8>>
where
Ipv4Handler: FnMut(&[u8], &Ipv4Addr, &Ipv4Addr) -> Result<Option<Vec<u8>>, PacketHandlingError>,
Ipv6Handler: FnMut(&[u8], &Ipv6Addr, &Ipv6Addr) -> Result<Option<Vec<u8>>, PacketHandlingError>,
{
/// Get the layer 3 protocol of a packet
pub fn get_layer_3_proto(packet: &[u8]) -> Option<u8> {
// If the packet is empty, return nothing
if packet.is_empty() {
return None;
@ -26,40 +19,31 @@ where
// Switch on the layer 3 protocol number to call the correct handler
let layer_3_proto = packet[0] >> 4;
log::trace!("New packet with layer 3 protocol: {}", layer_3_proto);
let handler_response = match layer_3_proto {
// IPv4
4 => {
// Extract the source and destination addresses
let source_addr =
Ipv4Addr::from(u32::from_be_bytes(packet[12..16].try_into().unwrap()));
let destination_addr =
Ipv4Addr::from(u32::from_be_bytes(packet[16..20].try_into().unwrap()));
Some(layer_3_proto)
}
// Call the handler
ipv4_handler(packet, &source_addr, &destination_addr)
}
/// Get the source and destination addresses of an IPv4 packet
pub fn get_ipv4_src_dst(packet: &[u8]) -> (Ipv4Addr, Ipv4Addr) {
let source_addr = Ipv4Addr::from(u32::from_be_bytes(packet[12..16].try_into().unwrap()));
let destination_addr = Ipv4Addr::from(u32::from_be_bytes(packet[16..20].try_into().unwrap()));
// IPv6
6 => {
// Extract the source and destination addresses
let source_addr =
Ipv6Addr::from(u128::from_be_bytes(packet[8..24].try_into().unwrap()));
let destination_addr =
Ipv6Addr::from(u128::from_be_bytes(packet[24..40].try_into().unwrap()));
(source_addr, destination_addr)
}
// Call the handler
ipv6_handler(packet, &source_addr, &destination_addr)
}
/// Get the source and destination addresses of an IPv6 packet
pub fn get_ipv6_src_dst(packet: &[u8]) -> (Ipv6Addr, Ipv6Addr) {
let source_addr = Ipv6Addr::from(u128::from_be_bytes(packet[8..24].try_into().unwrap()));
let destination_addr = Ipv6Addr::from(u128::from_be_bytes(packet[24..40].try_into().unwrap()));
// Unknown protocol numbers can't be handled
proto => {
log::warn!("Unknown Layer 3 protocol: {}", proto);
return None;
}
};
(source_addr, destination_addr)
}
// The response from the handler may or may not be a warn-able error
match handler_response {
/// Appropriately handle a translation error
pub fn handle_translation_error(
result: Result<Option<Vec<u8>>, PacketHandlingError>,
) -> Option<Vec<u8>> {
// We may or may not have a warn-able error
match result {
// If we get data, return it
Ok(data) => data,
// If we get an error, handle it and return None
@ -101,3 +85,97 @@ where
},
}
}
// /// Handles checking the version number of an IP packet and calling the correct handler with needed data
// pub fn handle_packet<Ipv4Handler, Ipv6Handler>(
// packet: &[u8],
// mut ipv4_handler: Ipv4Handler,
// mut ipv6_handler: Ipv6Handler,
// ) -> Option<Vec<u8>>
// where
// Ipv4Handler: FnMut(&[u8], &Ipv4Addr, &Ipv4Addr) -> Result<Option<Vec<u8>>, PacketHandlingError>,
// Ipv6Handler: FnMut(&[u8], &Ipv6Addr, &Ipv6Addr) -> Result<Option<Vec<u8>>, PacketHandlingError>,
// {
// // If the packet is empty, return nothing
// if packet.is_empty() {
// return None;
// }
// // Switch on the layer 3 protocol number to call the correct handler
// let layer_3_proto = packet[0] >> 4;
// log::trace!("New packet with layer 3 protocol: {}", layer_3_proto);
// let handler_response = match layer_3_proto {
// // IPv4
// 4 => {
// // Extract the source and destination addresses
// let source_addr =
// Ipv4Addr::from(u32::from_be_bytes(packet[12..16].try_into().unwrap()));
// let destination_addr =
// Ipv4Addr::from(u32::from_be_bytes(packet[16..20].try_into().unwrap()));
// // Call the handler
// ipv4_handler(packet, &source_addr, &destination_addr)
// }
// // IPv6
// 6 => {
// // Extract the source and destination addresses
// let source_addr =
// Ipv6Addr::from(u128::from_be_bytes(packet[8..24].try_into().unwrap()));
// let destination_addr =
// Ipv6Addr::from(u128::from_be_bytes(packet[24..40].try_into().unwrap()));
// // Call the handler
// ipv6_handler(packet, &source_addr, &destination_addr)
// }
// // Unknown protocol numbers can't be handled
// proto => {
// log::warn!("Unknown Layer 3 protocol: {}", proto);
// return None;
// }
// };
// // The response from the handler may or may not be a warn-able error
// match handler_response {
// // If we get data, return it
// Ok(data) => data,
// // If we get an error, handle it and return None
// Err(error) => match error {
// PacketHandlingError::InterprotoError(interproto::error::Error::PacketTooShort {
// expected,
// actual,
// }) => {
// log::warn!(
// "Got packet with length {} when expecting at least {} bytes",
// actual,
// expected
// );
// None
// }
// PacketHandlingError::InterprotoError(
// interproto::error::Error::UnsupportedIcmpType(icmp_type),
// ) => {
// log::warn!("Got a packet with an unsupported ICMP type: {}", icmp_type);
// None
// }
// PacketHandlingError::InterprotoError(
// interproto::error::Error::UnsupportedIcmpv6Type(icmpv6_type),
// ) => {
// log::warn!(
// "Got a packet with an unsupported ICMPv6 type: {}",
// icmpv6_type
// );
// None
// }
// PacketHandlingError::FastNatError(fast_nat::error::Error::Ipv4PoolExhausted) => {
// log::warn!("IPv4 pool exhausted. Dropping packet.");
// None
// }
// PacketHandlingError::FastNatError(fast_nat::error::Error::InvalidIpv4Address(addr)) => {
// log::warn!("Invalid IPv4 address: {}", addr);
// None
// }
// },
// }
// }

20
src/common/profiler.rs Normal file
View File

@ -0,0 +1,20 @@
use cfg_if::cfg_if;
use crate::args::ProfilerArgs;
cfg_if! {
if #[cfg(feature = "profiler")] {
pub fn start_puffin_server(args: &ProfilerArgs) -> Option<puffin_http::Server> {
if let Some(endpoint) = args.puffin_endpoint {
log::info!("Starting puffin server on {}", endpoint);
puffin::set_scopes_on(true);
Some(puffin_http::Server::new(&endpoint.to_string()).unwrap())
} else {
None
}
}
} else {
#[allow(dead_code)]
pub fn start_puffin_server(_args: &ProfilerArgs){}
}
}

75
src/protomask-clat.rs
View File

@ -3,7 +3,11 @@
//! This binary is a Customer-side transLATor (CLAT) that translates all native
//! IPv4 traffic to IPv6 traffic for transmission over an IPv6-only ISP network.
use crate::common::packet_handler::handle_packet;
use crate::common::packet_handler::{
get_ipv4_src_dst, get_ipv6_src_dst, get_layer_3_proto, handle_translation_error,
PacketHandlingError,
};
use crate::common::profiler::start_puffin_server;
use crate::{args::protomask_clat::Args, common::permissions::ensure_root};
use clap::Parser;
use common::logging::enable_logger;
@ -30,6 +34,10 @@ pub async fn main() {
// We must be root to continue program execution
ensure_root();
// Start profiling
#[allow(clippy::let_unit_value)]
let _server = start_puffin_server(&args.profiler_args);
// Bring up a TUN interface
let mut tun = Tun::new(&args.interface).unwrap();
@ -81,34 +89,51 @@ pub async fn main() {
log::info!("Translating packets on {}", tun.name());
let mut buffer = vec![0u8; 1500];
loop {
// Indicate to the profiler that we are starting a new packet
profiling::finish_frame!();
profiling::scope!("packet");
// Read a packet
let len = tun.read(&mut buffer).unwrap();
// Translate it based on the Layer 3 protocol number
if let Some(output) = handle_packet(
&buffer[..len],
// IPv4 -> IPv6
|packet, source, dest| {
Ok(translate_ipv4_to_ipv6(
packet,
unsafe { embed_ipv4_addr_unchecked(*source, config.embed_prefix) },
unsafe { embed_ipv4_addr_unchecked(*dest, config.embed_prefix) },
)
.map(Some)?)
},
// IPv6 -> IPv4
|packet, source, dest| {
Ok(translate_ipv6_to_ipv4(
packet,
unsafe {
extract_ipv4_addr_unchecked(*source, config.embed_prefix.prefix_len())
},
unsafe { extract_ipv4_addr_unchecked(*dest, config.embed_prefix.prefix_len()) },
)
.map(Some)?)
},
) {
// Write the packet if we get one back from the handler functions
let translation_result: Result<Option<Vec<u8>>, PacketHandlingError> =
match get_layer_3_proto(&buffer[..len]) {
Some(4) => {
let (source, dest) = get_ipv4_src_dst(&buffer[..len]);
translate_ipv4_to_ipv6(
&buffer[..len],
unsafe { embed_ipv4_addr_unchecked(source, config.embed_prefix) },
unsafe { embed_ipv4_addr_unchecked(dest, config.embed_prefix) },
)
.map(Some)
.map_err(PacketHandlingError::from)
}
Some(6) => {
let (source, dest) = get_ipv6_src_dst(&buffer[..len]);
translate_ipv6_to_ipv4(
&buffer[..len],
unsafe {
extract_ipv4_addr_unchecked(source, config.embed_prefix.prefix_len())
},
unsafe {
extract_ipv4_addr_unchecked(dest, config.embed_prefix.prefix_len())
},
)
.map(Some)
.map_err(PacketHandlingError::from)
}
Some(proto) => {
log::warn!("Unknown Layer 3 protocol: {}", proto);
continue;
}
None => {
continue;
}
};
// Handle any errors and write
if let Some(output) = handle_translation_error(translation_result) {
tun.write_all(&output).unwrap();
}
}

100
src/protomask.rs
View File

@ -1,4 +1,11 @@
use crate::common::{packet_handler::handle_packet, permissions::ensure_root};
use crate::common::{
packet_handler::{
get_ipv4_src_dst, get_ipv6_src_dst, get_layer_3_proto, handle_translation_error,
PacketHandlingError,
},
permissions::ensure_root,
profiler::start_puffin_server,
};
use clap::Parser;
use common::logging::enable_logger;
use easy_tun::Tun;
@ -29,6 +36,10 @@ pub async fn main() {
// We must be root to continue program execution
ensure_root();
// Start profiling
#[allow(clippy::let_unit_value)]
let _server = start_puffin_server(&args.profiler_args);
// Bring up a TUN interface
log::debug!("Creating new TUN interface");
let mut tun = Tun::new(&args.interface).unwrap();
@ -88,38 +99,71 @@ pub async fn main() {
log::info!("Translating packets on {}", tun.name());
let mut buffer = vec![0u8; 1500];
loop {
// Indicate to the profiler that we are starting a new packet
profiling::finish_frame!();
profiling::scope!("packet");
// Read a packet
let len = tun.read(&mut buffer).unwrap();
// Translate it based on the Layer 3 protocol number
if let Some(output) = handle_packet(
&buffer[..len],
// IPv4 -> IPv6
|packet, source, dest| match addr_table.borrow().get_ipv6(dest) {
Some(new_destination) => Ok(translate_ipv4_to_ipv6(
packet,
unsafe { embed_ipv4_addr_unchecked(*source, config.translation_prefix) },
new_destination,
)
.map(Some)?),
None => {
protomask_metrics::metric!(PACKET_COUNTER, PROTOCOL_IPV4, STATUS_DROPPED);
Ok(None)
let translation_result: Result<Option<Vec<u8>>, PacketHandlingError> =
match get_layer_3_proto(&buffer[..len]) {
Some(4) => {
let (source, dest) = get_ipv4_src_dst(&buffer[..len]);
match addr_table.borrow().get_ipv6(&dest) {
Some(new_destination) => translate_ipv4_to_ipv6(
&buffer[..len],
unsafe { embed_ipv4_addr_unchecked(source, config.translation_prefix) },
new_destination,
)
.map(Some)
.map_err(PacketHandlingError::from),
None => {
protomask_metrics::metric!(
PACKET_COUNTER,
PROTOCOL_IPV4,
STATUS_DROPPED
);
Ok(None)
}
}
}
},
// IPv6 -> IPv4
|packet, source, dest| {
Ok(translate_ipv6_to_ipv4(
packet,
addr_table.borrow_mut().get_or_create_ipv4(source)?,
unsafe {
extract_ipv4_addr_unchecked(*dest, config.translation_prefix.prefix_len())
},
)
.map(Some)?)
},
) {
// Write the packet if we get one back from the handler functions
Some(6) => {
let (source, dest) = get_ipv6_src_dst(&buffer[..len]);
match addr_table.borrow_mut().get_or_create_ipv4(&source) {
Ok(new_source) => {
translate_ipv6_to_ipv4(&buffer[..len], new_source, unsafe {
extract_ipv4_addr_unchecked(
dest,
config.translation_prefix.prefix_len(),
)
})
.map(Some)
.map_err(PacketHandlingError::from)
}
Err(error) => {
log::error!("Error getting IPv4 address: {}", error);
protomask_metrics::metric!(
PACKET_COUNTER,
PROTOCOL_IPV6,
STATUS_DROPPED
);
Ok(None)
}
}
}
Some(proto) => {
log::warn!("Unknown Layer 3 protocol: {}", proto);
continue;
}
None => {
continue;
}
};
// Handle any errors and write
if let Some(output) = handle_translation_error(translation_result) {
tun.write_all(&output).unwrap();
}
}