sample ChainMethods impl, updated example

4 files changed, 165 insertions, 57 deletions

diff --git a/Cargo.toml b/Cargo.toml
index 7df5b6c..e2a2039 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -22,6 +22,7 @@ log = "0.4"
 libc = "0.2.43"
 mnl = "0.2"
 ipnetwork = "0.16"
+serde = { version = "1.0", features = ["derive"] }
 
 [dev-dependencies]
 rustables = { path = ".", features = ["query"] }
diff --git a/examples/firewall.rs b/examples/firewall.rs
index f852618..0b5ea0c 100644
--- a/examples/firewall.rs
+++ b/examples/firewall.rs
@@ -1,10 +1,13 @@
-use rustables::{Batch, Chain, Hook, Match, MatchError, Policy, Rule, Protocol, ProtoFamily, Table, MsgType, expr::LogGroup};
+use rustables::{Batch, Chain, ChainMethods, Direction, MatchError, ProtoFamily,
+    Protocol, Rule, RuleMethods, Table, MsgType, Verdict};
 use rustables::query::{send_batch, Error as QueryError};
+use rustables::expr::{LogGroup, LogPrefix, LogPrefixError};
 use ipnetwork::IpNetwork;
 use std::ffi::{CString, NulError};
 use std::str::Utf8Error;
 use std::rc::Rc;
 
