Add some initial support for tests with the full-rust engine

6 files changed, 70 insertions, 173 deletions

diff --git a/src/batch.rs b/src/batch.rs
index 714dc55..a9529a3 100644
--- a/src/batch.rs
+++ b/src/batch.rs
@@ -1,19 +1,16 @@
-use crate::nlmsg::{NfNetlinkObject, NfNetlinkWriter};
-use crate::sys::{self};
-use crate::{MsgType, ProtoFamily};
 use libc;
-use std::ffi::c_void;
-use std::os::raw::c_char;
-use std::ptr;
+
 use thiserror::Error;
 
+use crate::nlmsg::{NfNetlinkObject, NfNetlinkWriter};
+use crate::{MsgType, ProtoFamily};
+
 /// Error while communicating with netlink.
 #[derive(Error, Debug)]
 #[error("Error while communicating with netlink")]
 pub struct NetlinkError(());
 
-/// A batch of netfilter messages to be performed in one atomic operation. Corresponds to
-/// `nftnl_batch` in libnftnl.
+/// A batch of netfilter messages to be performed in one atomic operation.
 pub struct Batch {
     buf: Box<Vec<u8>>,
     // the 'static lifetime here is a cheat, as the writer can only be used as long
@@ -40,6 +37,7 @@ impl Batch {
             0,
             Some(libc::NFNL_SUBSYS_NFTABLES as u16),
         );
+        writer.finalize_writing_object();
         Batch {
             buf,
             writer,
@@ -71,7 +69,7 @@ impl Batch {
     /// Return None if there is no object in the batch (this could block forever).
     ///
     /// [`FinalizedBatch`]: struct.FinalizedBatch.html
-    pub fn finalize(mut self) -> FinalizedBatch {
+    pub fn finalize(mut self) -> Vec<u8> {
         self.writer.write_header(
             libc::NFNL_MSG_BATCH_END as u16,
             ProtoFamily::Unspec,
@@ -79,104 +77,10 @@ impl Batch {
             self.seq,
             Some(libc::NFNL_SUBSYS_NFTABLES as u16),
         );
-        FinalizedBatch { batch: self }
-    }
-
-    /*
-        fn current(&self) -> *mut c_void {
-            unsafe { sys::nftnl_batch_buffer(self.batch) }
-        }
-
-        fn next(&mut self) {
-            if unsafe { sys::nftnl_batch_update(self.batch) } < 0 {
-                // See try_alloc definition.
-                std::process::abort();
-            }
-            self.seq += 1;
-        }
-
-        fn write_begin_msg(&mut self) {
-            unsafe { sys::nftnl_batch_begin(self.current() as *mut c_char, self.seq) };
-            self.next();
-        }
-
-        fn write_end_msg(&mut self) {
-            unsafe { sys::nftnl_batch_end(self.current() as *mut c_char, self.seq) };
-            self.next();
-        }
-
-        #[cfg(feature = "unsafe-raw-handles")]
-        /// Returns the raw handle.
-        pub fn as_ptr(&self) -> *const sys::nftnl_batch {
-            self.batch as *const sys::nftnl_batch
-        }
-
-        #[cfg(feature = "unsafe-raw-handles")]
-        /// Returns a mutable version of the raw handle.
-        pub fn as_mut_ptr(&mut self) -> *mut sys::nftnl_batch {
-            self.batch
-        }
-    */
-}
-
-/// A wrapper over [`Batch`], guaranteed to start with a proper batch begin and end with a proper
-/// batch end message. Created from [`Batch::finalize`].
-///
-/// Can be turned into an iterator of the byte buffers to send to netlink to execute this batch.
-///
-/// [`Batch`]: struct.Batch.html
-/// [`Batch::finalize`]: struct.Batch.html#method.finalize
-pub struct FinalizedBatch {
-    batch: Batch,
-}
-
-/*
-impl FinalizedBatch {
-    /// Returns the iterator over byte buffers to send to netlink.
-    pub fn iter(&mut self) -> Iter<'_> {
-        let num_pages = unsafe { sys::nftnl_batch_iovec_len(self.batch.batch) as usize };
-        let mut iovecs = vec![
-            libc::iovec {
-                iov_base: ptr::null_mut(),
-                iov_len: 0,
-            };
-            num_pages
-        ];
-        let iovecs_ptr = iovecs.as_mut_ptr();
-        unsafe {
-            sys::nftnl_batch_iovec(self.batch.batch, iovecs_ptr, num_pages as u32);
-        }
-        Iter {
-            iovecs: iovecs.into_iter(),
-            _marker: ::std::marker::PhantomData,
-        }
-    }
-}
-
-impl<'a> IntoIterator for &'a mut FinalizedBatch {
-    type Item = &'a [u8];
-    type IntoIter = Iter<'a>;
-
-    fn into_iter(self) -> Iter<'a> {
-        self.iter()
-    }
-}
-
-pub struct Iter<'a> {
-    iovecs: ::std::vec::IntoIter<libc::iovec>,
-    _marker: ::std::marker::PhantomData<&'a ()>,
-}
-
-impl<'a> Iterator for Iter<'a> {
-    type Item = &'a [u8];
-
-    fn next(&mut self) -> Option<&'a [u8]> {
-        self.iovecs.next().map(|iovec| unsafe {
-            ::std::slice::from_raw_parts(iovec.iov_base as *const u8, iovec.iov_len)
-        })
+        self.writer.finalize_writing_object();
+        *self.buf
     }
 }
-*/
 
 /// Selected batch page is 256 Kbytes long to load ruleset of half a million rules without hitting
 /// -EMSGSIZE due to large iovec.
diff --git a/src/lib.rs b/src/lib.rs
index 60643fe..de24c09 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -92,9 +92,7 @@ macro_rules! try_alloc {
 }
 
 mod batch;
-//#[cfg(feature = "query")]
-//pub use batch::{batch_is_supported, default_batch_page_size};
-//pub use batch::{Batch, FinalizedBatch, NetlinkError};
+pub use batch::{default_batch_page_size, Batch};
 
 //pub mod expr;
 
diff --git a/src/nlmsg.rs b/src/nlmsg.rs
index 70a4f01..97b3b02 100644
--- a/src/nlmsg.rs
+++ b/src/nlmsg.rs
@@ -13,24 +13,6 @@ use crate::{
     MsgType, ProtoFamily,
 };
 
-/*
-/// Trait for all types in this crate that can serialize to a Netlink message.
-pub trait NlMsg {
-    /// Serializes the Netlink message to the buffer at `buf`.
-    fn write(&self, buf: &mut Vec<u8>, msg_type: MsgType, seq: u32);
-}
-
-impl<T, R> NlMsg for T
-where
-    T: Deref<Target = R>,
-    R: NlMsg,
-{
-    fn write(&self, buf: &mut Vec<u8>, msg_type: MsgType, seq: u32) {
-        self.deref().write(buf, msg_type, seq);
-    }
-}
-*/
-
 pub struct NfNetlinkWriter<'a> {
     buf: &'a mut Vec<u8>,
     headers: HeaderStack<'a>,
@@ -99,6 +81,10 @@ impl<'a> NfNetlinkWriter<'a> {
             None
         }
     }
+
+    pub fn finalize_writing_object(&mut self) {
+        self.headers.pop_level();
+    }
 }
 
 struct HeaderStack<'a> {
@@ -126,6 +112,10 @@ impl<'a> HeaderStack<'a> {
     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();
+    }
 }
 
 pub trait NfNetlinkObject: Sized {
@@ -133,5 +123,5 @@ pub trait NfNetlinkObject: Sized {
 
     fn decode_attribute(attr_type: u16, buf: &[u8]) -> Result<Attribute, DecodeError>;
 
-    fn deserialize(buf: &[u8]) -> Result<(&[u8], Self), DecodeError>;
+    fn deserialize(buf: &[u8]) -> Result<(Self, &[u8]), DecodeError>;
 }
diff --git a/src/parser.rs b/src/parser.rs
index ddcfbf4..a8df855 100644
--- a/src/parser.rs
+++ b/src/parser.rs
@@ -314,9 +314,11 @@ impl<'a> NfNetlinkAttributeReader<'a> {
         })
     }
 
