rescan/responder/src/main.rs

use std::{path::PathBuf, time::Instant, fs::File, io::{Read, self, Write}, os::unix::prelude::AsRawFd, process};

use aya::{Bpf, maps::Array as BpfArray};
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 responder_common::filter;
use responder_common::filter::bloom_filter;

#[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(long, default_value = "bitmap")]
    filter_type: String,
    #[clap(long)]
    filter_path: PathBuf,
    #[clap(default_value = DEFAULT_TARGET, long)]
    target: PathBuf,
    #[clap(short, long)]
    fd_info_out_path: Option<PathBuf>
}

#[tokio::main]
async fn main() -> Result<(), anyhow::Error> {
    let start = Instant::now();
    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 = format!("{}-{}", opt.scan_type, opt.filter_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; FD: {}", xdp_name, program.fd().unwrap().as_raw_fd());


    if let Some(path) = opt.fd_info_out_path {
        info!("Emitting pid and fd to {:?}", &path);
        let info = format!("{{\"fd\": {}, \"pid\": {}}}\n", program.fd().unwrap().as_raw_fd(), process::id());
        File::create(path)?.write_all(info.as_bytes())?;
    }

    info!("Installing filter rules from {:?}", opt.filter_path);
    let mut filter_map: BpfArray<_, [u8; filter::CHUNK_BYTES]> =
        BpfArray::try_from(bpf.map_mut("FILTER_MAP")
                           .ok_or(anyhow!("Could not construct mutable FILTER_MAP"))?)?;
    let mut filter_binary = File::open(opt.filter_path)?;
    let mut dimensions_buf: [u8; 8] = [0u8; 8];

    filter_binary.read_exact(&mut dimensions_buf)?;
    let file_address_bits = u64::from_be_bytes(dimensions_buf);
    if file_address_bits != (filter::ADDRESS_BITS as u64) {
        return Err(anyhow!(
            "Filter dimension ADDRESS_BITS does not match binary. Please rebuild the filter or recompile the binaries\n\tFile:\t{}\n\tBinary:\t{}",
            file_address_bits, filter::ADDRESS_BITS
        ));
    }

    filter_binary.read_exact(&mut dimensions_buf)?;
    let file_address_bits_chunk = u64::from_be_bytes(dimensions_buf);
    if file_address_bits_chunk != (filter::ADDRESS_BITS_CHUNK as u64) {
        return Err(anyhow!(
            "Filter dimension ADDRESS_BITS_CHUNK does not match binary. Please rebuild the filter or recompile the binaries\n\tFile:\t{}\n\tBinary:\t{}",
            file_address_bits_chunk, filter::ADDRESS_BITS_CHUNK
        ));
    }

    if opt.filter_type == "bloom" {
        filter_binary.read_exact(&mut dimensions_buf)?;
        let file_hash_count = u64::from_be_bytes(dimensions_buf);
        if file_hash_count != (bloom_filter::HASH_COUNT as u64) {
            return Err(anyhow!(
                "Filter dimension HASH_COUNT does not match binary. Please rebuild the filter or recompile the binaries\n\tFile:\t{}\n\tBinary:\t{}",
                file_hash_count, bloom_filter::HASH_COUNT
            ));
        }
    }

    let mut chunk = [0 as u8; filter::CHUNK_BYTES];

    let mut i = 0;
    loop {
        let read = filter_binary.read(&mut chunk)?;
        if read > 0 {
            filter_map.set(i, chunk, 0)?;
            i += 1;
        }

        if read < filter::CHUNK_BYTES {
            break;
        }
    }
    if (i as usize) != filter::MAP_SIZE {
        return Err(anyhow!(
            "Number of chunks in file ({}) do not match expected number ({})",
            i, filter::MAP_SIZE
        ))
    }

    info!("Loading XDP program took {:?}", start.elapsed());
    println!("press any key to exit");
    let mut buf = [0u8];
    io::stdin().read_exact(&mut buf)?;
    // signal::ctrl_c().await?;
    info!("Exiting...");

    Ok(())
}