Refactor: Use enums for subsystem and message types

This refactors the crate to use type-safe enums for netfilter subsystems and message types, for a safer and more idiomatic API.

- Introduces a `Subsystem` enum to replace raw `u8` identifiers for `NfLog` and `Conntrack` subsystems.

- Introduces `NfLogMessageType` and `ConntrackMessageType` enums to provide type safety for messages within each subsystem.

- Makes the top-level `NetfilterMessage::message_type()` function private to guide users towards the safer pattern of matching on `NetfilterMessageInner`.

- Updates the internal parsing logic in `buffer.rs` to use the new `Subsystem` enum.

Signed-off-by: Shivang K Raghuvanshi <[email protected]>

Author: shivkr6

src/buffer.rs

@@ -3,7 +3,7 @@
 use crate::{
     conntrack::ConntrackMessage,
     message::{
-        NetfilterHeader, NetfilterMessage, NetfilterMessageInner,
+        NetfilterHeader, NetfilterMessage, NetfilterMessageInner, Subsystem,
         NETFILTER_HEADER_LEN,
     },
     nflog::NfLogMessage,
@@ -53,17 +53,17 @@ impl<'a, T: AsRef<[u8]> + ?Sized>
             .context("failed to parse netfilter header")?;
         let subsys = (message_type >> 8) as u8;
         let message_type = message_type as u8;
-        let inner = match subsys {
-            NfLogMessage::SUBSYS => NetfilterMessageInner::NfLog(
+        let inner = match Subsystem::from(subsys) {
+            Subsystem::NfLog => NetfilterMessageInner::NfLog(
                 NfLogMessage::parse_with_param(buf, message_type)
                     .context("failed to parse nflog payload")?,
             ),
-            ConntrackMessage::SUBSYS => NetfilterMessageInner::Conntrack(
+            Subsystem::Conntrack => NetfilterMessageInner::Conntrack(
                 ConntrackMessage::parse_with_param(buf, message_type)
                     .context("failed to parse conntrack payload")?,
             ),
-            _ => NetfilterMessageInner::Other {
-                subsys,
+            subsys_enum @ Subsystem::Other(_) => NetfilterMessageInner::Other {
+                subsys: subsys_enum,
                 message_type,
                 attributes: buf.default_nlas()?,
             },

src/conntrack/message.rs

@@ -1,10 +1,6 @@
 // SPDX-License-Identifier: MIT
 
-use crate::{
-    buffer::NetfilterBuffer,
-    conntrack::attributes::ConntrackNla,
-    constants::{IPCTNL_MSG_CT_GET, NFNL_SUBSYS_CTNETLINK},
-};
+use crate::{buffer::NetfilterBuffer, conntrack::attributes::ConntrackNla};
 use netlink_packet_core::{
     DecodeError, DefaultNla, Emitable, Parseable, ParseableParametrized,
 };
@@ -19,13 +15,40 @@ pub enum ConntrackMessage {
     },
 }
 
-impl ConntrackMessage {
-    pub(crate) const SUBSYS: u8 = NFNL_SUBSYS_CTNETLINK;
+const IPCTNL_MSG_CT_GET: u8 = 1;
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+#[non_exhaustive]
+pub enum ConntrackMessageType {
+    Get,
+    Other(u8),
+}
+
+impl From<u8> for ConntrackMessageType {
+    fn from(value: u8) -> Self {
+        match value {
+            IPCTNL_MSG_CT_GET => Self::Get,
+            v => Self::Other(v),
+        }
+    }
+}
+
+impl From<ConntrackMessageType> for u8 {
+    fn from(value: ConntrackMessageType) -> Self {
+        match value {
+            ConntrackMessageType::Get => IPCTNL_MSG_CT_GET,
+            ConntrackMessageType::Other(v) => v,
+        }
+    }
+}
 
-    pub fn message_type(&self) -> u8 {
+impl ConntrackMessage {
+    pub fn message_type(&self) -> ConntrackMessageType {
         match self {
-            ConntrackMessage::Get(_) => IPCTNL_MSG_CT_GET,
-            ConntrackMessage::Other { message_type, .. } => *message_type,
+            ConntrackMessage::Get(_) => ConntrackMessageType::Get,
+            ConntrackMessage::Other { message_type, .. } => {
+                (*message_type).into()
+            }
         }
     }
 }
@@ -61,16 +84,18 @@ impl<'a, T: AsRef<[u8]> + ?Sized>
         buf: &NetfilterBuffer<&'a T>,
         message_type: u8,
     ) -> Result<Self, DecodeError> {
-        Ok(match message_type {
-            IPCTNL_MSG_CT_GET => {
+        Ok(match ConntrackMessageType::from(message_type) {
+            ConntrackMessageType::Get => {
                 let attributes = buf
                     .parse_all_nlas(|nla_buf| ConntrackNla::parse(&nla_buf))?;
                 ConntrackMessage::Get(attributes)
             }
-            _ => ConntrackMessage::Other {
-                message_type,
-                attributes: buf.default_nlas()?,
-            },
+            ConntrackMessageType::Other(message_type) => {
+                ConntrackMessage::Other {
+                    message_type,
+                    attributes: buf.default_nlas()?,
+                }
+            }
         })
     }
 }

src/conntrack/mod.rs

@@ -1,7 +1,7 @@
 // SPDX-License-Identifier: MIT
 
 mod message;
-pub use message::ConntrackMessage;
+pub use message::{ConntrackMessage, ConntrackMessageType};
 mod attributes;
 pub use attributes::{
     ConntrackNla, IPTuple, ProtoInfo, ProtoInfoTCP, ProtoTuple, Protocol,

src/constants.rs

@@ -43,10 +43,8 @@ pub const AF_ALG: u8 = libc::AF_ALG as u8;
 pub const NFNETLINK_V0: u8 = libc::NFNETLINK_V0 as u8;
 
 pub const NFNL_SUBSYS_NONE: u8 = libc::NFNL_SUBSYS_NONE as u8;
-pub const NFNL_SUBSYS_CTNETLINK: u8 = libc::NFNL_SUBSYS_CTNETLINK as u8;
 pub const NFNL_SUBSYS_CTNETLINK_EXP: u8 = libc::NFNL_SUBSYS_CTNETLINK_EXP as u8;
 pub const NFNL_SUBSYS_QUEUE: u8 = libc::NFNL_SUBSYS_QUEUE as u8;
-pub const NFNL_SUBSYS_ULOG: u8 = libc::NFNL_SUBSYS_ULOG as u8;
 pub const NFNL_SUBSYS_OSF: u8 = libc::NFNL_SUBSYS_OSF as u8;
 pub const NFNL_SUBSYS_IPSET: u8 = libc::NFNL_SUBSYS_IPSET as u8;
 pub const NFNL_SUBSYS_ACCT: u8 = libc::NFNL_SUBSYS_ACCT as u8;
@@ -83,8 +81,3 @@ pub const NFULA_HWHEADER: u16 = libc::NFULA_HWHEADER as u16;
 pub const NFULA_HWLEN: u16 = libc::NFULA_HWLEN as u16;
 pub const NFULA_CT: u16 = libc::NFULA_CT as u16;
 pub const NFULA_CT_INFO: u16 = libc::NFULA_CT_INFO as u16;
-
-pub const NFULNL_MSG_CONFIG: u8 = libc::NFULNL_MSG_CONFIG as u8;
-pub const NFULNL_MSG_PACKET: u8 = libc::NFULNL_MSG_PACKET as u8;
-
-pub(crate) const IPCTNL_MSG_CT_GET: u8 = 1;

src/lib.rs

@@ -3,7 +3,9 @@
 pub(crate) mod buffer;
 pub mod constants;
 mod message;
-pub use message::{NetfilterHeader, NetfilterMessage, NetfilterMessageInner};
+pub use message::{
+    NetfilterHeader, NetfilterMessage, NetfilterMessageInner, Subsystem,
+};
 pub mod conntrack;
 pub mod nflog;
 #[cfg(test)]

src/message.rs

@@ -60,13 +60,44 @@ impl<T: AsRef<[u8]>> Parseable<NetfilterHeaderBuffer<T>> for NetfilterHeader {
     }
 }
 
+const NFNL_SUBSYS_CTNETLINK: u8 = libc::NFNL_SUBSYS_CTNETLINK as u8;
+const NFNL_SUBSYS_ULOG: u8 = libc::NFNL_SUBSYS_ULOG as u8;
+
+#[derive(Debug, PartialEq, Eq, Clone, Copy)]
+#[non_exhaustive]
+pub enum Subsystem {
+    NfLog,
+    Conntrack,
+    Other(u8),
+}
+
+impl From<u8> for Subsystem {
+    fn from(value: u8) -> Self {
+        match value {
+            NFNL_SUBSYS_ULOG => Self::NfLog,
+            NFNL_SUBSYS_CTNETLINK => Self::Conntrack,
+            v => Self::Other(v),
+        }
+    }
+}
+
+impl From<Subsystem> for u8 {
+    fn from(value: Subsystem) -> Self {
+        match value {
+            Subsystem::NfLog => NFNL_SUBSYS_ULOG,
+            Subsystem::Conntrack => NFNL_SUBSYS_CTNETLINK,
+            Subsystem::Other(v) => v,
+        }
+    }
+}
+
 #[derive(Debug, PartialEq, Eq, Clone)]
 #[non_exhaustive]
 pub enum NetfilterMessageInner {
     NfLog(NfLogMessage),
     Conntrack(ConntrackMessage),
     Other {
-        subsys: u8,
+        subsys: Subsystem,
         message_type: u8,
         attributes: Vec<DefaultNla>,
     },
@@ -123,19 +154,21 @@ impl NetfilterMessage {
         }
     }
 
-    pub fn subsys(&self) -> u8 {
+    pub fn subsys(&self) -> Subsystem {
         match self.inner {
-            NetfilterMessageInner::NfLog(_) => NfLogMessage::SUBSYS,
-            NetfilterMessageInner::Conntrack(_) => ConntrackMessage::SUBSYS,
+            NetfilterMessageInner::NfLog(_) => Subsystem::NfLog,
+            NetfilterMessageInner::Conntrack(_) => Subsystem::Conntrack,
             NetfilterMessageInner::Other { subsys, .. } => subsys,
         }
     }
 
-    pub fn message_type(&self) -> u8 {
+    fn message_type(&self) -> u8 {
         match self.inner {
-            NetfilterMessageInner::NfLog(ref message) => message.message_type(),
+            NetfilterMessageInner::NfLog(ref message) => {
+                message.message_type().into()
+            }
             NetfilterMessageInner::Conntrack(ref message) => {
-                message.message_type()
+                message.message_type().into()
             }
             NetfilterMessageInner::Other { message_type, .. } => message_type,
         }
@@ -155,7 +188,7 @@ impl Emitable for NetfilterMessage {
 
 impl NetlinkSerializable for NetfilterMessage {
     fn message_type(&self) -> u16 {
-        ((self.subsys() as u16) << 8) | self.message_type() as u16
+        ((u8::from(self.subsys()) as u16) << 8) | self.message_type() as u16
     }
 
     fn buffer_len(&self) -> usize {

src/nflog/message.rs

@@ -6,7 +6,6 @@ use netlink_packet_core::{
 
 use crate::{
     buffer::NetfilterBuffer,
-    constants::{NFNL_SUBSYS_ULOG, NFULNL_MSG_CONFIG, NFULNL_MSG_PACKET},
     nflog::nlas::{config::ConfigNla, packet::PacketNla},
 };
 
@@ -20,14 +19,43 @@ pub enum NfLogMessage {
     },
 }
 
-impl NfLogMessage {
-    pub const SUBSYS: u8 = NFNL_SUBSYS_ULOG;
+const NFULNL_MSG_CONFIG: u8 = libc::NFULNL_MSG_CONFIG as u8;
+const NFULNL_MSG_PACKET: u8 = libc::NFULNL_MSG_PACKET as u8;
+
+#[derive(Debug, PartialEq, Eq, Clone, Copy)]
+#[non_exhaustive]
+pub enum NfLogMessageType {
+    Config,
+    Packet,
+    Other(u8),
+}
+
+impl From<u8> for NfLogMessageType {
+    fn from(value: u8) -> Self {
+        match value {
+            NFULNL_MSG_CONFIG => Self::Config,
+            NFULNL_MSG_PACKET => Self::Packet,
+            v => Self::Other(v),
+        }
+    }
+}
+
+impl From<NfLogMessageType> for u8 {
+    fn from(value: NfLogMessageType) -> Self {
+        match value {
+            NfLogMessageType::Config => NFULNL_MSG_CONFIG,
+            NfLogMessageType::Packet => NFULNL_MSG_PACKET,
+            NfLogMessageType::Other(v) => v,
+        }
+    }
+}
 
-    pub fn message_type(&self) -> u8 {
+impl NfLogMessage {
+    pub fn message_type(&self) -> NfLogMessageType {
         match self {
-            NfLogMessage::Config(_) => NFULNL_MSG_CONFIG,
-            NfLogMessage::Packet(_) => NFULNL_MSG_PACKET,
-            NfLogMessage::Other { message_type, .. } => *message_type,
+            NfLogMessage::Config(_) => NfLogMessageType::Config,
+            NfLogMessage::Packet(_) => NfLogMessageType::Packet,
+            NfLogMessage::Other { message_type, .. } => (*message_type).into(),
         }
     }
 }
@@ -57,18 +85,18 @@ impl<'a, T: AsRef<[u8]> + ?Sized>
         buf: &NetfilterBuffer<&'a T>,
         message_type: u8,
     ) -> Result<Self, DecodeError> {
-        Ok(match message_type {
-            NFULNL_MSG_CONFIG => {
+        Ok(match NfLogMessageType::from(message_type) {
+            NfLogMessageType::Config => {
                 let nlas =
                     buf.parse_all_nlas(|nla_buf| ConfigNla::parse(&nla_buf))?;
                 NfLogMessage::Config(nlas)
             }
-            NFULNL_MSG_PACKET => {
+            NfLogMessageType::Packet => {
                 let nlas =
                     buf.parse_all_nlas(|nla_buf| PacketNla::parse(&nla_buf))?;
                 NfLogMessage::Packet(nlas)
             }
-            _ => NfLogMessage::Other {
+            NfLogMessageType::Other(message_type) => NfLogMessage::Other {
                 message_type,
                 nlas: buf.default_nlas()?,
             },

src/nflog/mod.rs

@@ -1,7 +1,7 @@
 // SPDX-License-Identifier: MIT
 
 mod message;
-pub use message::NfLogMessage;
+pub use message::{NfLogMessage, NfLogMessageType};
 pub mod nlas;
 
 use netlink_packet_core::{

src/tests.rs

@@ -2,16 +2,16 @@
 
 use std::net::IpAddr;
 
-use libc::NFNL_SUBSYS_CTNETLINK;
 use netlink_packet_core::{Emitable, ParseableParametrized};
 
 use crate::{
     buffer::NetfilterBuffer,
     conntrack::{
-        ConntrackMessage, ConntrackNla, IPTuple, ProtoInfo, ProtoInfoTCP,
-        ProtoTuple, Protocol, TCPFlags, Tuple,
+        ConntrackMessage, ConntrackMessageType, ConntrackNla, IPTuple,
+        ProtoInfo, ProtoInfoTCP, ProtoTuple, Protocol, TCPFlags, Tuple,
     },
-    constants::{AF_INET, AF_INET6, AF_UNSPEC, IPCTNL_MSG_CT_GET},
+    constants::{AF_INET, AF_INET6, AF_UNSPEC},
+    message::Subsystem,
     NetfilterHeader, NetfilterMessage,
 };
 
@@ -32,8 +32,8 @@ fn test_dump_conntrack() {
     // Check if the serialization was correct
     assert_eq!(buffer, raw);
 
-    let message_type =
-        ((NFNL_SUBSYS_CTNETLINK as u16) << 8) | (IPCTNL_MSG_CT_GET as u16);
+    let message_type = ((u8::from(Subsystem::Conntrack) as u16) << 8)
+        | (u8::from(ConntrackMessageType::Get) as u16);
     // Check if the deserialization was correct
     assert_eq!(
         NetfilterMessage::parse_with_param(
@@ -100,8 +100,8 @@ fn test_get_conntrack_tcp_ipv4() {
     // Check if the serialization was correct
     assert_eq!(buffer, raw);
 
-    let message_type =
-        ((NFNL_SUBSYS_CTNETLINK as u16) << 8) | (IPCTNL_MSG_CT_GET as u16);
+    let message_type = ((u8::from(Subsystem::Conntrack) as u16) << 8)
+        | (u8::from(ConntrackMessageType::Get) as u16);
     // Check if the deserialization was correct
     assert_eq!(
         NetfilterMessage::parse_with_param(
@@ -156,8 +156,8 @@ fn test_get_conntrack_udp_ipv6() {
     // Check if the serialization was correct
     assert_eq!(buffer, raw);
 
-    let message_type =
-        ((NFNL_SUBSYS_CTNETLINK as u16) << 8) | (IPCTNL_MSG_CT_GET as u16);
+    let message_type = ((u8::from(Subsystem::Conntrack) as u16) << 8)
+        | (u8::from(ConntrackMessageType::Get) as u16);
     // Check if the deserialization was correct
     assert_eq!(
         NetfilterMessage::parse_with_param(