+
 #[derive(thiserror::Error, Debug)]
 pub enum Error {
     #[error("Unable to open netlink socket to netfilter")]
@@ -21,96 +24,119 @@ pub enum Error {
     BatchError(#[from] std::io::Error),
     #[error("Error applying batch")]
     QueryError(#[from] QueryError),
+    #[error("Error encoding the prefix")]
+    LogPrefixError(#[from] LogPrefixError),
 }
 
 const TABLE_NAME: &str = "main-table";
 
+
 fn main() -> Result<(), Error> {
-    let mut fw = Firewall::new()?;
+    let fw = Firewall::new()?;
     fw.start()?;
     Ok(())
 }
 
+
+/// An example firewall. See the source of its `start()` method.
 pub struct Firewall {
-    table: Rc<Table>
+    batch: Batch,
+    inbound: Rc<Chain>,
+    _outbound: Rc<Chain>,
+    _forward: Rc<Chain>,
+    _table: Rc<Table>,
 }
 
 impl Firewall {
     pub fn new() -> Result<Self, Error> {
-        let table = Rc::new(Table::new(
-            &CString::new(TABLE_NAME)?,
-            ProtoFamily::Inet
-        ));
-        Ok(Firewall { table })
-    }
-    /// Attempt to use the batch from the struct holding the table.
-    pub fn allow_port(&mut self, port: &str, protocol: &Protocol, chain: Rc<Chain>, batch: &mut Batch) -> Result<(), Error> {
-        let rule = Rule::new(chain).dport(port, protocol)?.accept().add_to_batch(batch);
-        batch.add(&rule, MsgType::Add);
-        Ok(())
-    }
-    /// Apply the current realm's batch.
-    pub fn start(&mut self) -> Result<(), Error> {
         let mut batch = Batch::new();
-        batch.add(&self.table, MsgType::Add);
+        let _table = Rc::new(
+            Table::new(&CString::new(TABLE_NAME)?, ProtoFamily::Inet)
+        );
+        batch.add(&_table, MsgType::Add);
 
-        let local_net = IpNetwork::new([192, 168, 1, 0].into(), 24).unwrap();
-        let mut inbound = Chain::new(&CString::new("in")?, Rc::clone(&self.table));
-        inbound.set_hook(Hook::In, 0);
-        inbound.set_policy(Policy::Drop);
-        let inbound = Rc::new(inbound);
-        batch.add(&inbound, MsgType::Add);
-        let mut outbound = Chain::new(&CString::new("out")?, Rc::clone(&self.table));
-        outbound.set_hook(Hook::Out, 0);
-        outbound.set_policy(Policy::Accept);
-        batch.add(&outbound, MsgType::Add);
-        let mut forward = Chain::new(&CString::new("forward")?, Rc::clone(&self.table));
-        forward.set_hook(Hook::Forward, 0);
-        forward.set_policy(Policy::Accept);
-        batch.add(&forward, MsgType::Add);
-        Rule::new(Rc::clone(&inbound))
+        // Create base chains. Base chains are hooked into a Direction/Hook.
+        let inbound = Rc::new(
+            Chain::from_direction(&Direction::Inbound, Rc::clone(&_table))?
+                  .verdict(&Verdict::Drop)
+                  .add_to_batch(&mut batch)
+        );
+        let _outbound = Rc::new(
+            Chain::from_direction(&Direction::Outbound, Rc::clone(&_table))?
+                  .verdict(&Verdict::Accept)
+                  .add_to_batch(&mut batch)
+        );
+        let _forward = Rc::new(
+            Chain::from_direction(&Direction::Forward, Rc::clone(&_table))?
+                  .verdict(&Verdict::Accept)
+                  .add_to_batch(&mut batch)
+        );
+
+        Ok(Firewall {
+            _table,
+            batch,
+            inbound,
+            _outbound,
+            _forward
+        })
+    }
+    /// Allow some common-sense exceptions to inbound drop, and accept outbound and forward.
+    pub fn start(mut self) -> Result<(), Error> {
+        // Allow all established connections to get in.
+        Rule::new(Rc::clone(&self.inbound))
              .established()
              .accept()
-             .add_to_batch(&mut batch);
-        Rule::new(Rc::clone(&inbound))
+             .add_to_batch(&mut self.batch);
+        // Allow all traffic on the loopback interface.
+        Rule::new(Rc::clone(&self.inbound))
              .iface("lo")?
              .accept()
-             .add_to_batch(&mut batch);
-        self.allow_port("22", &Protocol::TCP, Rc::clone(&inbound), &mut batch)?;
-        Rule::new(Rc::clone(&inbound))
+             .add_to_batch(&mut self.batch);
+        // Allow ssh from anywhere, and log to dmesg with a prefix.
+        Rule::new(Rc::clone(&self.inbound))
+             .dport("22", &Protocol::TCP)?
+             .accept()
+             .log(None, Some(LogPrefix::new("allow ssh connection:")?))
+             .add_to_batch(&mut self.batch);
+
+        // Allow http from all IPs in 192.168.1.255/24 .
+        let local_net = IpNetwork::new([192, 168, 1, 0].into(), 24).unwrap();
+        Rule::new(Rc::clone(&self.inbound))
              .dport("80", &Protocol::TCP)?
              .snetwork(local_net)
              .accept()
-             .add_to_batch(&mut batch);
-        Rule::new(Rc::clone(&inbound))
+             .add_to_batch(&mut self.batch);
+
+        // Allow ICMP traffic, drop IGMP.
+        Rule::new(Rc::clone(&self.inbound))
              .icmp()
              .accept()
-             .add_to_batch(&mut batch);
-        Rule::new(Rc::clone(&inbound))
+             .add_to_batch(&mut self.batch);
+        Rule::new(Rc::clone(&self.inbound))
              .igmp()
              .drop()
-             .add_to_batch(&mut batch);
+             .add_to_batch(&mut self.batch);
 
-        //use nftnl::expr::LogPrefix;
-        //let prefix = "REALM=".to_string() + &self.realm_def.name;
-        Rule::new(Rc::clone(&inbound))
+        // Log all traffic not accepted to NF_LOG group 1, accessible with ulogd.
+        Rule::new(Rc::clone(&self.inbound))
              .log(Some(LogGroup(1)), None)
-             //.log( Some(LogGroup(1)), Some(LogPrefix::new(&prefix)
-             //       .expect("Could not convert log prefix string to CString")))
-             .add_to_batch(&mut batch);
+             .add_to_batch(&mut self.batch);
 
-        let mut finalized_batch = batch.finalize().unwrap();
+        let mut finalized_batch = self.batch.finalize().unwrap();
         send_batch(&mut finalized_batch)?;
-        println!("ruleset applied");
+        println!("table {} commited", TABLE_NAME);
         Ok(())
     }
-    /// If there are any rulesets applied, remove them.
-    pub fn stop(&mut self) -> Result<(), Error> {
-        let table = Table::new(&CString::new(TABLE_NAME)?, ProtoFamily::Inet);
-        let mut batch = Batch::new();
-        batch.add(&table, MsgType::Add);
-        batch.add(&table, MsgType::Del);
+    /// If there is any table with name TABLE_NAME, remove it.
+    pub fn stop(mut self) -> Result<(), Error> {
+        self.batch.add(&self._table, MsgType::Add);
+        self.batch.add(&self._table, MsgType::Del);
+
+        let mut finalized_batch = self.batch.finalize().unwrap();
+        send_batch(&mut finalized_batch)?;
+        println!("table {} destroyed", TABLE_NAME);
         Ok(())
     }
 }
 
+
diff --git a/src/chain_methods.rs b/src/chain_methods.rs
new file mode 100644
index 0000000..81b5fd1
--- /dev/null
+++ b/src/chain_methods.rs
@@ -0,0 +1,78 @@
+use crate::{Batch, Chain, Hook, MsgType, Policy, Table};
+use std::ffi::{CString, NulError};
+use std::rc::Rc;
+
+use serde::{Deserialize, Serialize};
+
+/// A helper trait over [`rustables::Chain`].
+pub trait ChainMethods {
+    /// Create a new Chain instance from a [`Direction`] over a [`rustables::Table`].
+    fn from_direction(direction: &Direction, table: Rc<Table>) -> Result<Self, NulError> where Self: std::marker::Sized;
+    /// Add a [`Verdict`] to the current Chain.
+    fn verdict(self, verdict: &Verdict) -> Self;
+    fn add_to_batch(self, batch: &mut Batch) -> Self;
+}
+
+impl ChainMethods for Chain {
+    fn from_direction(direction: &Direction, table: Rc<Table>) -> Result<Self, NulError> {
+        let chain_name = CString::new(direction.display())?;
+        let mut chain = Chain::new(&chain_name, table);
+        chain.set_hook(direction.get_hook(), 0);
+        Ok(chain)
+    }
+    fn verdict(mut self, verdict: &Verdict) -> Self {
+        self.set_policy(verdict.get());
+        self
+    }
+    fn add_to_batch(self, batch: &mut Batch) -> Self {
+        batch.add(&self, MsgType::Add);
+        self
+    }
+}
+
+/// A Serializable wrapper type around [`rustables::Hook`].
+#[derive(Serialize, Deserialize, Debug, Clone, Eq, PartialEq, Hash)]
+#[serde(rename_all = "snake_case")]
+pub enum Direction {
+    Inbound,
+    Outbound,
+    Forward
+}
+impl Direction {
+    /// Return the Direction's [`rustables::Hook`], ie its representation inside rustables. Note that
+    /// there are Hooks not represented here, namely Prerouting and Postrouting. File a bug if
+    /// you need those.
+    pub fn get_hook(&self) -> Hook {
+        match self {
+            Direction::Inbound => Hook::In,
+            Direction::Outbound => Hook::Out,
+            Direction::Forward => Hook::Forward,
+        }
+    }
+    /// Return a string representation of the Direction.
+    pub fn display(&self) -> String {
+        let s = match self {
+            Direction::Inbound => "inbound",
+            Direction::Outbound => "outbound",
+            Direction::Forward => "forward",
+        };
+        s.to_string()
+    }
+}
+/// A Serializable wrapper type around [`rustables::Policy`].
+#[derive(Serialize, Deserialize, Debug, Clone)]
+#[serde(rename_all = "snake_case")]
+pub enum Verdict {
+    Accept,
+    Drop
+}
+impl Verdict {
+    /// Return the rustables representation of a Verdict (ie, a [`rustables::Policy`]).
+    pub fn get(&self) -> Policy {
+        match self {
+            Verdict::Accept => Policy::Accept,
+            Verdict::Drop => Policy::Drop,
+        }
+    }
+}
+
diff --git a/src/lib.rs b/src/lib.rs
index 33bf12c..511b0ba 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -109,6 +109,9 @@ mod chain;
 pub use chain::{get_chains_cb, list_chains_for_table};
 pub use chain::{Chain, ChainType, Hook, Policy, Priority};
 
+mod chain_methods;
+pub use chain_methods::{ChainMethods, Direction, Verdict};
+
 pub mod query;
 
 mod rule;