diff --git a/.cargo/config.toml b/.cargo/config.toml new file mode 100644 index 0000000..f0ccbc9 --- /dev/null +++ b/.cargo/config.toml @@ -0,0 +1,2 @@ +[alias] +xtask = "run --package xtask --" \ No newline at end of file diff --git a/.github/workflows/rust.yml b/.github/workflows/rust.yml index 31000a2..a5c7012 100644 --- a/.github/workflows/rust.yml +++ b/.github/workflows/rust.yml @@ -1,4 +1,4 @@ -name: Rust +name: Rust Aya eBPF CI on: push: @@ -16,7 +16,23 @@ jobs: steps: - uses: actions/checkout@v3 - - name: Build + + - name: Set up Rust + uses: actions-rs/toolchain@v1 + with: + toolchain: nightly + profile: minimal + components: rust-src, llvm-tools-preview + override: true + + - name: Install bpf-linker + run: cargo install bpf-linker + + - name: Build ingress eBPF program + run: cargo xtask ingress-ebpf + + - name: Build egress eBPF program + run: cargo xtask egress-ebpf + + - name: Build userspace program run: cargo build --verbose - - name: Run tests - run: cargo test --verbose diff --git a/.gitignore b/.gitignore index 088ba6b..767dae2 100644 --- a/.gitignore +++ b/.gitignore @@ -1,6 +1,6 @@ # Generated by Cargo # will have compiled files and executables -/target/ +target/ # Remove Cargo.lock from gitignore if creating an executable, leave it for libraries # More information here https://doc.rust-lang.org/cargo/guide/cargo-toml-vs-cargo-lock.html diff --git a/Cargo.toml b/Cargo.toml index 4c14ac8..db89b9e 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,11 +1,2 @@ -[package] -name = "nids-feature-extraction-tool" -version = "0.1.0" -edition = "2021" - -# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html - -[dependencies] -clap = { version = "4.5.0", features = ["derive"] } -csv = "1.3.0" -serde = { version = "1.0.196", features = ["derive"] } +[workspace] +members = ["feature-extraction-tool", "common", "xtask"] \ No newline at end of file diff --git a/README.md b/README.md index eb69f56..57bb729 100644 --- a/README.md +++ b/README.md @@ -1,2 +1,70 @@ # Real-Time Adaptive Feature Extraction for ML-Based Network Intrusion Detection + This is a feature extraction tool build in Rust using eBPF for network intrusion detection + +## Install: + +### Prerequisites + +Make sure you have Rust installed: + +```bash +curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh +``` + +Installing nightly: + +```bash +rustup install stable +rustup toolchain install nightly --component rust-src +``` + +Installing the bpf linker +This is highly dependent on your operating system, just follow the error messages and install the requirements. For llvm you need version 18, make sure that Polly is installed with it. + +```bash +sudo apt install llvm +sudo apt install llvm-dev +sudo apt install libzstd-dev +``` + +Make sure you are in the project root directory +```bash +cargo install --no-default-features bpf-linker +``` + +When you are running Ubuntu 20.04 LTS you need to run this command to avoid bugs: + +```bash +sudo apt install linux-tools-5.8.0-63-generic +export PATH=/usr/lib/linux-tools/5.8.0-63-generic:$PATH +``` + +### Building the project + +To build the eBPF programs: + +```bash +cargo xtask ingress-ebpf +cargo xtask egress-ebpf +``` + +To build the user space programs: + +```bash +cargo build +``` + +### Running the project + +To run the program: + +```bash +RUST_LOG=info cargo xtask run -- realtime +``` + +To now the other possibilities, run this command: + +```bash +RUST_LOG=info cargo xtask run -- help +``` diff --git a/common/Cargo.toml b/common/Cargo.toml new file mode 100644 index 0000000..2e194a0 --- /dev/null +++ b/common/Cargo.toml @@ -0,0 +1,14 @@ +[package] +name = "common" +version = "0.1.0" +edition = "2021" + +[features] +default = [] +user = [ "aya" ] + +[dependencies] +aya = { git = "https://github.com/aya-rs/aya", optional=true } + +[lib] +path = "src/lib.rs" \ No newline at end of file diff --git a/common/src/lib.rs b/common/src/lib.rs new file mode 100644 index 0000000..33a7657 --- /dev/null +++ b/common/src/lib.rs @@ -0,0 +1,13 @@ +#![no_std] + +#[repr(C)] +#[derive(Clone, Copy)] +pub struct PacketLog { + pub ipv4_destination: u32, + pub ipv4_source: u32, + pub port_destination: u16, + pub port_source: u16, +} + +#[cfg(feature = "user")] +unsafe impl aya::Pod for PacketLog {} diff --git a/egress-ebpf/.cargo/config.toml b/egress-ebpf/.cargo/config.toml new file mode 100644 index 0000000..5d7e591 --- /dev/null +++ b/egress-ebpf/.cargo/config.toml @@ -0,0 +1,6 @@ +[build] +target-dir = "../target" +target = "bpfel-unknown-none" + +[unstable] +build-std = ["core"] \ No newline at end of file diff --git a/egress-ebpf/Cargo.toml b/egress-ebpf/Cargo.toml new file mode 100644 index 0000000..833ff57 --- /dev/null +++ b/egress-ebpf/Cargo.toml @@ -0,0 +1,34 @@ +[package] +name = "egress-ebpf" +version = "0.1.0" +edition = "2021" + +[dependencies] +aya-bpf = { git = "https://github.com/aya-rs/aya" } +aya-log-ebpf = { git = "https://github.com/aya-rs/aya" } +common = { path = "../common" } +memoffset = "0.8" +network-types = "0.0.4" + +[[bin]] +name = "feature-extraction-tool-egress" +path = "src/main.rs" + +[profile.dev] +opt-level = 3 +debug = false +debug-assertions = false +overflow-checks = false +lto = true +panic = "abort" +incremental = false +codegen-units = 1 +rpath = false + +[profile.release] +lto = true +panic = "abort" +codegen-units = 1 + +[workspace] +members = [] \ No newline at end of file diff --git a/egress-ebpf/rust-toolchain.toml b/egress-ebpf/rust-toolchain.toml new file mode 100644 index 0000000..271800c --- /dev/null +++ b/egress-ebpf/rust-toolchain.toml @@ -0,0 +1,2 @@ +[toolchain] +channel = "nightly" \ No newline at end of file diff --git a/egress-ebpf/src/main.rs b/egress-ebpf/src/main.rs new file mode 100644 index 0000000..0893035 --- /dev/null +++ b/egress-ebpf/src/main.rs @@ -0,0 +1,75 @@ +#![no_std] +#![no_main] +#![allow(nonstandard_style, dead_code)] + +use aya_bpf::{ + bindings::TC_ACT_PIPE, + macros::{classifier, map}, + maps::PerfEventArray, + programs::TcContext, +}; + +use common::PacketLog; + +use network_types::{ + eth::{EthHdr, EtherType}, + ip::{IpProto, Ipv4Hdr}, + tcp::TcpHdr, + udp::UdpHdr, +}; + +#[panic_handler] +fn panic(_info: &core::panic::PanicInfo) -> ! { + unsafe { core::hint::unreachable_unchecked() } +} + +#[map] +static EVENTS_EGRESS: PerfEventArray = PerfEventArray::with_max_entries(1024, 0); + +#[classifier] +pub fn tc_flow_track(ctx: TcContext) -> i32 { + match try_tc_flow_track(ctx) { + Ok(ret) => ret, + Err(_) => TC_ACT_PIPE, + } +} + +fn try_tc_flow_track(ctx: TcContext) -> Result { + let ethhdr: EthHdr = ctx.load(0).map_err(|_| ())?; + match ethhdr.ether_type { + EtherType::Ipv4 => {} + _ => return Ok(TC_ACT_PIPE), + } + + let ipv4hdr: Ipv4Hdr = ctx.load(EthHdr::LEN).map_err(|_| ())?; + let ipv4_destination = u32::from_be(ipv4hdr.dst_addr); + let ipv4_source = u32::from_be(ipv4hdr.src_addr); + + let source_port; + let destination_port; + match ipv4hdr.proto { + IpProto::Tcp => { + let tcphdr: TcpHdr = ctx.load(EthHdr::LEN + Ipv4Hdr::LEN).map_err(|_| ())?; + source_port = u16::from_be(tcphdr.source); + destination_port = u16::from_be(tcphdr.dest); + } + IpProto::Udp => { + let udphdr: UdpHdr = ctx.load(EthHdr::LEN + Ipv4Hdr::LEN).map_err(|_| ())?; + source_port = u16::from_be(udphdr.source); + destination_port = u16::from_be(udphdr.dest); + } + _ => return Ok(TC_ACT_PIPE), + }; + + let flow = PacketLog { + ipv4_destination: ipv4_destination, + ipv4_source: ipv4_source, + port_destination: destination_port, + port_source: source_port, + }; + + // the zero value is a flag + EVENTS_EGRESS.output(&ctx, &flow, 0); + + Ok(TC_ACT_PIPE) +} diff --git a/feature-extraction-tool/Cargo.toml b/feature-extraction-tool/Cargo.toml new file mode 100644 index 0000000..c4e51c6 --- /dev/null +++ b/feature-extraction-tool/Cargo.toml @@ -0,0 +1,30 @@ +[package] +name = "nids-feature-extraction-tool" +version = "0.1.0" +edition = "2021" +publish = false + +# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html + +[dependencies] +clap = { version = "4.5.0", features = ["derive"] } +csv = "1.3.0" +serde = { version = "1.0.196", features = ["derive"] } +aya = { git = "https://github.com/aya-rs/aya", features = ["async_tokio"] } +aya-log = { git = "https://github.com/aya-rs/aya"} +common = { path = "../common", features = ["user"] } +anyhow = "1" +log = "0.4" +tokio = { version = "1.25", features = [ + "macros", + "rt", + "rt-multi-thread", + "net", + "signal", +] } +bytes = "1" +env_logger = "0.11" + +[[bin]] +name = "feature-extraction-tool" +path = "src/main.rs" \ No newline at end of file diff --git a/src/args.rs b/feature-extraction-tool/src/args.rs similarity index 91% rename from src/args.rs rename to feature-extraction-tool/src/args.rs index be34910..9e726e9 100644 --- a/src/args.rs +++ b/feature-extraction-tool/src/args.rs @@ -10,7 +10,10 @@ pub struct Cli { #[derive(Debug, Subcommand)] pub enum Commands { /// Real-time feature extraction - Realtime, + Realtime { + /// The network interface to capture packets from + interface: String, + }, /// Feature extraction from a dataset Dataset { diff --git a/feature-extraction-tool/src/main.rs b/feature-extraction-tool/src/main.rs new file mode 100644 index 0000000..0e0139d --- /dev/null +++ b/feature-extraction-tool/src/main.rs @@ -0,0 +1,209 @@ +mod args; +mod parsers; +mod records; + +use crate::{ + parsers::csv_parser::CsvParser, + records::{cic_record::CicRecord, print::Print}, +}; + +use anyhow::Context; +use args::{Cli, Commands, Dataset}; +use aya::{ + include_bytes_aligned, + maps::AsyncPerfEventArray, + programs::{tc, SchedClassifier, TcAttachType, Xdp, XdpFlags}, + util::online_cpus, + Bpf, +}; +use aya_log::BpfLogger; +use bytes::BytesMut; +use clap::Parser; +use common::PacketLog; +use core::panic; +use log::{info, warn}; +use std::net::Ipv4Addr; +use tokio::{signal, task}; + +#[tokio::main] +async fn main() { + let cli = Cli::parse(); + + match cli.command { + Commands::Realtime { interface } => { + if let Err(err) = handle_realtime(interface).await { + eprintln!("Error: {:?}", err); + } + } + Commands::Dataset { dataset, path } => { + handle_dataset(dataset, &path); + } + } +} + +async fn handle_realtime(interface: String) -> Result<(), anyhow::Error> { + env_logger::init(); + + // This will include your eBPF object file as raw bytes at compile-time and load it at + // runtime. This approach is recommended for most real-world use cases. If you would + // like to specify the eBPF program at runtime rather than at compile-time, you can + // reach for `Bpf::load_file` instead. + + // Loading the eBPF program for egress, the macros make sure the correct file is loaded + #[cfg(debug_assertions)] + let mut bpf_egress = Bpf::load(include_bytes_aligned!( + "../../target/bpfel-unknown-none/debug/feature-extraction-tool-egress" + ))?; + #[cfg(not(debug_assertions))] + let mut bpf_egress = Bpf::load(include_bytes_aligned!( + "../../target/bpfel-unknown-none/release/feature-extraction-tool-egress" + ))?; + + // Loading the eBPF program for ingress, the macros make sure the correct file is loaded + #[cfg(debug_assertions)] + let mut bpf_ingress = Bpf::load(include_bytes_aligned!( + "../../target/bpfel-unknown-none/debug/feature-extraction-tool-ingress" + ))?; + #[cfg(not(debug_assertions))] + let mut bpf_ingress = Bpf::load(include_bytes_aligned!( + "../../target/bpfel-unknown-none/release/feature-extraction-tool-ingress" + ))?; + + // You can remove this when you don't log anything in your egress eBPF program. + if let Err(e) = BpfLogger::init(&mut bpf_egress) { + warn!("failed to initialize the egress eBPF logger: {}", e); + } + + // You can remove this when you don't log anything in your ingress eBPF program. + if let Err(e) = BpfLogger::init(&mut bpf_ingress) { + warn!("failed to initialize the ingress eBPF logger: {}", e); + } + + // Loading and attaching the eBPF program function for egress + let _ = tc::qdisc_add_clsact(interface.as_str()); + let program_egress: &mut SchedClassifier = bpf_egress + .program_mut("tc_flow_track") + .unwrap() + .try_into()?; + program_egress.load()?; + program_egress.attach(&interface, TcAttachType::Egress)?; + + // Loading and attaching the eBPF program function for ingress + let program_ingress: &mut Xdp = bpf_ingress + .program_mut("xdp_flow_track") + .unwrap() + .try_into()?; + program_ingress.load()?; + program_ingress.attach(&interface, XdpFlags::default()) + .context("failed to attach the XDP program with default flags - try changing XdpFlags::default() to XdpFlags::SKB_MODE")?; + + // Attach to the event arrays + let mut flows_egress = + AsyncPerfEventArray::try_from(bpf_egress.take_map("EVENTS_EGRESS").unwrap())?; + + let mut flows_ingress = + AsyncPerfEventArray::try_from(bpf_ingress.take_map("EVENTS_INGRESS").unwrap())?; + + // Use all online CPUs to process the events in the user space + for cpu_id in online_cpus()? { + let mut buf_egress = flows_egress.open(cpu_id, None)?; + + task::spawn(async move { + let mut buffers = (0..10) + .map(|_| BytesMut::with_capacity(1024)) + .collect::>(); + + loop { + let events = buf_egress.read_events(&mut buffers).await.unwrap(); + for buf in buffers.iter_mut().take(events.read) { + let ptr = buf.as_ptr() as *const PacketLog; + let data = unsafe { ptr.read_unaligned() }; + + let src_addr = Ipv4Addr::from(data.ipv4_source); + let dst_addr = Ipv4Addr::from(data.ipv4_destination); + let src_port = data.port_source; + let dst_port = data.port_destination; + + info!( + "LOG: SRC {}:{}, DST {}:{}", + src_addr, src_port, dst_addr, dst_port + ); + } + } + }); + + let mut buf_ingress = flows_ingress.open(cpu_id, None)?; + task::spawn(async move { + let mut buffers = (0..10) + .map(|_| BytesMut::with_capacity(1024)) + .collect::>(); + + loop { + let events = buf_ingress.read_events(&mut buffers).await.unwrap(); + for buf in buffers.iter_mut().take(events.read) { + let ptr = buf.as_ptr() as *const PacketLog; + let data = unsafe { ptr.read_unaligned() }; + + let src_addr = Ipv4Addr::from(data.ipv4_source); + let dst_addr = Ipv4Addr::from(data.ipv4_destination); + let src_port = data.port_source; + let dst_port = data.port_destination; + + info!( + "LOG: SRC {}:{}, DST {}:{}", + src_addr, src_port, dst_addr, dst_port + ); + } + } + }); + } + + info!("Waiting for Ctrl-C..."); + signal::ctrl_c().await?; + info!("Exiting..."); + + Ok(()) +} + +fn handle_dataset(dataset: Dataset, path: &str) { + println!( + "Dataset feature extraction for {:?} from path: {}", + dataset, path + ); + + match dataset { + Dataset::CicIds2017 => { + if path.ends_with(".csv") { + let parser = CsvParser; + + match parser.parse::(path) { + Ok(records) => { + for record in records { + match record { + Ok(record) => { + record.print(); + } + Err(err) => { + // TODO: Will we output to stderr, drop the record or use default values? + eprintln!("Error: {:?}", err); + } + } + } + } + Err(err) => { + eprintln!("Error: {:?}", err); + } + } + } else if path.ends_with(".pcap") { + panic!("This file format is not supported yet..."); + } else if path.ends_with(".parquet") { + panic!("This file format is not supported yet..."); + } else { + panic!("This file format is not supported..."); + } + } + _ => { + panic!("This is not implemented yet..."); + } + } +} diff --git a/src/parsers/csv_parser.rs b/feature-extraction-tool/src/parsers/csv_parser.rs similarity index 100% rename from src/parsers/csv_parser.rs rename to feature-extraction-tool/src/parsers/csv_parser.rs diff --git a/src/parsers/mod.rs b/feature-extraction-tool/src/parsers/mod.rs similarity index 100% rename from src/parsers/mod.rs rename to feature-extraction-tool/src/parsers/mod.rs diff --git a/src/parsers/parser.rs b/feature-extraction-tool/src/parsers/parser.rs similarity index 100% rename from src/parsers/parser.rs rename to feature-extraction-tool/src/parsers/parser.rs diff --git a/src/records/cic_record.rs b/feature-extraction-tool/src/records/cic_record.rs similarity index 100% rename from src/records/cic_record.rs rename to feature-extraction-tool/src/records/cic_record.rs diff --git a/src/records/mod.rs b/feature-extraction-tool/src/records/mod.rs similarity index 100% rename from src/records/mod.rs rename to feature-extraction-tool/src/records/mod.rs diff --git a/src/records/print.rs b/feature-extraction-tool/src/records/print.rs similarity index 100% rename from src/records/print.rs rename to feature-extraction-tool/src/records/print.rs diff --git a/ingress-ebpf/.cargo/config.toml b/ingress-ebpf/.cargo/config.toml new file mode 100644 index 0000000..5d7e591 --- /dev/null +++ b/ingress-ebpf/.cargo/config.toml @@ -0,0 +1,6 @@ +[build] +target-dir = "../target" +target = "bpfel-unknown-none" + +[unstable] +build-std = ["core"] \ No newline at end of file diff --git a/ingress-ebpf/Cargo.toml b/ingress-ebpf/Cargo.toml new file mode 100644 index 0000000..06ec343 --- /dev/null +++ b/ingress-ebpf/Cargo.toml @@ -0,0 +1,33 @@ +[package] +name = "ingress-ebpf" +version = "0.1.0" +edition = "2021" + +[dependencies] +aya-bpf = { git = "https://github.com/aya-rs/aya" } +aya-log-ebpf = { git = "https://github.com/aya-rs/aya" } +common = { path = "../common" } +network-types = "0.0.4" + +[[bin]] +name = "feature-extraction-tool-ingress" +path = "src/main.rs" + +[profile.dev] +opt-level = 3 +debug = false +debug-assertions = false +overflow-checks = false +lto = true +panic = "abort" +incremental = false +codegen-units = 1 +rpath = false + +[profile.release] +lto = true +panic = "abort" +codegen-units = 1 + +[workspace] +members = [] \ No newline at end of file diff --git a/ingress-ebpf/rust-toolchain.toml b/ingress-ebpf/rust-toolchain.toml new file mode 100644 index 0000000..271800c --- /dev/null +++ b/ingress-ebpf/rust-toolchain.toml @@ -0,0 +1,2 @@ +[toolchain] +channel = "nightly" \ No newline at end of file diff --git a/ingress-ebpf/rustfmt.toml b/ingress-ebpf/rustfmt.toml new file mode 100644 index 0000000..39f97b0 --- /dev/null +++ b/ingress-ebpf/rustfmt.toml @@ -0,0 +1 @@ +../rustfmt.toml \ No newline at end of file diff --git a/ingress-ebpf/src/main.rs b/ingress-ebpf/src/main.rs new file mode 100644 index 0000000..37270a3 --- /dev/null +++ b/ingress-ebpf/src/main.rs @@ -0,0 +1,93 @@ +#![no_std] +#![no_main] +#![allow(nonstandard_style, dead_code)] + +use aya_bpf::{ + bindings::xdp_action, + macros::{map, xdp}, + programs::XdpContext, + maps::PerfEventArray +}; + +use common::PacketLog; + +use core::mem; +use network_types::{ + eth::{EthHdr, EtherType}, + ip::{Ipv4Hdr, IpProto}, + tcp::TcpHdr, + udp::UdpHdr, +}; + +#[panic_handler] +fn panic(_info: &core::panic::PanicInfo) -> ! { + unsafe { core::hint::unreachable_unchecked() } +} + +#[map] +static EVENTS_INGRESS: PerfEventArray = + PerfEventArray::with_max_entries(1024, 0); + +#[xdp] +pub fn xdp_flow_track(ctx: XdpContext) -> u32{ + match try_xdp_flow_track(ctx) { + Ok(ret) => ret, + Err(_) => xdp_action::XDP_ABORTED, + } +} + +#[inline(always)] +unsafe fn ptr_at(ctx: &XdpContext, offset: usize) -> Result<*const T, ()> { + let start = ctx.data(); + let end = ctx.data_end(); + let len = mem::size_of::(); + + if start + offset + len > end { + return Err(()); + } + + let ptr = (start + offset) as *const T; + Ok(&*ptr) +} + +fn try_xdp_flow_track(ctx: XdpContext) -> Result{ + let ethhdr: *const EthHdr = unsafe { ptr_at(&ctx, 0)? }; + match unsafe { (*ethhdr).ether_type } { + EtherType::Ipv4 => {} + _ => return Ok(xdp_action::XDP_PASS), + } + + let ipv4hdr: *const Ipv4Hdr = unsafe { ptr_at(&ctx, EthHdr::LEN)? }; + let ipv4_source = u32::from_be(unsafe { (*ipv4hdr).src_addr }); + let ipv4_destination = u32::from_be(unsafe { (*ipv4hdr).dst_addr }); + + let source_port; + let destination_port; + match unsafe { *ipv4hdr }.proto { + IpProto::Tcp => { + let tcphdr: *const TcpHdr = + unsafe { ptr_at(&ctx, EthHdr::LEN + Ipv4Hdr::LEN) }?; + source_port = u16::from_be(unsafe { *tcphdr }.source); + destination_port = u16::from_be(unsafe { *tcphdr }.dest); + } + IpProto::Udp => { + let udphdr: *const UdpHdr = + unsafe { ptr_at(&ctx, EthHdr::LEN + Ipv4Hdr::LEN) }?; + source_port = u16::from_be(unsafe { *udphdr }.source); + destination_port = u16::from_be(unsafe { *udphdr }.dest); + } + _ => return Ok(xdp_action::XDP_ABORTED), + }; + + let flow = PacketLog { + ipv4_destination: ipv4_destination, + ipv4_source: ipv4_source, + port_destination: destination_port, + port_source: source_port, + }; + + // the zero value is a flag + EVENTS_INGRESS.output(&ctx, &flow, 0); + + Ok(xdp_action::XDP_PASS) +} \ No newline at end of file diff --git a/src/main.rs b/src/main.rs deleted file mode 100644 index 8f8d540..0000000 --- a/src/main.rs +++ /dev/null @@ -1,75 +0,0 @@ -mod args; -mod parsers; -mod records; - -use core::panic; - -use args::{Cli, Commands, Dataset}; -use clap::Parser; - -use crate::{ - parsers::csv_parser::CsvParser, - records::{cic_record::CicRecord, print::Print}, -}; - -fn main() { - let cli = Cli::parse(); - - match cli.command { - Commands::Realtime => { - handle_realtime(); - } - Commands::Dataset { dataset, path } => { - handle_dataset(dataset, &path); - } - } -} - -fn handle_realtime() { - println!("Real-time feature extraction"); -} - -fn handle_dataset(dataset: Dataset, path: &str) { - println!( - "Dataset feature extraction for {:?} from path: {}", - dataset, path - ); - - match dataset { - Dataset::CicIds2017 => { - if path.ends_with(".csv") { - let parser = CsvParser; - - match parser.parse::(path) { - Ok(records) => { - for record in records { - match record { - Ok(record) => { - record.print(); - } - Err(err) => { - // TODO: Will we output to stderr, drop the record or use default values? - eprintln!("Error: {:?}", err); - } - } - } - } - Err(err) => { - eprintln!("Error: {:?}", err); - } - } - } else if path.ends_with(".pcap") { - panic!("This file format is not supported yet..."); - - } else if path.ends_with(".parquet") { - panic!("This file format is not supported yet..."); - - } else { - panic!("This file format is not supported..."); - } - } - _ => { - panic!("This is not implemented yet..."); - } - } -} diff --git a/xtask/Cargo.toml b/xtask/Cargo.toml new file mode 100644 index 0000000..bc13a13 --- /dev/null +++ b/xtask/Cargo.toml @@ -0,0 +1,10 @@ +[package] +name = "xtask" +version = "0.1.0" +edition = "2021" + +# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html +[dependencies] +anyhow = "1" +clap = { version = "4.1", features = ["derive"] } +aya-tool = { git = "https://github.com/aya-rs/aya" } \ No newline at end of file diff --git a/xtask/src/build_ebpf.rs b/xtask/src/build_ebpf.rs new file mode 100644 index 0000000..294279b --- /dev/null +++ b/xtask/src/build_ebpf.rs @@ -0,0 +1,63 @@ +use std::{path::PathBuf, process::Command}; + +use clap::Parser; + +#[derive(Debug, Copy, Clone)] +pub enum Architecture { + BpfEl, + BpfEb, +} + +impl std::str::FromStr for Architecture { + type Err = String; + + fn from_str(s: &str) -> Result { + Ok(match s { + "bpfel-unknown-none" => Architecture::BpfEl, + "bpfeb-unknown-none" => Architecture::BpfEb, + _ => return Err("invalid target".to_owned()), + }) + } +} + +impl std::fmt::Display for Architecture { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str(match self { + Architecture::BpfEl => "bpfel-unknown-none", + Architecture::BpfEb => "bpfeb-unknown-none", + }) + } +} + +#[derive(Debug, Parser)] +pub struct Options { + /// Set the endianness of the BPF target + #[clap(default_value = "bpfel-unknown-none", long)] + pub target: Architecture, + /// Build the release target + #[clap(long)] + pub release: bool, +} + +pub fn build_ebpf(opts: Options, target: String) -> Result<(), anyhow::Error> { + let dir = PathBuf::from(target); + let target = format!("--target={}", opts.target); + let mut args = vec!["build", target.as_str(), "-Z", "build-std=core"]; + if opts.release { + args.push("--release") + } + + // Command::new creates a child process which inherits all env variables. This means env + // vars set by the cargo xtask command are also inherited. RUSTUP_TOOLCHAIN is removed + // so the rust-toolchain.toml file in the -ebpf folder is honored. + + let status = Command::new("cargo") + .current_dir(&dir) + .env_remove("RUSTUP_TOOLCHAIN") + .args(&args) + .status() + .expect("failed to build bpf program"); + assert!(status.success()); + + Ok(()) +} diff --git a/xtask/src/main.rs b/xtask/src/main.rs new file mode 100644 index 0000000..cf2c431 --- /dev/null +++ b/xtask/src/main.rs @@ -0,0 +1,37 @@ +mod build_ebpf; +mod run; + +use std::process::exit; + +use clap::Parser; + +#[derive(Debug, Parser)] +pub struct Options { + #[clap(subcommand)] + command: Command, +} + +#[derive(Debug, Parser)] +enum Command { + #[clap(name = "ingress-ebpf")] + BuildIngressEbpf(build_ebpf::Options), + #[clap(name = "egress-ebpf")] + BuildEgressEbpf(build_ebpf::Options), + Run(run::Options), +} + +fn main() { + let opts = Options::parse(); + + use Command::*; + let ret = match opts.command { + BuildIngressEbpf(opts) => build_ebpf::build_ebpf(opts, "ingress-ebpf".to_string()), + BuildEgressEbpf(opts) => build_ebpf::build_ebpf(opts, "egress-ebpf".to_string()), + Run(opts) => run::run(opts), + }; + + if let Err(e) = ret { + eprintln!("{e:#}"); + exit(1); + } +} diff --git a/xtask/src/run.rs b/xtask/src/run.rs new file mode 100644 index 0000000..9ea19c7 --- /dev/null +++ b/xtask/src/run.rs @@ -0,0 +1,83 @@ +use std::process::Command; + +use anyhow::Context as _; +use clap::Parser; + +use crate::build_ebpf::{build_ebpf, Architecture, Options as BuildOptions}; + +#[derive(Debug, Parser)] +pub struct Options { + /// Set the endianness of the BPF target + #[clap(default_value = "bpfel-unknown-none", long)] + pub bpf_target: Architecture, + /// Build and run the release target + #[clap(long)] + pub release: bool, + /// The command used to wrap your application + #[clap(short, long, default_value = "sudo -E")] + pub runner: String, + /// Arguments to pass to your application + #[clap(name = "args", last = true)] + pub run_args: Vec, +} + +/// Build the project +fn build(opts: &Options) -> Result<(), anyhow::Error> { + let mut args = vec!["build"]; + if opts.release { + args.push("--release") + } + let status = Command::new("cargo") + .args(&args) + .status() + .expect("failed to build userspace"); + assert!(status.success()); + Ok(()) +} + +/// Build and run the project +pub fn run(opts: Options) -> Result<(), anyhow::Error> { + // build our ebpf program followed by our application + build_ebpf( + BuildOptions { + target: opts.bpf_target, + release: opts.release, + }, + "ingress-ebpf".to_string(), + ) + .context("Error while building the ingress-eBPF program")?; + build(&opts).context("Error while building userspace application")?; + + build_ebpf( + BuildOptions { + target: opts.bpf_target, + release: opts.release, + }, + "egress-ebpf".to_string(), + ) + .context("Error while building the egress-eBPF program")?; + build(&opts).context("Error while building userspace application")?; + + // profile we are building (release or debug) + let profile = if opts.release { "release" } else { "debug" }; + let bin_path = format!("target/{profile}/feature-extraction-tool"); + + // arguments to pass to the application + let mut run_args: Vec<_> = opts.run_args.iter().map(String::as_str).collect(); + + // configure args + let mut args: Vec<_> = opts.runner.trim().split_terminator(' ').collect(); + args.push(bin_path.as_str()); + args.append(&mut run_args); + + // run the command + let status = Command::new(args.first().expect("No first argument")) + .args(args.iter().skip(1)) + .status() + .expect("failed to run the command"); + + if !status.success() { + anyhow::bail!("Failed to run `{}`", args.join(" ")); + } + Ok(()) +}