(re-)add table listing

5 files changed, 141 insertions, 169 deletions

diff --git a/src/batch.rs b/src/batch.rs
index c7fb8f3..01f0c31 100644
--- a/src/batch.rs
+++ b/src/batch.rs
@@ -87,9 +87,8 @@ impl Batch {
         *self.buf
     }
 
-    #[cfg(feature = "query")]
     pub fn send(mut self) -> Result<(), Error> {
-        use crate::query::{recv_and_process_until_seq, socket_close_wrapper};
+        use crate::query::{recv_and_process, socket_close_wrapper};
 
         let sock = socket::socket(
             AddressFamily::Netlink,
@@ -114,7 +113,7 @@ impl Batch {
         }
 
         Ok(socket_close_wrapper(sock, move |sock| {
-            recv_and_process_until_seq(sock, max_seq, None, &mut ())
+            recv_and_process(sock, Some(max_seq), None, &mut ())
         })?)
     }
 }
diff --git a/src/lib.rs b/src/lib.rs
index de24c09..6208bb5 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -98,11 +98,9 @@ pub use batch::{default_batch_page_size, Batch};
 
 pub mod table;
 pub use table::Table;
-//#[cfg(feature = "query")]
 //pub use table::{get_tables_cb, list_tables};
 //
 //mod chain;
-//#[cfg(feature = "query")]
 //pub use chain::{get_chains_cb, list_chains_for_table};
 //pub use chain::{Chain, ChainType, Hook, Policy, Priority};
 
@@ -116,7 +114,6 @@ pub mod parser;
 
 //mod rule;
 //pub use rule::Rule;
-//#[cfg(feature = "query")]
 //pub use rule::{get_rules_cb, list_rules_for_chain};
 
 //mod rule_methods;
diff --git a/src/parser.rs b/src/parser.rs
index 23b5213..a01c3cd 100644
--- a/src/parser.rs
+++ b/src/parser.rs
@@ -1,20 +1,19 @@
 use std::{
     collections::HashMap,
     fmt::Debug,
-    mem::{self, size_of, transmute},
-    str::Utf8Error,
+    mem::{size_of, transmute},
     string::FromUtf8Error,
 };
 
-use libc::{
-    nlattr, nlmsgerr, nlmsghdr, NFNETLINK_V0, NFNL_MSG_BATCH_BEGIN, NFNL_MSG_BATCH_END,
-    NFNL_SUBSYS_NFTABLES, NLA_TYPE_MASK, NLMSG_DONE, NLMSG_ERROR, NLMSG_MIN_TYPE, NLMSG_NOOP,
-    NLM_F_DUMP_INTR,
-};
 use thiserror::Error;
 
 use crate::{
     nlmsg::{NfNetlinkObject, NfNetlinkWriter},
+    sys::{
+        nlattr, nlmsgerr, nlmsghdr, NFNETLINK_V0, NFNL_MSG_BATCH_BEGIN, NFNL_MSG_BATCH_END,
+        NFNL_SUBSYS_NFTABLES, NLA_TYPE_MASK, NLMSG_ALIGNTO, NLMSG_DONE, NLMSG_ERROR,
+        NLMSG_MIN_TYPE, NLMSG_NOOP, NLM_F_DUMP_INTR,
+    },
     InvalidProtocolFamily, ProtoFamily,
 };
 
