1 files changed, 24 insertions, 59 deletions

diff --git a/src/nlmsg.rs b/src/nlmsg.rs
index 435fed3..a1bb200 100644
--- a/src/nlmsg.rs
+++ b/src/nlmsg.rs
@@ -8,26 +8,22 @@ use std::{
 use crate::{
     parser::{
         pad_netlink_object, pad_netlink_object_with_variable_size, AttributeType, DecodeError,
-        Nfgenmsg,
     },
     sys::{
-        nlattr, nlmsghdr, NFNETLINK_V0, NFNL_MSG_BATCH_BEGIN, NFNL_MSG_BATCH_END,
-        NFNL_SUBSYS_NFTABLES, NLA_TYPE_MASK,
+        nfgenmsg, nlattr, nlmsghdr, NFNETLINK_V0, NFNL_MSG_BATCH_BEGIN, NFNL_MSG_BATCH_END,
+        NFNL_SUBSYS_NFTABLES,
     },
     MsgType, ProtoFamily,
 };
 
 pub struct NfNetlinkWriter<'a> {
     buf: &'a mut Vec<u8>,
-    headers: HeaderStack<'a>,
+    headers: Option<(usize, usize)>,
 }
 
 impl<'a> NfNetlinkWriter<'a> {
     pub fn new(buf: &'a mut Vec<u8>) -> NfNetlinkWriter<'a> {
-        NfNetlinkWriter {
-            buf,
-            headers: HeaderStack::new(),
-        }
+        NfNetlinkWriter { buf, headers: None }
     }
 
     pub fn add_data_zeroed<'b>(&'b mut self, size: usize) -> &'b mut [u8] {
@@ -35,7 +31,12 @@ impl<'a> NfNetlinkWriter<'a> {
         let start = self.buf.len();
         self.buf.resize(start + padded_size, 0);
 
-        self.headers.add_size(padded_size as u32);
+        if let Some((msghdr_idx, _nfgenmsg_idx)) = self.headers {
+            let mut hdr: &mut nlmsghdr = unsafe {
+                std::mem::transmute(self.buf[msghdr_idx..].as_mut_ptr() as *mut nlmsghdr)
+            };
+            hdr.nlmsg_len += padded_size as u32;
+        }
 
         &mut self.buf[start..start + size]
     }
@@ -53,14 +54,16 @@ impl<'a> NfNetlinkWriter<'a> {
         seq: u32,
         ressource_id: Option<u16>,
     ) {
+        if self.headers.is_some() {
+            error!("Calling write_header while still holding headers open!?");
+        }
+
         let nlmsghdr_len = pad_netlink_object::<nlmsghdr>();
-        let nfgenmsg_len = pad_netlink_object::<Nfgenmsg>();
-        let nlmsghdr_buf = self.add_data_zeroed(nlmsghdr_len);
+        let nfgenmsg_len = pad_netlink_object::<nfgenmsg>();
 
+        let nlmsghdr_buf = self.add_data_zeroed(nlmsghdr_len);
         let mut hdr: &mut nlmsghdr =
             unsafe { std::mem::transmute(nlmsghdr_buf.as_mut_ptr() as *mut nlmsghdr) };
-        //let mut hdr = &mut unsafe { *(nlmsghdr_buf.as_mut_ptr() as *mut nlmsghdr) };
-
         hdr.nlmsg_len = (nlmsghdr_len + nfgenmsg_len) as u32;
         hdr.nlmsg_type = msg_type;
         // batch messages are not specific to the nftables subsystem
@@ -71,58 +74,20 @@ impl<'a> NfNetlinkWriter<'a> {
         hdr.nlmsg_seq = seq;
 
         let nfgenmsg_buf = self.add_data_zeroed(nfgenmsg_len);
-
-        let mut nfgenmsg: &mut Nfgenmsg =
-            unsafe { std::mem::transmute(nfgenmsg_buf.as_mut_ptr() as *mut Nfgenmsg) };
-        nfgenmsg.family = family as u8;
+        let mut nfgenmsg: &mut nfgenmsg =
+            unsafe { std::mem::transmute(nfgenmsg_buf.as_mut_ptr() as *mut nfgenmsg) };
+        nfgenmsg.nfgen_family = family as u8;
         nfgenmsg.version = NFNETLINK_V0 as u8;
         nfgenmsg.res_id = ressource_id.unwrap_or(0);
 
-        self.headers.add_level(hdr, Some(nfgenmsg));
-    }
-
-    pub fn get_current_header(&mut self) -> Option<&mut nlmsghdr> {
-        let stack_size = self.headers.stack.len();
-        if stack_size > 0 {
-            Some(unsafe { std::mem::transmute(self.headers.stack[stack_size - 1].0) })
-        } else {
-            None
-        }
+        self.headers = Some((
+            self.buf.len() - (nlmsghdr_len + nfgenmsg_len),
+            self.buf.len() - nfgenmsg_len,
+        ));
     }
 
     pub fn finalize_writing_object(&mut self) {
-        self.headers.pop_level();
-    }
-}
-
-struct HeaderStack<'a> {
-    stack: Vec<(*mut nlmsghdr, Option<*mut Nfgenmsg>)>,
-    lifetime: PhantomData<&'a ()>,
-}
-
-impl<'a> HeaderStack<'a> {
-    fn new() -> HeaderStack<'a> {
-        HeaderStack {
-            stack: Vec::new(),
-            lifetime: PhantomData,
-        }
-    }
-
-    /// resize all the stacked netlink containers to hold additional_size new bytes
-    fn add_size(&mut self, additional_size: u32) {
-        for (hdr, _) in &mut self.stack {
-            unsafe {
-                (**hdr).nlmsg_len = (**hdr).nlmsg_len + additional_size;
-            }
-        }
-    }
-
-    fn add_level(&mut self, hdr: *mut nlmsghdr, nfgenmsg: Option<*mut Nfgenmsg>) {
-        self.stack.push((hdr, nfgenmsg));
-    }
-
-    fn pop_level(&mut self) {
-        self.stack.pop();
+        self.headers = None;
     }
 }