From 4937441e1540734d2974f5a2de0a1e8d5e387ad2 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?David=20L=C3=B6nnhager?= <david.l@mullvad.net>
Date: Wed, 24 Jul 2024 14:19:33 +0200
Subject: [PATCH] fixup: improve safety docs

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

diff --git a/src/lib.rs b/src/lib.rs
index c9527c6..de73a92 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -438,7 +438,12 @@ impl PfCtl {
         let wrapped_states = self
             .get_states_inner()?
             .into_iter()
-            .map(State::new)
+            .map(|state| {
+                // SAFETY: `DIOCGETSTATE` calls `pf_state_export`, where `pfsync_state` is zero-initialized.
+                // https://github.com/apple-oss-distributions/xnu/blob/main/bsd/net/pf_ioctl.c#L1247
+                // https://github.com/apple-oss-distributions/xnu/blob/main/bsd/net/pf_ioctl.c#L3583
+                unsafe { State::new(state) }
+            })
             .collect();
         Ok(wrapped_states)
     }
diff --git a/src/state.rs b/src/state.rs
index 0293487..7f30086 100644
--- a/src/state.rs
+++ b/src/state.rs
@@ -3,7 +3,7 @@ use std::net::{Ipv4Addr, Ipv6Addr, SocketAddr};
 
 use crate::ffi::pfvar::pfsync_state_host;
 use crate::{ffi::pfvar::pfsync_state, Direction, Proto};
-use crate::{AddrFamily, Error, ErrorInternal, Result};
+use crate::{Error, ErrorInternal, Result};
 
 /// PF connection state created by a stateful rule
 #[derive(Clone)]
@@ -24,7 +24,15 @@ impl fmt::Debug for State {
 }
 
 impl State {
-    pub(crate) fn new(sync_state: pfsync_state) -> State {
+    /// Wrap `pfsync_state` so that it can be accessed safely.
+    ///
+    /// # Safety
+    ///
+    /// `sync_state.lan` and `sync_state.ext_lan` must set an address and port:
+    /// * If `sync_state.af_lan == PF_INET`, then `host.addr.pfa._v4addr` must be initialized.
+    /// * If `sync_state.af_lan == PF_INET6`, then `host.addr.pfa._v6addr` must be initialized.
+    /// * `host.xport.port` must be initialized.
+    pub(crate) unsafe fn new(sync_state: pfsync_state) -> State {
         State { sync_state }
     }
 
@@ -40,23 +48,13 @@ impl State {
 
     /// Return the local socket address for this state
     pub fn local_address(&self) -> Result<SocketAddr> {
-        // SAFETY: `pf_state_export` sets all fields. The port is zero-initialized if the transport
-        // protocol does not contain ports.
-        //
-        // https://github.com/apple-oss-distributions/xnu/blob/main/bsd/net/pf_ioctl.c#L1247
-        // https://github.com/apple-oss-distributions/xnu/blob/main/bsd/net/pf.c#L4474
-
+        // SAFETY: The address and port are initialized according to the contract of `Self::new`.
         unsafe { parse_address(self.sync_state.af_lan, self.sync_state.lan) }
     }
 
     /// Return the remote socket address for this state
     pub fn remote_address(&self) -> Result<SocketAddr> {
-        // SAFETY: `pf_state_export` sets all fields. The port is zero-initialized if the transport
-        // protocol does not contain ports.
-        //
-        // https://github.com/apple-oss-distributions/xnu/blob/main/bsd/net/pf_ioctl.c#L1247
-        // https://github.com/apple-oss-distributions/xnu/blob/main/bsd/net/pf.c#L4474
-
+        // SAFETY: The address and port are initialized according to the contract of `Self::new`.
         unsafe { parse_address(self.sync_state.af_lan, self.sync_state.ext_lan) }
     }
 
@@ -71,18 +69,16 @@ impl State {
 /// # Safety
 ///
 /// `host` must contain a valid address and a port:
-/// * If `AddrFamily::try_from(family) == Ok(AddrFamily::Ipv4)`, then `host.addr.pfa._v4addr` must
-///   be set.
-/// * If `AddrFamily::try_from(family) == Ok(AddrFamily::Ipv6)`, then `host.addr.pfa._v6addr` must
-///   be set.
+/// * If `family == PF_INET`, then `host.addr.pfa._v4addr` must be initialized.
+/// * If `family == PF_INET6`, then `host.addr.pfa._v6addr` must be initialized.
 /// * `host.xport.port` must always be initialized.
 unsafe fn parse_address(family: u8, host: pfsync_state_host) -> Result<SocketAddr> {
-    let ip = match AddrFamily::try_from(family) {
-        Ok(AddrFamily::Ipv4) => {
+    let ip = match u32::from(family) {
+        crate::ffi::pfvar::PF_INET => {
             // SAFETY: The caller has initialized this memory
             Ipv4Addr::from(u32::from_be(unsafe { host.addr.pfa._v4addr.s_addr })).into()
         }
-        Ok(AddrFamily::Ipv6) => {
+        crate::ffi::pfvar::PF_INET6 => {
             // SAFETY: The caller has initialized this memory
             Ipv6Addr::from(unsafe { host.addr.pfa._v6addr.__u6_addr.__u6_addr8 }).into()
         }