Add support for ct marks

2 files changed, 17 insertions, 2 deletions

diff --git a/nftnl/src/expr/ct.rs b/nftnl/src/expr/ct.rs
index 1bae280..65b47bc 100644
--- a/nftnl/src/expr/ct.rs
+++ b/nftnl/src/expr/ct.rs
@@ -14,12 +14,14 @@ bitflags::bitflags! {
 
 pub enum Conntrack {
     State,
+    Mark { set: bool },
 }
 
 impl Conntrack {
     fn raw_key(&self) -> u32 {
         match *self {
             Conntrack::State => libc::NFT_CT_STATE as u32,
+            Conntrack::Mark { .. } => libc::NFT_CT_MARK as u32,
         }
     }
 }
@@ -29,7 +31,11 @@ impl Expression for Conntrack {
         unsafe {
             let expr = try_alloc!(sys::nftnl_expr_alloc(b"ct\0" as *const _ as *const c_char));
 
-            sys::nftnl_expr_set_u32(expr, sys::NFTNL_EXPR_CT_DREG as u16, libc::NFT_REG_1 as u32);
+            if let Conntrack::Mark { set: true } = self {
+                sys::nftnl_expr_set_u32(expr, sys::NFTNL_EXPR_CT_SREG as u16, libc::NFT_REG_1 as u32);
+            } else {
+                sys::nftnl_expr_set_u32(expr, sys::NFTNL_EXPR_CT_DREG as u16, libc::NFT_REG_1 as u32);
+            }
             sys::nftnl_expr_set_u32(expr, sys::NFTNL_EXPR_CT_KEY as u16, self.raw_key());
 
             expr
@@ -42,4 +48,10 @@ macro_rules! nft_expr_ct {
     (state) => {
         $crate::expr::Conntrack::State
     };
+    (mark set) => {
+        $crate::expr::Conntrack::Mark { set: true }
+    };
+    (mark) => {
+        $crate::expr::Conntrack::Mark { set: false }
+    };
 }
diff --git a/nftnl/src/expr/mod.rs b/nftnl/src/expr/mod.rs
index acce9ed..fd0cc4c 100644
--- a/nftnl/src/expr/mod.rs
+++ b/nftnl/src/expr/mod.rs
@@ -22,7 +22,7 @@ mod counter;
 pub use self::counter::*;
 
 pub mod ct;
-pub use self::ct::Conntrack;
+pub use self::ct::*;
 
 mod immediate;
 pub use self::immediate::*;
@@ -47,6 +47,9 @@ macro_rules! nft_expr {
     (counter) => {
         $crate::expr::Counter
     };
+    (ct $key:ident set) => {
+        nft_expr_ct!($key set)
+    };
     (ct $key:ident) => {
         nft_expr_ct!($key)
     };