ipblc/src/fw.rs

use crate::ip::IpData;
use crate::ipblc::PKG_NAME;

use nftnl::{nft_expr, Batch, Chain, FinalizedBatch, ProtoFamily, Rule, Table};
use std::{
    ffi::CString,
    io::Error,
    net::{Ipv4Addr, Ipv6Addr},
};

pub enum FwTableType {
    IPv4,
    IPv6,
}

pub fn fwglobalinit<'a>() -> ((Batch, Table), (Batch, Table)) {
    let (batch4, table4) = fwinit(FwTableType::IPv4);
    let (batch6, table6) = fwinit(FwTableType::IPv6);
    ((batch4, table4), (batch6, table6))
}

fn fwinit(t: FwTableType) -> (Batch, Table) {
    let table_name: String;
    let table: Table;
    match t {
        FwTableType::IPv4 => {
            table_name = format!("{PKG_NAME}4");
            table = Table::new(
                &CString::new(format!("{table_name}")).unwrap(),
                ProtoFamily::Ipv4,
            );
        }
        FwTableType::IPv6 => {
            table_name = format!("{PKG_NAME}6");
            table = Table::new(
                &CString::new(format!("{table_name}")).unwrap(),
                ProtoFamily::Ipv6,
            );
        }
    }
    let mut batch = Batch::new();

    batch.add(&table, nftnl::MsgType::Add);
    batch.add(&table, nftnl::MsgType::Del);

    batch.add(&table, nftnl::MsgType::Add);
    (batch, table)
}

pub fn fwblock(
    ips_add: &Vec<IpData>,
    ret: &mut Vec<String>,
    fwlen: &mut usize,
) -> std::result::Result<(), Error> {
    let ((mut batch4, table4), (mut batch6, table6)) = fwglobalinit();

    // build chain for ipv4
    let mut chain4 = Chain::new(&CString::new(PKG_NAME).unwrap(), &table4);
    chain4.set_hook(nftnl::Hook::In, 1);
    chain4.set_policy(nftnl::Policy::Accept);

    // add chain
    batch4.add(&chain4, nftnl::MsgType::Add);

    batch4.add(&Rule::new(&chain4), nftnl::MsgType::Del);

    let mut rule4 = Rule::new(&chain4);
    rule4.add_expr(&nft_expr!(ct state));
    rule4.add_expr(&nft_expr!(bitwise mask 4u32, xor 0u32));
    rule4.add_expr(&nft_expr!(cmp != 0u32));
    rule4.add_expr(&nft_expr!(counter));
    rule4.add_expr(&nft_expr!(verdict accept));
    batch4.add(&rule4, nftnl::MsgType::Add);

    // build chain for ipv6
    let mut chain6 = Chain::new(&CString::new(PKG_NAME).unwrap(), &table6);
    chain6.set_hook(nftnl::Hook::In, 1);
    chain6.set_policy(nftnl::Policy::Accept);

    // add chain
    batch6.add(&chain6, nftnl::MsgType::Add);

    batch6.add(&Rule::new(&chain6), nftnl::MsgType::Del);

    let mut rule6 = Rule::new(&chain6);
    rule6.add_expr(&nft_expr!(ct state));
    rule6.add_expr(&nft_expr!(bitwise mask 4u32, xor 0u32));
    rule6.add_expr(&nft_expr!(cmp != 0u32));
    rule6.add_expr(&nft_expr!(counter));
    rule6.add_expr(&nft_expr!(verdict accept));
    batch6.add(&rule6, nftnl::MsgType::Add);

    // build and add rules
    for ipdata in ips_add.clone() {
        match ipdata.t {
            4 => {
                let ip = ipdata.ip.parse::<Ipv4Addr>().unwrap();
                let mut rule = Rule::new(&chain4);
                rule.add_expr(&nft_expr!(payload ipv4 saddr));
                rule.add_expr(&nft_expr!(cmp == ip));
                rule.add_expr(&nft_expr!(ct state));
                rule.add_expr(&nft_expr!(bitwise mask 10u32, xor 0u32));
                rule.add_expr(&nft_expr!(cmp != 0u32));
                rule.add_expr(&nft_expr!(counter));
                rule.add_expr(&nft_expr!(verdict drop));
                batch4.add(&rule, nftnl::MsgType::Add);
            }
            6 => {
                let ip = ipdata.ip.parse::<Ipv6Addr>().unwrap();
                let mut rule = Rule::new(&chain6);
                rule.add_expr(&nft_expr!(payload ipv6 saddr));
                rule.add_expr(&nft_expr!(cmp == ip));
                rule.add_expr(&nft_expr!(ct state));
                rule.add_expr(&nft_expr!(bitwise mask 10u32, xor 0u32));
                rule.add_expr(&nft_expr!(cmp != 0u32));
                rule.add_expr(&nft_expr!(counter));
                rule.add_expr(&nft_expr!(verdict drop));
                batch6.add(&rule, nftnl::MsgType::Add);
            }
            _ => {
                todo!()
            }
        }
    }

    // validate and send batch
    for b in [batch4, batch6] {
        let bf = b.finalize();
        send_and_process(&bf).unwrap();
    }
    if fwlen != &mut ips_add.len() {
        ret.push(format!("{length} ip in firewall", length = ips_add.len()));
    }
    *fwlen = ips_add.len();
    Ok(())
}

fn send_and_process(batch: &FinalizedBatch) -> std::result::Result<(), Error> {
    let seq: u32 = 2;
    let socket = mnl::Socket::new(mnl::Bus::Netfilter)?;
    socket.send_all(batch)?;
    let mut buffer = vec![0; nftnl::nft_nlmsg_maxsize() as usize];
    while let Some(message) = socket_recv(&socket, &mut buffer[..])? {
        match mnl::cb_run(message, seq, socket.portid())? {
            mnl::CbResult::Stop => {
                break;
            }
            mnl::CbResult::Ok => (),
        }
    }
    Ok(())
}

fn socket_recv<'a>(
    socket: &mnl::Socket,
    buf: &'a mut [u8],
) -> std::result::Result<Option<&'a [u8]>, Error> {
    let ret = socket.recv(buf)?;
    if ret > 0 {
        Ok(Some(&buf[..ret]))
    } else {
        Ok(None)
    }
}