1 files changed, 9 insertions, 105 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.