Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

pcap-rewrite: improve filtering key types #35

Merged
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
35 changes: 24 additions & 11 deletions pcap-rewrite/src/container/ipaddr_proto_port_container.rs
Original file line number Diff line number Diff line change
@@ -1,28 +1,44 @@
use std::collections::HashSet;
use std::error::Error;
use std::fs::File;
use std::iter::FromIterator;
use std::net::IpAddr;
use std::path::Path;
use std::str::FromStr;

use csv::ReaderBuilder;
use pnet_packet::ip::IpNextHeaderProtocol;

#[derive(Debug, PartialEq, Eq, Hash)]
pub struct IpAddrProtoPort {
ipaddr: IpAddr,
proto: IpNextHeaderProtocol,
port: u16,
}

impl IpAddrProtoPort {
pub fn new(ipaddr: IpAddr, proto: IpNextHeaderProtocol, port: u16) -> IpAddrProtoPort {
IpAddrProtoPort {
ipaddr,
proto,
port,
}
}
}

pub struct IpAddrProtoPortC {
s: HashSet<(IpAddr, IpNextHeaderProtocol, u16)>,
s: HashSet<IpAddrProtoPort>,
}

impl IpAddrProtoPortC {
pub fn new(s: HashSet<(IpAddr, IpNextHeaderProtocol, u16)>) -> IpAddrProtoPortC {
pub fn new(s: HashSet<IpAddrProtoPort>) -> IpAddrProtoPortC {
IpAddrProtoPortC { s }
}

pub fn of_file_path(ip_file_path: &Path) -> Result<IpAddrProtoPortC, Box<dyn Error>> {
let file = File::open(ip_file_path)?;

let mut rdr = ReaderBuilder::new().has_headers(false).from_reader(file);
let ipaddr_proto_port_v = rdr
let hs = rdr
.records()
.map(|l| {
let record = l?;
Expand All @@ -45,11 +61,9 @@ impl IpAddrProtoPortC {
.ok_or_else(|| "Missing port value in dispatch filter key file".to_string())?;
let port = port_s.parse()?;

Ok((ipaddr, protocol, port))
Ok(IpAddrProtoPort::new(ipaddr, protocol, port))
})
.collect::<Result<Vec<(IpAddr, IpNextHeaderProtocol, u16)>, Box<dyn Error>>>()?;

let hs = HashSet::from_iter(ipaddr_proto_port_v.iter().cloned());
.collect::<Result<HashSet<IpAddrProtoPort>, Box<dyn Error>>>()?;

Ok(IpAddrProtoPortC::new(hs))
}
Expand All @@ -62,8 +76,7 @@ impl IpAddrProtoPortC {
// self.s.len()
// }

pub fn contains(&self, ipaddr: &IpAddr, proto: &IpNextHeaderProtocol, port: u16) -> bool {
let t = (*ipaddr, *proto, port);
self.s.contains(&t)
pub fn contains(&self, ipaddr_proto_port: &IpAddrProtoPort) -> bool {
self.s.contains(ipaddr_proto_port)
}
}
31 changes: 15 additions & 16 deletions pcap-rewrite/src/filters/dispatch_filter.rs
Original file line number Diff line number Diff line change
Expand Up @@ -6,16 +6,16 @@ use std::path::Path;
use libpcap_tools::{Error, FiveTuple, ParseContext};
use pcap_parser::data::PacketData;
use pnet_packet::ethernet::{EtherType, EtherTypes};
use pnet_packet::ip::IpNextHeaderProtocol;
use pnet_packet::PrimitiveValues;

use crate::container::five_tuple_container::FiveTupleC;
use crate::container::ipaddr_container::IpAddrC;
use crate::container::ipaddr_proto_port_container::IpAddrProtoPortC;
use crate::container::ipaddr_proto_port_container::{IpAddrProtoPort, IpAddrProtoPortC};
use crate::filters::filter::{FResult, Filter, Verdict};
use crate::filters::filter_utils;
use crate::filters::filtering_action::FilteringAction;
use crate::filters::filtering_key::FilteringKey;
use crate::filters::ipaddr_pair::IpAddrPair;
use crate::filters::key_parser_ipv4;
use crate::filters::key_parser_ipv6;

Expand Down Expand Up @@ -159,12 +159,12 @@ impl DispatchFilterBuilder {
let ipaddr_container = IpAddrC::of_file_path(Path::new(key_file_path))
.map_err(|e| io::Error::new(io::ErrorKind::Other, e.to_string()))?;

let keep: KeepFn<IpAddrC, (IpAddr, IpAddr)> = match filtering_action {
FilteringAction::Keep => Box::new(|c, ipaddr_tuple| {
Ok(c.contains(&ipaddr_tuple.0) || c.contains(&ipaddr_tuple.1))
let keep: KeepFn<IpAddrC, IpAddrPair> = match filtering_action {
FilteringAction::Keep => Box::new(|c, ipaddr_pair| {
Ok(c.contains(&ipaddr_pair.0) || c.contains(&ipaddr_pair.1))
}),
FilteringAction::Drop => Box::new(|c, ipaddr_tuple| {
Ok(!c.contains(&ipaddr_tuple.0) && !c.contains(&ipaddr_tuple.1))
FilteringAction::Drop => Box::new(|c, ipaddr_pair| {
Ok(!c.contains(&ipaddr_pair.0) && !c.contains(&ipaddr_pair.1))
}),
};

Expand All @@ -180,15 +180,14 @@ impl DispatchFilterBuilder {
IpAddrProtoPortC::of_file_path(Path::new(key_file_path))
.map_err(|e| io::Error::new(io::ErrorKind::Other, e.to_string()))?;

let keep: KeepFn<IpAddrProtoPortC, (IpAddr, IpNextHeaderProtocol, u16)> =
match filtering_action {
FilteringAction::Keep => {
Box::new(|c, tuple| Ok(c.contains(&tuple.0, &tuple.1, tuple.2)))
}
FilteringAction::Drop => {
Box::new(|c, tuple| Ok(!c.contains(&tuple.0, &tuple.1, tuple.2)))
}
};
let keep: KeepFn<IpAddrProtoPortC, IpAddrProtoPort> = match filtering_action {
FilteringAction::Keep => {
Box::new(|c, ipaddr_proto_port| Ok(c.contains(ipaddr_proto_port)))
}
FilteringAction::Drop => {
Box::new(|c, ipaddr_proto_port| Ok(!c.contains(ipaddr_proto_port)))
}
};

Ok(Box::new(DispatchFilter::new(
ipaddr_proto_port_container,
Expand Down
6 changes: 3 additions & 3 deletions pcap-rewrite/src/filters/fragmentation/convert_fn.rs
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@ use libpcap_tools::FiveTuple;

use crate::container::five_tuple_container::FiveTupleC;
use crate::container::ipaddr_container::IpAddrC;
use crate::container::ipaddr_proto_port_container::IpAddrProtoPortC;
use crate::container::ipaddr_proto_port_container::{IpAddrProtoPort, IpAddrProtoPortC};
use crate::container::two_tuple_proto_ipid_container::TwoTupleProtoIpidC;
use crate::filters::fragmentation::two_tuple_proto_ipid_five_tuple::TwoTupleProtoIpidFiveTuple;

Expand Down Expand Up @@ -51,7 +51,7 @@ pub fn convert_data_hs_to_src_ipaddr_proto_dst_port_container(
.iter()
.filter_map(|t| t.get_five_tuple_option())
.map(|five_tuple| {
(
IpAddrProtoPort::new(
five_tuple.src,
IpNextHeaderProtocol::new(five_tuple.proto),
five_tuple.dst_port,
Expand All @@ -61,7 +61,7 @@ pub fn convert_data_hs_to_src_ipaddr_proto_dst_port_container(
.iter()
.filter_map(|t| t.get_five_tuple_option())
.map(|five_tuple| {
(
IpAddrProtoPort::new(
five_tuple.dst,
IpNextHeaderProtocol::new(five_tuple.proto),
five_tuple.src_port,
Expand Down
38 changes: 18 additions & 20 deletions pcap-rewrite/src/filters/fragmentation/fragmentation_filter.rs
Original file line number Diff line number Diff line change
Expand Up @@ -5,24 +5,23 @@ use std::net::IpAddr;
use log::warn;
use pcap_parser::data::PacketData;
use pnet_packet::ethernet::{EtherType, EtherTypes};
use pnet_packet::ip::IpNextHeaderProtocol;

use libpcap_tools::{Error, Packet, ParseContext};

use crate::container::five_tuple_container::FiveTupleC;
use crate::container::ipaddr_container::IpAddrC;
use crate::container::ipaddr_proto_port_container::IpAddrProtoPortC;
use crate::container::ipaddr_proto_port_container::{IpAddrProtoPort, IpAddrProtoPortC};
use crate::container::two_tuple_proto_ipid_container::TwoTupleProtoIpidC;
use crate::filters::filter::Filter;
use crate::filters::filter::{FResult, Verdict};
use crate::filters::filter_utils;
use crate::filters::filtering_action::FilteringAction;
use crate::filters::filtering_key::FilteringKey;
use crate::filters::key_parser_ipv4;
use crate::filters::key_parser_ipv6;

use crate::filters::fragmentation::fragmentation_test;
use crate::filters::fragmentation::two_tuple_proto_ipid_five_tuple::TwoTupleProtoIpidFiveTuple;
use crate::filters::ipaddr_pair::IpAddrPair;
use crate::filters::key_parser_ipv4;
use crate::filters::key_parser_ipv6;

use super::convert_fn;

Expand Down Expand Up @@ -298,12 +297,12 @@ impl FragmentationFilterBuilder {
FilteringKey::SrcDstIpaddr => {
let ipaddr_container = IpAddrC::new(HashSet::new());

let keep: KeepFn<IpAddrC, (IpAddr, IpAddr)> = match filtering_action {
FilteringAction::Keep => |c, ipaddr_tuple| {
Ok(c.contains(&ipaddr_tuple.0) || c.contains(&ipaddr_tuple.1))
},
FilteringAction::Drop => |c, ipaddr_tuple| {
Ok(!c.contains(&ipaddr_tuple.0) && !c.contains(&ipaddr_tuple.1))
let keep: KeepFn<IpAddrC, IpAddrPair> = match filtering_action {
FilteringAction::Keep => {
|c, ipaddr_pair| Ok(c.contains(&ipaddr_pair.0) || c.contains(&ipaddr_pair.1))
}
FilteringAction::Drop => |c, ipaddr_pair| {
Ok(!c.contains(&ipaddr_pair.0) && !c.contains(&ipaddr_pair.1))
},
};

Expand All @@ -319,15 +318,14 @@ impl FragmentationFilterBuilder {
FilteringKey::SrcIpaddrProtoDstPort => {
let ipaddr_proto_port_container = IpAddrProtoPortC::new(HashSet::new());

let keep: KeepFn<IpAddrProtoPortC, (IpAddr, IpNextHeaderProtocol, u16)> =
match filtering_action {
FilteringAction::Keep => {
|c, tuple| Ok(c.contains(&tuple.0, &tuple.1, tuple.2))
}
FilteringAction::Drop => {
|c, tuple| Ok(!c.contains(&tuple.0, &tuple.1, tuple.2))
}
};
let keep: KeepFn<IpAddrProtoPortC, IpAddrProtoPort> = match filtering_action {
FilteringAction::Keep => {
|c: &IpAddrProtoPortC, ipaddr_proto_port| Ok(c.contains(ipaddr_proto_port))
}
FilteringAction::Drop => {
|c, ipaddr_proto_port| Ok(!c.contains(ipaddr_proto_port))
}
};

Ok(Box::new(FragmentationFilter::new(
HashSet::new(),
Expand Down
9 changes: 9 additions & 0 deletions pcap-rewrite/src/filters/ipaddr_pair.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,9 @@
use std::net::IpAddr;

pub struct IpAddrPair (pub IpAddr,pub IpAddr);

impl IpAddrPair {
pub fn new(ipaddr_0: IpAddr, ipaddr_1: IpAddr) -> IpAddrPair {
IpAddrPair (ipaddr_0, ipaddr_1 )
}
}
30 changes: 20 additions & 10 deletions pcap-rewrite/src/filters/key_parser_ipv4.rs
Original file line number Diff line number Diff line change
@@ -1,7 +1,6 @@
use std::net::IpAddr;

use log::warn;
use pnet_packet::ip::IpNextHeaderProtocol;
use pnet_packet::ip::IpNextHeaderProtocols;
use pnet_packet::ipv4::Ipv4Packet;
use pnet_packet::tcp::TcpPacket;
Expand All @@ -12,6 +11,8 @@ use libpcap_tools::{Error, FiveTuple, ParseContext};

use super::fragmentation::two_tuple_proto_ipid::TwoTupleProtoIpid;
use super::fragmentation::two_tuple_proto_ipid_five_tuple::TwoTupleProtoIpidFiveTuple;
use crate::container::ipaddr_proto_port_container::IpAddrProtoPort;
use crate::filters::ipaddr_pair::IpAddrPair;

pub fn parse_src_ipaddr(ctx: &ParseContext, payload: &[u8]) -> Result<IpAddr, Error> {
let ipv4 = Ipv4Packet::new(payload).ok_or_else(|| {
Expand All @@ -35,10 +36,7 @@ pub fn parse_dst_ipaddr(ctx: &ParseContext, payload: &[u8]) -> Result<IpAddr, Er
Result::Ok(IpAddr::V4(ipv4.get_destination()))
}

pub fn parse_src_dst_ipaddr(
ctx: &ParseContext,
payload: &[u8],
) -> Result<(IpAddr, IpAddr), Error> {
pub fn parse_src_dst_ipaddr(ctx: &ParseContext, payload: &[u8]) -> Result<IpAddrPair, Error> {
let ipv4_packet = Ipv4Packet::new(payload).ok_or_else(|| {
warn!(
"Expected Ipv4 packet but could not parse at index {}",
Expand All @@ -48,13 +46,13 @@ pub fn parse_src_dst_ipaddr(
})?;
let src_ipaddr = IpAddr::V4(ipv4_packet.get_source());
let dst_ipaddr = IpAddr::V4(ipv4_packet.get_destination());
Result::Ok((src_ipaddr, dst_ipaddr))
Result::Ok(IpAddrPair::new(src_ipaddr, dst_ipaddr))
}

pub fn parse_src_ipaddr_proto_dst_port(
ctx: &ParseContext,
payload: &[u8],
) -> Result<(IpAddr, IpNextHeaderProtocol, u16), Error> {
) -> Result<IpAddrProtoPort, Error> {
let ipv4_packet = Ipv4Packet::new(payload).ok_or_else(|| {
warn!(
"Expected Ipv4 packet but could not parse at index {}",
Expand All @@ -71,7 +69,11 @@ pub fn parse_src_ipaddr_proto_dst_port(
match TcpPacket::new(ipv4_payload) {
Some(ref tcp) => {
let dst_port = tcp.get_destination();
Ok((src_ipaddr, IpNextHeaderProtocols::Tcp, dst_port))
Ok(IpAddrProtoPort::new(
src_ipaddr,
IpNextHeaderProtocols::Tcp,
dst_port,
))
}
None => {
warn!(
Expand All @@ -87,7 +89,11 @@ pub fn parse_src_ipaddr_proto_dst_port(
IpNextHeaderProtocols::Udp => match UdpPacket::new(ipv4_packet.payload()) {
Some(ref udp) => {
let dst_port = udp.get_destination();
Ok((src_ipaddr, IpNextHeaderProtocols::Udp, dst_port))
Ok(IpAddrProtoPort::new(
src_ipaddr,
IpNextHeaderProtocols::Udp,
dst_port,
))
}
None => {
warn!(
Expand All @@ -99,7 +105,11 @@ pub fn parse_src_ipaddr_proto_dst_port(
))
}
},
_ => Ok((src_ipaddr, ipv4_packet.get_next_level_protocol(), 0)),
_ => Ok(IpAddrProtoPort::new(
src_ipaddr,
ipv4_packet.get_next_level_protocol(),
0,
)),
}
}

Expand Down
Loading
Loading