@@ -74,13 +73,12 @@ pub fn nft_nlmsg_maxsize() -> u32 {
 #[inline]
 pub fn pad_netlink_object_with_variable_size(size: usize) -> usize {
     // align on a 4 bytes boundary
-    (size + 3) & (!3)
+    (size + (NLMSG_ALIGNTO as usize - 1)) & !(NLMSG_ALIGNTO as usize - 1)
 }
 
 #[inline]
 pub fn pad_netlink_object<T>() -> usize {
     let size = size_of::<T>();
-    // align on a 4 bytes boundary
     pad_netlink_object_with_variable_size(size)
 }
 
@@ -139,12 +137,10 @@ pub fn parse_nlmsg<'a>(buf: &'a [u8]) -> Result<(nlmsghdr, NlMsg<'a>), DecodeErr
     let size_of_hdr = pad_netlink_object::<nlmsghdr>();
 
     if hdr.nlmsg_type < NLMSG_MIN_TYPE as u16 {
-        match hdr.nlmsg_type as libc::c_int {
+        match hdr.nlmsg_type as u32 {
             x if x == NLMSG_NOOP => return Ok((hdr, NlMsg::Noop)),
             x if x == NLMSG_ERROR => {
-                if hdr.nlmsg_len as usize > buf.len()
-                    || (hdr.nlmsg_len as usize) < size_of_hdr + size_of::<nlmsgerr>()
-                {
+                if (hdr.nlmsg_len as usize) < size_of_hdr + size_of::<nlmsgerr>() {
                     return Err(DecodeError::NlMsgTooSmall);
                 }
                 let mut err = unsafe {
@@ -196,8 +192,8 @@ pub trait NfNetlinkAttribute: Debug + Sized {
         size_of::<Self>()
     }
 
-    unsafe fn write_payload(&self, addr: *mut u8);
     // example body: std::ptr::copy_nonoverlapping(self as *const Self as *const u8, addr, self.get_size());
+    unsafe fn write_payload(&self, addr: *mut u8);
 }
 
 /// Write the attribute, preceded by a `libc::nlattr`
@@ -438,7 +434,7 @@ impl_attribute!(
 
 #[macro_export]
 macro_rules! impl_attr_getters_and_setters {
-    ($struct:ident, [$(($getter_name:ident, $setter_name:ident, $attr_name:expr, $internal_name:ident, $type:ty)),+]) => {
+    ($struct:ident, [$(($getter_name:ident, $setter_name:ident, $in_place_edit_name:ident, $attr_name:expr, $internal_name:ident, $type:ty)),+]) => {
         impl $struct {
             $(
                 #[allow(dead_code)]
@@ -447,8 +443,14 @@ macro_rules! impl_attr_getters_and_setters {
                 }
 
                 #[allow(dead_code)]
-                pub fn $setter_name(&mut self, val: $type) {
-                    self.inner.set_attr($attr_name as $crate::parser::NetlinkType, $crate::parser::Attribute::$internal_name(val));
+                pub fn $setter_name(&mut self, val: impl Into<$type>) {
+                    self.inner.set_attr($attr_name as $crate::parser::NetlinkType, $crate::parser::Attribute::$internal_name(val.into()));
+                }
+
+                #[allow(dead_code)]
+                pub fn $in_place_edit_name(mut self, val: impl Into<$type>) -> Self {
+                    self.inner.set_attr($attr_name as $crate::parser::NetlinkType, $crate::parser::Attribute::$internal_name(val.into()));
+                    self
                 }
             )+
         }
diff --git a/src/query.rs b/src/query.rs
index d2409f2..f84586a 100644
--- a/src/query.rs
+++ b/src/query.rs
@@ -1,55 +1,18 @@
-use std::mem::size_of;
+use std::os::unix::prelude::RawFd;
 
 use crate::{
-    nlmsg::NfNetlinkWriter,
-    parser::{nft_nlmsg_maxsize, Nfgenmsg},
-    sys, ProtoFamily,
-};
-use libc::{
-    nlmsgerr, nlmsghdr, NFNETLINK_V0, NFNL_SUBSYS_NFTABLES, NLMSG_DONE, NLMSG_ERROR,
-    NLMSG_MIN_TYPE, NLMSG_NOOP, NLM_F_DUMP_INTR,
+    nlmsg::{NfNetlinkObject, NfNetlinkWriter},
+    parser::{nft_nlmsg_maxsize, pad_netlink_object_with_variable_size},
+    sys::{nlmsgerr, NLM_F_DUMP, NLM_F_MULTI},
+    ProtoFamily,
 };
 
-/// Returns a buffer containing a netlink message which requests a list of all the netfilter
-/// matching objects (e.g. tables, chains, rules, ...).
-/// Supply the type of objects to retrieve (e.g. libc::NFT_MSG_GETTABLE), and optionally a callback
-/// to execute on the header, to set parameters for example.
-/// To pass arbitrary data inside that callback, please use a closure.
-pub fn get_list_of_objects<Error>(
-    msg_type: u16,
-    seq: u32,
-    setup_cb: Option<&dyn Fn(&mut libc::nlmsghdr) -> Result<(), Error>>,
-) -> Result<Vec<u8>, Error> {
-    let mut buffer = vec![0; nft_nlmsg_maxsize() as usize];
-    let mut writer = NfNetlinkWriter::new(&mut buffer);
-    writer.write_header(
-        msg_type,
-        ProtoFamily::Unspec,
-        (libc::NLM_F_ROOT | libc::NLM_F_MATCH) as u16,
-        seq,
-        None,
-    );
-    if let Some(cb) = setup_cb {
-        cb(writer
-            .get_current_header()
-            .expect("Fatal error: mising header"))?;
-    }
-    Ok(buffer)
-}
-
-use std::os::unix::prelude::RawFd;
-
 use nix::{
     errno::Errno,
-    sys::socket::{
-        self, AddressFamily, MsgFlags, NetlinkAddr, SockAddr, SockFlag, SockProtocol, SockType,
-    },
+    sys::socket::{self, AddressFamily, MsgFlags, SockFlag, SockProtocol, SockType},
 };
 
-use crate::{
-    batch::Batch,
-    parser::{parse_nlmsg, DecodeError, NlMsg},
-};
+use crate::parser::{parse_nlmsg, DecodeError, NlMsg};
 
 #[derive(thiserror::Error, Debug)]
 pub enum Error {
@@ -83,27 +46,45 @@ pub enum Error {
     #[error("Only a part of the message was sent")]
     TruncatedSend,
 
+    #[error("Got a message without the NLM_F_MULTI flag, but a maximum sequence number was not specified")]
+    UndecidableMessageTermination,
+
     #[error("Couldn't close the socket")]
     CloseFailed(#[source] Errno),
 }
 
-pub(crate) fn recv_and_process_until_seq<'a, T>(
+pub(crate) fn recv_and_process<'a, T>(
     sock: RawFd,
-    max_seq: u32,
-    cb: Option<&dyn Fn(&nlmsghdr, &Nfgenmsg, &[u8], &mut T) -> Result<(), Error>>,
+    max_seq: Option<u32>,
+    cb: Option<&dyn Fn(&[u8], &mut T) -> Result<(), Error>>,
     working_data: &'a mut T,
 ) -> Result<(), Error> {
-    let mut msg_buffer = vec![0; nft_nlmsg_maxsize() as usize];
+    let mut msg_buffer = vec![0; 2 * nft_nlmsg_maxsize() as usize];
+    let mut buf_start = 0;
+    let mut end_pos = 0;
 
     loop {
-        let nb_recv = socket::recv(sock, &mut msg_buffer, MsgFlags::empty())
+        let nb_recv = socket::recv(sock, &mut msg_buffer[end_pos..], MsgFlags::empty())
             .map_err(Error::NetlinkRecvError)?;
         if nb_recv <= 0 {
             return Ok(());
         }
-        let mut buf = &msg_buffer.as_slice()[0..nb_recv];
+        end_pos += nb_recv;
         loop {
+            let buf = &msg_buffer.as_slice()[buf_start..end_pos];
+            // exit the loop and try to receive further messages when we consumed all the buffer
+            if buf.len() == 0 {
+                break;
+            }
+
+            debug!("calling parse_nlmsg");
             let (nlmsghdr, msg) = parse_nlmsg(&buf)?;
+            debug!("Got a valid netlink message: {:?} {:?}", nlmsghdr, msg);
+            // we cannot know when a message will end if we are not receiving messages ending with an
+            // NlMsg::Done marker, and if a maximum sequence number wasn't specified either
+            if max_seq.is_none() && nlmsghdr.nlmsg_flags & NLM_F_MULTI as u16 == 0 {
+                return Err(Error::UndecidableMessageTermination);
+            }
             match msg {
                 NlMsg::Done => {
                     return Ok(());
@@ -114,55 +95,88 @@ pub(crate) fn recv_and_process_until_seq<'a, T>(
                     }
                 }
                 NlMsg::Noop => {}
-                NlMsg::NfGenMsg(genmsg, data) => {
+                NlMsg::NfGenMsg(_genmsg, _data) => {
                     if let Some(cb) = cb {
-                        cb(&nlmsghdr, &genmsg, &data, working_data)?;
+                        cb(&buf[0..nlmsghdr.nlmsg_len as usize], working_data)?;
                     }
                 }
             }
 
-            // netlink messages are 4bytes aligned
-            let aligned_length = ((nlmsghdr.nlmsg_len + 3) & !3u32) as usize;
-
             // retrieve the next message
-            buf = &buf[aligned_length..];
-
-            if nlmsghdr.nlmsg_seq >= max_seq {
-                return Ok(());
+            if let Some(max_seq) = max_seq {
+                if nlmsghdr.nlmsg_seq >= max_seq {
+                    return Ok(());
+                }
             }
 
-            // exit the loop and try to receive further messages when we consumed all the buffer
-            if buf.len() == 0 {
-                break;
+            // netlink messages are 4bytes aligned
+            let aligned_length = pad_netlink_object_with_variable_size(nlmsghdr.nlmsg_len as usize);
+            buf_start += aligned_length;
+        }
+        // Ensure that we always have nft_nlmsg_maxsize() free space available in the buffer.
+        // We achieve this by relocating the buffer content at the beginning of the buffer
+        if end_pos >= nft_nlmsg_maxsize() as usize {
+            if buf_start < end_pos {
+                unsafe {
+                    std::ptr::copy(
+                        msg_buffer[buf_start..end_pos].as_ptr(),
+                        msg_buffer.as_mut_ptr(),
+                        end_pos - buf_start,
+                    );
+                }
             }
+            end_pos = end_pos - buf_start;
+            buf_start = 0;
         }
     }
 }
 
-pub(crate) fn socket_close_wrapper(
+pub(crate) fn socket_close_wrapper<E>(
     sock: RawFd,
-    cb: impl FnOnce(RawFd) -> Result<(), Error>,
-) -> Result<(), Error> {
+    cb: impl FnOnce(RawFd) -> Result<(), E>,
+) -> Result<(), Error>
+where
+    Error: From<E>,
+{
     let ret = cb(sock);
 
     // we don't need to shutdown the socket (in fact, Linux doesn't support that operation;
     // and return EOPNOTSUPP if we try)
     nix::unistd::close(sock).map_err(Error::CloseFailed)?;
 
-    ret
+    Ok(ret?)
+}
+
+/// Returns a buffer containing a netlink message which requests a list of all the netfilter
+/// matching objects (e.g. tables, chains, rules, ...).
+/// Supply the type of objects to retrieve (e.g. libc::NFT_MSG_GETTABLE), and a search filter.
+pub fn get_list_of_objects<T>(msg_type: u16, seq: u32, filter: Option<&T>) -> Result<Vec<u8>, Error>
+where
+    T: NfNetlinkObject,
+{
+    let mut buffer = Vec::new();
+    let mut writer = NfNetlinkWriter::new(&mut buffer);
+    writer.write_header(msg_type, ProtoFamily::Unspec, NLM_F_DUMP as u16, seq, None);
+    writer.finalize_writing_object();
+    if let Some(filter) = filter {
+        filter.add_or_remove(&mut writer, crate::MsgType::Add, 0);
+    }
+    Ok(buffer)
 }
 
-/*
 /// Lists objects of a certain type (e.g. libc::NFT_MSG_GETTABLE) with the help of a helper
 /// function called by mnl::cb_run2.
 /// The callback expects a tuple of additional data (supplied as an argument to this function)
 /// and of the output vector, to which it should append the parsed object it received.
-pub fn list_objects_with_data<'a, T>(
+pub fn list_objects_with_data<'a, Object, Accumulator>(
     data_type: u16,
-    cb: &dyn Fn(&libc::nlmsghdr, &Nfgenmsg, &[u8], &mut T) -> Result<(), Error>,
-    working_data: &'a mut T,
-    req_hdr_customize: Option<&dyn Fn(&mut libc::nlmsghdr) -> Result<(), Error>>,
-) -> Result<(), Error> {
+    cb: &dyn Fn(Object, &mut Accumulator) -> Result<(), Error>,
+    filter: Option<&Object>,
+    working_data: &'a mut Accumulator,
+) -> Result<(), Error>
+where
+    Object: NfNetlinkObject,
+{
     debug!("listing objects of kind {}", data_type);
     let sock = socket::socket(
         AddressFamily::Netlink,
@@ -174,11 +188,18 @@ pub fn list_objects_with_data<'a, T>(
 
     let seq = 0;
 
-    let chains_buf = get_list_of_objects(data_type, seq, req_hdr_customize)?;
+    let chains_buf = get_list_of_objects(data_type, seq, filter)?;
     socket::send(sock, &chains_buf, MsgFlags::empty()).map_err(Error::NetlinkSendError)?;
 
-    Ok(socket_close_wrapper(sock, move |sock| {
-        recv_and_process(sock, Some(cb), working_data)
-    })?)
+    socket_close_wrapper(sock, move |sock| {
+        // the kernel should return NLM_F_MULTI objects
+        recv_and_process(
+            sock,
+            None,
+            Some(&|buf: &[u8], working_data: &mut Accumulator| {
+                cb(Object::deserialize(buf)?.0, working_data)
+            }),
+            working_data,
+        )
+    })
 }
-*/
diff --git a/src/table.rs b/src/table.rs
index 42c4f86..23495a4 100644
--- a/src/table.rs
+++ b/src/table.rs
@@ -4,11 +4,13 @@ use std::fmt::Debug;
 use crate::nlmsg::{NfNetlinkObject, NfNetlinkWriter};
 use crate::parser::{
     get_operation_from_nlmsghdr_type, parse_nlmsg, parse_object, Attribute, DecodeError,
-    NfNetlinkAttributeReader, NfNetlinkAttributes, NlMsg, SerializeNfNetlink,
+    NfNetlinkAttributeReader, NfNetlinkAttributes, Nfgenmsg, NlMsg, SerializeNfNetlink,
+};
+use crate::sys::{
+    self, NFTA_OBJ_TABLE, NFTA_TABLE_FLAGS, NFTA_TABLE_NAME, NFT_MSG_DELTABLE, NFT_MSG_GETTABLE,
+    NFT_MSG_NEWTABLE, NLM_F_ACK,
 };
-use crate::sys::{self, NFTA_OBJ_TABLE, NFTA_TABLE_FLAGS, NFTA_TABLE_NAME};
 use crate::{impl_attr_getters_and_setters, MsgType, ProtoFamily};
-use libc;
 
 /// Abstraction of `nftnl_table`, the top level container in netfilter. A table has a protocol
 /// family and contains [`Chain`]s that in turn hold the rules.
@@ -21,17 +23,11 @@ pub struct Table {
 }
 
 impl Table {
-    /// Creates a new table instance with the given name and protocol family.
-    pub fn new<T: Into<String>>(name: T, family: ProtoFamily) -> Table {
-        let mut res = Table {
+    pub fn new(family: ProtoFamily) -> Table {
+        Table {
             inner: NfNetlinkAttributes::new(),
             family,
-        };
-
-        res.set_name(name.into());
-        res.set_flags(0);
-
-        res
+        }
     }
 
     /*
@@ -62,10 +58,10 @@ impl PartialEq for Table {
 impl NfNetlinkObject for Table {
     fn add_or_remove<'a>(&self, writer: &mut NfNetlinkWriter<'a>, msg_type: MsgType, seq: u32) {
         let raw_msg_type = match msg_type {
-            MsgType::Add => libc::NFT_MSG_NEWTABLE,
-            MsgType::Del => libc::NFT_MSG_DELTABLE,
+            MsgType::Add => NFT_MSG_NEWTABLE,
+            MsgType::Del => NFT_MSG_DELTABLE,
         } as u16;
-        writer.write_header(raw_msg_type, self.family, libc::NLM_F_ACK as u16, seq, None);
+        writer.write_header(raw_msg_type, self.family, NLM_F_ACK as u16, seq, None);
         self.inner.serialize(writer);
         writer.finalize_writing_object();
     }
@@ -86,9 +82,9 @@ impl NfNetlinkObject for Table {
     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;
+        let op = get_operation_from_nlmsghdr_type(hdr.nlmsg_type) as u32;
 
-        if op != libc::NFT_MSG_NEWTABLE && op != libc::NFT_MSG_DELTABLE {
+        if op != NFT_MSG_NEWTABLE && op != NFT_MSG_DELTABLE {
             return Err(DecodeError::UnexpectedType(hdr.nlmsg_type));
         }
 
@@ -109,72 +105,29 @@ impl NfNetlinkObject for Table {
 impl_attr_getters_and_setters!(
     Table,
     [
-        (get_name, set_name, sys::NFTA_TABLE_NAME, String, String),
+        (get_name, set_name, with_name, sys::NFTA_TABLE_NAME, String, String),
         (
             get_userdata,
             set_userdata,
+            with_userdata,
             sys::NFTA_TABLE_USERDATA,
             VecU8,
             Vec<u8>
         ),
-        (get_flags, set_flags, sys::NFTA_TABLE_FLAGS, U32, u32)
+        (get_flags, set_flags, with_flags, sys::NFTA_TABLE_FLAGS, U32, u32)
     ]
 );
 
-/*
-#[cfg(feature = "query")]
-/// A callback to parse the response for messages created with `get_tables_nlmsg`.
-pub fn get_tables_cb(
-    header: &libc::nlmsghdr,
-    _genmsg: &Nfgenmsg,
-    _data: &[u8],
-    tables: &mut Vec<Table>,
-) -> Result<(), crate::query::Error> {
-    unsafe {
-        let table = sys::nftnl_table_alloc();
-        if table == std::ptr::null_mut() {
-            return Err(ParseError::Custom(Box::new(std::io::Error::new(
-                std::io::ErrorKind::Other,
-                "Table allocation failed",
-            )))
-            .into());
-        }
-        let err = sys::nftnl_table_nlmsg_parse(header, table);
-        if err < 0 {
-            sys::nftnl_table_free(table);
-            return Err(ParseError::Custom(Box::new(std::io::Error::new(
-                std::io::ErrorKind::Other,
-                "The netlink table couldn't be parsed !?",
-            )))
-            .into());
-        }
-        let family = sys::nftnl_table_get_u32(table, sys::NFTNL_TABLE_FAMILY as u16);
-        match crate::ProtoFamily::try_from(family as i32) {
-            Ok(family) => {
-                tables.push(Table::from_raw(table, family));
-                Ok(())
-            }
-            Err(crate::InvalidProtocolFamily) => {
-                sys::nftnl_table_free(table);
-                Err(ParseError::Custom(Box::new(std::io::Error::new(
-                    std::io::ErrorKind::Other,
-                    "The netlink table didn't have a valid protocol family !?",
-                )))
-                .into())
-            }
-        }
-    }
-}
-
-#[cfg(feature = "query")]
 pub fn list_tables() -> Result<Vec<Table>, crate::query::Error> {
     let mut result = Vec::new();
     crate::query::list_objects_with_data(
-        libc::NFT_MSG_GETTABLE as u16,
-        &get_tables_cb,
-        &mut result,
+        NFT_MSG_GETTABLE as u16,
+        &|table: Table, tables: &mut Vec<Table>| {
+            tables.push(table);
+            Ok(())
+        },
         None,
+        &mut result,
     )?;
     Ok(result)
 }
-*/