From 1e2f04782436cc50d06c7b332c00a57ae149701d Mon Sep 17 00:00:00 2001
From: Paul Lecuq <paul@paulbsd.com>
Date: Wed, 3 Jan 2024 21:44:00 +0100
Subject: [PATCH] add ips in chunks to nftables

---
 src/fw.rs | 47 ++++++++++++++++++++++++++++++++---------------
 1 file changed, 32 insertions(+), 15 deletions(-)

diff --git a/src/fw.rs b/src/fw.rs
index e29dc86..81d5302 100644
--- a/src/fw.rs
+++ b/src/fw.rs
@@ -21,7 +21,6 @@ pub fn fwglobalinit<'a>() -> ((Batch, Table), (Batch, Table)) {
 
 macro_rules! initrules {
     ($batch:expr, $table:expr, $chain:ident) => {
-        let mut $chain = Chain::new(&CString::new(PKG_NAME).unwrap(), &$table);
         $chain.set_hook(nftnl::Hook::In, 1);
         $chain.set_policy(nftnl::Policy::Accept);
 
@@ -29,13 +28,11 @@ macro_rules! initrules {
         $batch.add(&Rule::new(&$chain), nftnl::MsgType::Del);
 
         let mut rule = Rule::new(&$chain);
-
         rule.add_expr(&nft_expr!(ct state));
         rule.add_expr(&nft_expr!(bitwise mask 4u32, xor 0u32));
         rule.add_expr(&nft_expr!(cmp != 0u32));
         rule.add_expr(&nft_expr!(counter));
         rule.add_expr(&nft_expr!(verdict accept));
-
         $batch.add(&rule, nftnl::MsgType::Add);
     };
 }
@@ -86,25 +83,42 @@ fn fwinit(t: FwTableType) -> (Batch, Table) {
 }
 
 pub fn fwblock(
-    ips_add: &Vec<IpData>,
+    ips_add_all: &Vec<IpData>,
     ret: &mut Vec<String>,
     fwlen: &mut usize,
 ) -> std::result::Result<(), Error> {
     let ((mut batch4, table4), (mut batch6, table6)) = fwglobalinit();
 
+    let mut chain4 = Chain::new(&CString::new(PKG_NAME).unwrap(), &table4);
+    let mut chain6 = Chain::new(&CString::new(PKG_NAME).unwrap(), &table6);
+
     initrules!(batch4, table4, chain4);
     initrules!(batch6, table6, chain6);
 
+    let mut factor = 1;
+    if ips_add_all.len() > 100 {
+        factor = (ips_add_all.len() / 10) as usize
+    }
+
+    let ips_add_tmp: Vec<IpData> = ips_add_all.clone().iter().map(|x| x.clone()).collect();
+    let mut ips_add_iter = ips_add_tmp.chunks(factor);
+    let mut ips_add: Vec<&[IpData]> = vec![];
+    while let Some(x) = ips_add_iter.next() {
+        ips_add.push(x);
+    }
+
     // build and add rules
-    for ipdata in ips_add.clone() {
-        match ipdata.t {
-            4 => {
-                createrules!(ipdata, chain4, batch4, Ipv4Addr, ipv4);
+    for ipdata_group in ips_add.clone() {
+        for ipdata in ipdata_group {
+            match ipdata.t {
+                4 => {
+                    createrules!(ipdata, chain4, batch4, Ipv4Addr, ipv4);
+                }
+                6 => {
+                    createrules!(ipdata, chain6, batch6, Ipv6Addr, ipv6);
+                }
+                _ => {}
             }
-            6 => {
-                createrules!(ipdata, chain6, batch6, Ipv6Addr, ipv6);
-            }
-            _ => {}
         }
     }
 
@@ -118,10 +132,13 @@ pub fn fwblock(
             }
         };
     }
-    if fwlen != &mut ips_add.len() {
-        ret.push(format!("{length} ip in firewall", length = ips_add.len()));
+    if fwlen != &mut ips_add_all.len() {
+        ret.push(format!(
+            "{length} ip in firewall",
+            length = ips_add_all.len()
+        ));
     }
-    *fwlen = ips_add.len();
+    *fwlen = ips_add_all.len();
     Ok(())
 }