rescan/responder/src/main.rs

use std::{net::Ipv4Addr, path::PathBuf};

use aya::{Bpf, maps::HashMap};
use anyhow::{anyhow, Context};
use aya::programs::{Xdp, XdpFlags};
use aya_log::BpfLogger;
use clap::Parser;
use env_logger::Env;
use log::{info, warn};
use tokio::signal;
use csv::ReaderBuilder;
use serde::Deserialize;

#[cfg(all(debug_assertions, not(feature = "default_artifact_build")))]
const DEFAULT_TARGET: &str = "target/bpfel-unknown-none/debug";

#[cfg(all(not(debug_assertions), not(feature = "default_artifact_build")))]
const DEFAULT_TARGET: &str = "target/bpfel-unknown-none/release";

#[cfg(feature = "default_artifact_build")]
const DEFAULT_TARGET: &str = "bin/ebpf";

#[derive(Debug, Parser)]
struct Opt {
    #[clap(short, long, default_value = "lo")]
    iface: String,
    #[clap(short, long, default_value = "syn")]
    scan_type: String,
    #[clap(short, long)]
    csv: Option<String>,
    #[clap(default_value = DEFAULT_TARGET, long)]
    target: PathBuf,
}

#[derive(Debug, Deserialize, Eq, PartialEq)]
struct CsvRow {
    saddr: Ipv4Addr,
}

#[tokio::main]
async fn main() -> Result<(), anyhow::Error> {
    let opt = Opt::parse();

    env_logger::Builder::from_env(Env::default().default_filter_or("info")).init();

    let mut bpf_path = PathBuf::new();
    bpf_path.push(opt.target);

    let xdp_name = opt.scan_type;
    bpf_path.push(xdp_name.clone());

    let mut bpf = Bpf::load_file(bpf_path)?;

    if let Err(e) = BpfLogger::init(&mut bpf) {
        warn!("failed to initialize eBPF logger: {}
This can happen if the loaded eBPF program has no log statements.", e);
    }

    // Obtain and load the XDP program called "responder" defined in the XDP file loaded above
    let program: &mut Xdp = bpf.program_mut("responder").unwrap().try_into()?;
    program.load()?;
    program.attach(&opt.iface, XdpFlags::default())
           .context(format!(
               "failed to attach the {} XDP program with default flags \
 - try changing XdpFlags::default() to XdpFlags::SKB_MODE",
               xdp_name
           ))?;

    info!("Loaded {} XDP program", xdp_name);

    let mut filter_map: HashMap<_, u32, u8> =
        HashMap::try_from(bpf.map_mut("FILTER_MAP")
                          .ok_or(anyhow!("Could not construct mutable FILTER_MAP"))?)?;

    if let Some(csv_path) = opt.csv {
        info!("Installing filter rules from {}", csv_path);
        let mut reader = ReaderBuilder::new()
            .has_headers(true)
            .from_path(csv_path)?;
        for record in reader.deserialize() {
            let row: CsvRow = record?;
            filter_map.insert(u32::from(row.saddr), 1u8, 0)?;
            // TODO replace with BPF_MAP_UPDATE_BATCH
        }
    };

    info!("press ctrl-c to exit");
    signal::ctrl_c().await?;
    info!("Exiting...");

    Ok(())
}