-    pub fn decode<T: NfNetlinkObject>(
-        mut self,
-    ) -> Result<(&'a [u8], NfNetlinkAttributes), DecodeError> {
+    pub fn get_raw_data(&self) -> &'a [u8] {
+        &self.buf[self.pos..]
+    }
+
+    pub fn decode<T: NfNetlinkObject>(mut self) -> Result<NfNetlinkAttributes, DecodeError> {
         while self.remaining_size > pad_netlink_object::<nlattr>() {
             let nlattr =
                 unsafe { *transmute::<*const u8, *const nlattr>(self.buf[self.pos..].as_ptr()) };
@@ -337,29 +339,25 @@ impl<'a> NfNetlinkAttributeReader<'a> {
             self.remaining_size -= pad_netlink_object_with_variable_size(nlattr.nla_len as usize);
         }
 
-        Ok((&self.buf[self.pos..], self.attrs))
+        Ok(self.attrs)
     }
 }
 
-pub fn expect_msgtype_in_nlmsg<'a>(
+pub fn parse_object<'a>(
+    hdr: nlmsghdr,
+    msg: NlMsg<'a>,
     buf: &'a [u8],
-    nlmsg_type: u8,
-) -> Result<(nlmsghdr, Nfgenmsg, &'a [u8], NfNetlinkAttributeReader<'a>), DecodeError> {
-    let (hdr, msg) = parse_nlmsg(buf)?;
-
-    if get_operation_from_nlmsghdr_type(hdr.nlmsg_type) != nlmsg_type {
-        return Err(DecodeError::UnexpectedType(hdr.nlmsg_type));
-    }
-
+) -> Result<(Nfgenmsg, NfNetlinkAttributeReader<'a>, &'a [u8]), DecodeError> {
     let remaining_size = hdr.nlmsg_len as usize
         - pad_netlink_object_with_variable_size(size_of::<nlmsghdr>() + size_of::<Nfgenmsg>());
 
+    let remaining_data = &buf[pad_netlink_object_with_variable_size(hdr.nlmsg_len as usize)..];
+
     match msg {
         NlMsg::NfGenMsg(nfgenmsg, content) => Ok((
-            hdr,
             nfgenmsg,
-            content,
             NfNetlinkAttributeReader::new(content, remaining_size)?,
+            remaining_data,
         )),
         _ => Err(DecodeError::UnexpectedType(hdr.nlmsg_type)),
     }
diff --git a/src/query.rs b/src/query.rs
index 80fdc75..5065436 100644
--- a/src/query.rs
+++ b/src/query.rs
@@ -43,12 +43,15 @@ mod inner {
 
     use nix::{
         errno::Errno,
-        sys::socket::{self, AddressFamily, MsgFlags, SockFlag, SockProtocol, SockType},
+        sys::socket::{
+            self, AddressFamily, MsgFlags, NetlinkAddr, SockAddr, SockFlag, SockProtocol, SockType,
+        },
     };
 
-    //use crate::FinalizedBatch;
-
-    use crate::nlmsg::{parse_nlmsg, DecodeError, NlMsg};
+    use crate::{
+        batch::Batch,
+        parser::{parse_nlmsg, DecodeError, NlMsg},
+    };
 
     use super::*;
 
@@ -90,10 +93,8 @@ mod inner {
 
     fn recv_and_process<'a, T>(
         sock: RawFd,
-        cb: &dyn Fn(&nlmsghdr, &Nfgenmsg, &[u8], &mut T) -> Result<(), Error>,
+        cb: Option<&dyn Fn(&nlmsghdr, &Nfgenmsg, &[u8], &mut T) -> Result<(), Error>>,
         working_data: &'a mut T,
-        seq: u32,
-        portid: u32,
     ) -> Result<(), Error> {
         let mut msg_buffer = vec![0; nft_nlmsg_maxsize() as usize];
 
@@ -105,7 +106,7 @@ mod inner {
             }
             let mut buf = &msg_buffer.as_slice()[0..nb_recv];
             loop {
-                let (nlmsghdr, msg) = unsafe { parse_nlmsg(&buf, seq, portid) }?;
+                let (nlmsghdr, msg) = unsafe { parse_nlmsg(&buf) }?;
                 match msg {
                     NlMsg::Done => {
                         return Ok(());
@@ -116,7 +117,11 @@ mod inner {
                         }
                     }
                     NlMsg::Noop => {}
-                    NlMsg::NfGenMsg(genmsg, data) => cb(&nlmsghdr, &genmsg, &data, working_data)?,
+                    NlMsg::NfGenMsg(genmsg, data) => {
+                        if let Some(cb) = cb {
+                            cb(&nlmsghdr, &genmsg, &data, working_data);
+                        }
+                    }
                 }
 
                 // netlink messages are 4bytes aligned
@@ -166,18 +171,16 @@ mod inner {
         .map_err(Error::NetlinkOpenError)?;
 
         let seq = 0;
-        let portid = 0;
 
         let chains_buf = get_list_of_objects(data_type, seq, req_hdr_customize)?;
         socket::send(sock, &chains_buf, MsgFlags::empty()).map_err(Error::NetlinkSendError)?;
 
         Ok(socket_close_wrapper(sock, move |sock| {
-            recv_and_process(sock, cb, working_data, seq, portid)
+            recv_and_process(sock, Some(cb), working_data)
         })?)
     }
 
-    /*
-    pub fn send_batch(batch: &mut FinalizedBatch) -> Result<(), Error> {
+    pub fn send_batch(batch: Batch) -> Result<(), Error> {
         let sock = socket::socket(
             AddressFamily::Netlink,
             SockType::Raw,
@@ -189,28 +192,26 @@ mod inner {
         let seq = 0;
         let portid = 0;
 
+        let addr = SockAddr::Netlink(NetlinkAddr::new(portid, 0));
         // while this bind() is not strictly necessary, strace have trouble decoding the messages
         // if we don't
-        let addr = SockAddr::Netlink(NetlinkAddr::new(portid, 0));
         socket::bind(sock, &addr).expect("bind");
         //match socket::getsockname(sock).map_err(|_| Error::RetrievingSocketInfoFailed)? {
         //    SockAddr::Netlink(addr) => addr.0.nl_pid,
         //    _ => return Err(Error::NotNetlinkSocket),
         //};
 
-        for data in batch {
-            if socket::send(sock, data, MsgFlags::empty()).map_err(Error::NetlinkSendError)?
-                < data.len()
-            {
-                return Err(Error::TruncatedSend);
-            }
+        let to_send = batch.finalize();
+        let mut sent = 0;
+        while sent != to_send.len() {
+            sent += socket::send(sock, &to_send[sent..], MsgFlags::empty())
+                .map_err(Error::NetlinkSendError)?;
         }
 
         Ok(socket_close_wrapper(sock, move |sock| {
-            recv_and_process(sock, &|_, _, _, _| Ok(()), &mut (), seq, portid)
+            recv_and_process(sock, None, &mut ())
         })?)
     }
-    */
 }
 
 #[cfg(feature = "query")]
diff --git a/src/table.rs b/src/table.rs
index 6b34291..42c4f86 100644
--- a/src/table.rs
+++ b/src/table.rs
@@ -3,13 +3,11 @@ use std::fmt::Debug;
 
 use crate::nlmsg::{NfNetlinkObject, NfNetlinkWriter};
 use crate::parser::{
-    expect_msgtype_in_nlmsg, parse_nlmsg, Attribute, DecodeError, NfNetlinkAttributeReader,
-    NfNetlinkAttributes, NlMsg, SerializeNfNetlink,
+    get_operation_from_nlmsghdr_type, parse_nlmsg, parse_object, Attribute, DecodeError,
+    NfNetlinkAttributeReader, NfNetlinkAttributes, NlMsg, SerializeNfNetlink,
 };
 use crate::sys::{self, NFTA_OBJ_TABLE, NFTA_TABLE_FLAGS, NFTA_TABLE_NAME};
 use crate::{impl_attr_getters_and_setters, MsgType, ProtoFamily};
-#[cfg(feature = "query")]
-use crate::{parser::Nfgenmsg, query::ParseError};
 use libc;
 
 /// Abstraction of `nftnl_table`, the top level container in netfilter. A table has a protocol
@@ -69,6 +67,7 @@ impl NfNetlinkObject for Table {
         } as u16;
         writer.write_header(raw_msg_type, self.family, libc::NLM_F_ACK as u16, seq, None);
         self.inner.serialize(writer);
+        writer.finalize_writing_object();
     }
 
     fn decode_attribute(attr_type: u16, buf: &[u8]) -> Result<Attribute, DecodeError> {
@@ -84,18 +83,25 @@ impl NfNetlinkObject for Table {
         }
     }
 
-    fn deserialize(buf: &[u8]) -> Result<(&[u8], Self), DecodeError> {
-        let (hdr, nfgenmsg, content, mut attrs) =
-            expect_msgtype_in_nlmsg(buf, libc::NFT_MSG_NEWTABLE as u8)?;
+    fn deserialize(buf: &[u8]) -> Result<(Self, &[u8]), DecodeError> {
+        let (hdr, msg) = parse_nlmsg(buf)?;
+
+        let op = get_operation_from_nlmsghdr_type(hdr.nlmsg_type) as i32;
+
+        if op != libc::NFT_MSG_NEWTABLE && op != libc::NFT_MSG_DELTABLE {
+            return Err(DecodeError::UnexpectedType(hdr.nlmsg_type));
+        }
+
+        let (nfgenmsg, attrs, remaining_data) = parse_object(hdr, msg, buf)?;
 
-        let (remaining_buf, inner) = attrs.decode::<Table>()?;
+        let inner = attrs.decode::<Table>()?;
 
         Ok((
-            remaining_buf,
             Table {
                 inner,
                 family: ProtoFamily::try_from(nfgenmsg.family as i32)?,
             },
+            remaining_data,
         ))
     }
 }