From a56581903203dbe0e2b80e1c76842103f4f4be87 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?David=20L=C3=B6nnhager?= <david.l@mullvad.net>
Date: Tue, 23 Jul 2024 12:39:54 +0200
Subject: [PATCH] fixup: kill_state + get_states instead of filter_states

---
 src/lib.rs   | 38 ++++++++++++++++++--------------------
 src/state.rs |  5 +++++
 2 files changed, 23 insertions(+), 20 deletions(-)

diff --git a/src/lib.rs b/src/lib.rs
index 2b0576a..117dbf9 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -385,7 +385,7 @@ impl PfCtl {
     /// Returns total number of removed states upon success, otherwise
     /// ErrorKind::AnchorDoesNotExist if anchor does not exist.
     pub fn clear_states(&mut self, anchor_name: &str, kind: AnchorKind) -> Result<u32> {
-        let pfsync_states = self.get_states()?;
+        let pfsync_states = self.get_states_inner()?;
         if !pfsync_states.is_empty() {
             self.with_anchor_rule(anchor_name, kind, |anchor_rule| {
                 pfsync_states
@@ -418,29 +418,27 @@ impl PfCtl {
         Ok(pfioc_state_kill.psk_af as u32)
     }
 
-    /// Keep all states for which `filter` returns true.
-    /// Returns total number of removed states upon success.
-    pub fn filter_states(&mut self, filter: impl Fn(&State) -> bool) -> Result<u32> {
-        let states = self.get_states()?;
-        states
+    /// Get all states created by stateful rules
+    pub fn get_states(&mut self) -> Result<Vec<State>> {
+        let wrapped_states = self
+            .get_states_inner()?
             .into_iter()
-            .filter(|state| {
-                let parsed_state = State::new(*state);
-                !filter(&parsed_state)
-            })
-            .map(|pfsync_state| {
-                let mut pfioc_state_kill = unsafe { mem::zeroed::<ffi::pfvar::pfioc_state_kill>() };
-                setup_pfioc_state_kill(&pfsync_state, &mut pfioc_state_kill);
-                ioctl_guard!(ffi::pf_kill_states(self.fd(), &mut pfioc_state_kill))?;
-                // psk_af holds the number of killed states
-                Ok(pfioc_state_kill.psk_af as u32)
-            })
-            .collect::<Result<Vec<_>>>()
-            .map(|v| v.iter().sum())
+            .map(State::new)
+            .collect();
+        Ok(wrapped_states)
+    }
+
+    /// Remove the specified state
+    pub fn kill_state(&mut self, state: &State) -> Result<()> {
+        let mut pfioc_state_kill = unsafe { mem::zeroed::<ffi::pfvar::pfioc_state_kill>() };
+        setup_pfioc_state_kill(state.as_raw(), &mut pfioc_state_kill);
+        ioctl_guard!(ffi::pf_kill_states(self.fd(), &mut pfioc_state_kill))?;
+        // psk_af holds the number of killed states, but it should be zero
+        Ok(())
     }
 
     /// Get all states created by stateful rules
-    fn get_states(&mut self) -> Result<Vec<ffi::pfvar::pfsync_state>> {
+    fn get_states_inner(&mut self) -> Result<Vec<ffi::pfvar::pfsync_state>> {
         let num_states = self.get_num_states()?;
         if num_states > 0 {
             let (mut pfioc_states, pfsync_states) = setup_pfioc_states(num_states);
diff --git a/src/state.rs b/src/state.rs
index 8c39b58..e7dd959 100644
--- a/src/state.rs
+++ b/src/state.rs
@@ -34,6 +34,11 @@ impl State {
     pub fn remote_address(&self) -> Result<SocketAddr> {
         parse_address(self.sync_state.af_lan, self.sync_state.ext_lan)
     }
+
+    /// Return a reference to the inner `pfsync_state` state
+    pub(crate) fn as_raw(&self) -> &pfsync_state {
+        &self.sync_state
+    }
 }
 
 fn parse_address(family: u8, host: pfsync_state_host) -> Result<SocketAddr> {