rust/ffi: add safe flow storage wrapper

Add a typed FlowStorage<T> wrapper around the flow storage bindings.

Update example and docs.

Ticket: OISF#8447

Author: jasonish

doc/userguide/devguide/extending/flow-lifecycle-callbacks.rst

@@ -131,3 +131,45 @@ The Rust wrappers register closures or function items and return
 The raw pointers passed into callbacks are only valid for the duration
 of the callback invocation and must not be stored. Rust callbacks must
 not panic.
+
+Flow Storage
+============
+
+``flow::FlowStorage<T>`` provides typed, per-flow storage backed by
+Suricata's flow storage API. Each registered slot holds an independent value
+of type ``T`` for every flow.
+
+Register a slot once during initialization with
+``FlowStorage::<T>::register``. Registration must happen before Suricata
+finalizes its storage registration, which is the case during plugin
+initialization.
+
+.. code-block:: rust
+
+   use suricata_ffi::flow::{self, Flow, FlowStorage};
+   use suricata_ffi::thread::ThreadVars;
+   use suricata_sys::sys::Packet;
+
+   #[derive(Default)]
+   struct FlowState {
+       packets: u64,
+   }
+
+   fn register(storage: FlowStorage<FlowState>) -> Result<(), &'static str> {
+       flow::register_update_callback(move |tv, f, p| on_flow_update(storage, tv, f, p))
+   }
+
+Values are owned by Suricata's flow storage and are dropped automatically when
+the flow's storage is freed.
+
+.. code-block:: rust
+
+   fn on_flow_init(storage: FlowStorage<FlowState>, _tv: &mut ThreadVars, f: &mut Flow, _p: *const Packet) {
+       let _ = storage.get_or_insert_with(f, FlowState::default);
+   }
+
+   fn on_flow_update(storage: FlowStorage<FlowState>, _tv: &mut ThreadVars, f: &mut Flow, _p: *mut Packet) {
+       if let Some(state) = storage.get_mut(f) {
+           state.packets += 1;
+       }
+   }

examples/plugins/rust/src/mod.rs

@@ -1,7 +1,7 @@
 use std::ptr::null_mut;
 
 use suricata_ffi::eve::{self, SCJsonBuilder};
-use suricata_ffi::flow::{self, Flow};
+use suricata_ffi::flow::{self, Flow, FlowStorage};
 use suricata_ffi::jsonbuilder::JsonBuilder;
 use suricata_ffi::thread::{self, ThreadStorage, ThreadVars};
 use suricata_ffi::{SCLogError, SCLogNotice, SCLogWarning};
@@ -13,6 +13,12 @@ struct ThreadState {
     flows: u64,
 }
 
+/// Per-flow state stored in Suricata flow storage.
+#[derive(Default)]
+struct FlowState {
+    packets: u64,
+}
+
 unsafe extern "C" fn init() {
     suricata_ffi::plugin::init();
     SCLogNotice!("Initializing rust example plugin");
@@ -26,11 +32,18 @@ unsafe extern "C" fn init() {
             return;
         }
     };
+    let flow_storage = match FlowStorage::<FlowState>::register("rust-example-flow") {
+        Ok(storage) => storage,
+        Err(err) => {
+            SCLogError!("Failed to register rust example flow storage: {}", err);
+            return;
+        }
+    };
 
     if let Err(err) = register_eve_callbacks() {
         SCLogError!("Failed to register rust example EVE callbacks: {}", err);
     }
-    if let Err(err) = register_flow_callbacks(thread_storage) {
+    if let Err(err) = register_flow_callbacks(thread_storage, flow_storage) {
         SCLogError!("Failed to register rust example flow callbacks: {}", err);
     }
     if let Err(err) = register_thread_callbacks(thread_storage) {
@@ -45,10 +58,15 @@ fn register_eve_callbacks() -> Result<(), &'static str> {
     eve::register_callback(log_eve_wrapped)
 }
 
-fn register_flow_callbacks(storage: ThreadStorage<ThreadState>) -> Result<(), &'static str> {
-    flow::register_init_callback(move |tv, f, p| log_flow_init(storage, tv, f, p))?;
-    flow::register_update_callback(log_flow_update)?;
-    flow::register_finish_callback(log_flow_finish)?;
+fn register_flow_callbacks(
+    thread_storage: ThreadStorage<ThreadState>,
+    flow_storage: FlowStorage<FlowState>,
+) -> Result<(), &'static str> {
+    flow::register_init_callback(move |tv, f, p| {
+        log_flow_init(thread_storage, flow_storage, tv, f, p)
+    })?;
+    flow::register_update_callback(move |tv, f, p| log_flow_update(flow_storage, tv, f, p))?;
+    flow::register_finish_callback(move |tv, f| log_flow_finish(flow_storage, tv, f))?;
     Ok(())
 }
 
@@ -94,13 +112,14 @@ fn on_thread_init(storage: ThreadStorage<ThreadState>, tv: &mut ThreadVars) {
 }
 
 fn log_flow_init(
-    storage: ThreadStorage<ThreadState>,
+    thread_storage: ThreadStorage<ThreadState>,
+    flow_storage: FlowStorage<FlowState>,
     tv: &mut ThreadVars,
     f: &mut Flow,
     _p: *const Packet,
 ) {
     // Count flows seen by this thread using the per-thread storage.
-    let flows = match storage.get_mut(tv) {
+    let flows = match thread_storage.get_mut(tv) {
         Some(state) => {
             state.flows += 1;
             state.flows
@@ -110,23 +129,48 @@ fn log_flow_init(
             0
         }
     };
+    // Initialize the per-flow storage for this flow.
+    if let Err(err) = flow_storage.get_or_insert_with(f, FlowState::default) {
+        SCLogError!("failed to initialize rust example flow storage: {}", err);
+    }
     SCLogNotice!(
         "rust example flow init callback: flow={:p}, thread_flows={}",
         f.as_ptr(),
         flows
     );
 }
 
-fn log_flow_update(_tv: &mut ThreadVars, f: &mut Flow, _p: *mut Packet) {
+fn log_flow_update(
+    flow_storage: FlowStorage<FlowState>,
+    _tv: &mut ThreadVars,
+    f: &mut Flow,
+    _p: *mut Packet,
+) {
+    // Count packets seen on this flow using the per-flow storage.
+    let packets = match flow_storage.get_mut(f) {
+        Some(state) => {
+            state.packets += 1;
+            state.packets
+        }
+        None => {
+            SCLogWarning!("rust example flow storage was not initialized");
+            0
+        }
+    };
     SCLogNotice!(
-        "rust example flow update callback: flow={:p}, packet={:p}",
+        "rust example flow update callback: flow={:p}, flow_packets={}",
         f.as_ptr(),
-        _p
+        packets
     );
 }
 
-fn log_flow_finish(_tv: &mut ThreadVars, f: &mut Flow) {
-    SCLogNotice!("rust example flow finish callback: flow={:p}", f.as_ptr());
+fn log_flow_finish(flow_storage: FlowStorage<FlowState>, _tv: &mut ThreadVars, f: &mut Flow) {
+    let packets = flow_storage.get(f).map(|state| state.packets).unwrap_or(0);
+    SCLogNotice!(
+        "rust example flow finish callback: flow={:p}, flow_packets={}",
+        f.as_ptr(),
+        packets
+    );
 }
 
 #[no_mangle]

rust/ffi/src/flow.rs

@@ -15,12 +15,13 @@
  * 02110-1301, USA.
  */
 
+use std::ffi::CString;
 use std::marker::PhantomData;
 use std::os::raw::c_void;
 
 use suricata_sys::sys::{
-    self, Packet, SCFlowRegisterFinishCallback, SCFlowRegisterInitCallback,
-    SCFlowRegisterUpdateCallback,
+    self, Packet, SCFlowGetStorageById, SCFlowRegisterFinishCallback, SCFlowRegisterInitCallback,
+    SCFlowRegisterUpdateCallback, SCFlowSetStorageById, SCFlowStorageId, SCFlowStorageRegister,
 };
 
 use crate::thread::ThreadVars;
@@ -62,6 +63,119 @@ impl<'a> Flow<'a> {
     }
 }
 
+/// A typed handle to a per-flow storage slot.
+///
+/// `FlowStorage<T>` wraps the `SCFlowStorageId` returned when registering flow
+/// storage with Suricata. Values are stored as a `Box<T>` owned by Suricata's
+/// flow storage and are dropped automatically when the flow's storage is freed.
+///
+/// The handle only holds the storage id, so it is `Copy` and `Send`/`Sync`
+/// regardless of `T`, and can be passed by value into the callbacks that need
+/// it.
+pub struct FlowStorage<T> {
+    id: SCFlowStorageId,
+    _marker: PhantomData<fn() -> T>,
+}
+
+// Manual `Copy`/`Clone` impls so the handle is copyable regardless of whether
+// `T` is; it only holds the storage id.
+impl<T> Clone for FlowStorage<T> {
+    fn clone(&self) -> Self {
+        *self
+    }
+}
+
+impl<T> Copy for FlowStorage<T> {}
+
+impl<T: Send + 'static> FlowStorage<T> {
+    /// Register a new flow storage slot for values of type `T`.
+    ///
+    /// `name` must be unique among registered flow storage. Registration has to
+    /// happen during initialization, before Suricata finalizes storage
+    /// registration (`SCStorageFinalize`).
+    ///
+    /// Returns an error if `name` contains an interior nul byte or if Suricata
+    /// rejects the registration.
+    pub fn register(name: &str) -> Result<Self, &'static str> {
+        let name = CString::new(name).map_err(|_| "flow storage name contains a nul byte")?;
+        let id = unsafe { SCFlowStorageRegister(name.as_ptr(), Some(Self::free)) };
+        if id.id < 0 {
+            return Err("Failed to register flow storage");
+        }
+
+        // Suricata keeps the storage name pointer in its storage mapping for
+        // the lifetime of the process, so the CString is intentionally leaked.
+        std::mem::forget(name);
+
+        Ok(Self {
+            id,
+            _marker: PhantomData,
+        })
+    }
+
+    /// Return a reference to the value stored for `f`, if any.
+    pub fn get<'f>(&self, f: &'f Flow<'_>) -> Option<&'f T> {
+        let ptr = unsafe { SCFlowGetStorageById(f.as_ptr(), self.id) };
+        if ptr.is_null() {
+            None
+        } else {
+            Some(unsafe { &*(ptr as *const T) })
+        }
+    }
+
+    /// Return a mutable reference to the value stored for `f`, if any.
+    ///
+    /// Takes `&mut Flow` so the returned `&mut T` is the only live reference to
+    /// the stored value for the duration of the borrow.
+    pub fn get_mut<'f>(&self, f: &'f mut Flow<'_>) -> Option<&'f mut T> {
+        let ptr = unsafe { SCFlowGetStorageById(f.as_ptr(), self.id) };
+        if ptr.is_null() {
+            None
+        } else {
+            Some(unsafe { &mut *(ptr as *mut T) })
+        }
+    }
+
+    /// Return a mutable reference to the value stored for `f`, inserting the
+    /// value produced by `init` if none is present yet.
+    ///
+    /// Takes `&mut Flow` so the returned `&mut T` is the only live reference to
+    /// the stored value for the duration of the borrow.
+    pub fn get_or_insert_with<'f>(
+        &self, f: &'f mut Flow<'_>, init: impl FnOnce() -> T,
+    ) -> Result<&'f mut T, &'static str> {
+        let ptr = unsafe { SCFlowGetStorageById(f.as_ptr(), self.id) };
+        if !ptr.is_null() {
+            return Ok(unsafe { &mut *(ptr as *mut T) });
+        }
+
+        // `SCFlowSetStorageById` overwrites the slot without freeing any
+        // previous value; we only reach here when the slot is empty.
+        let ptr = Box::into_raw(Box::new(init()));
+        let rc = unsafe { SCFlowSetStorageById(f.as_mut_ptr(), self.id, ptr.cast()) };
+        if rc != 0 {
+            unsafe {
+                drop(Box::from_raw(ptr));
+            }
+            return Err("Failed to set flow storage");
+        }
+
+        Ok(unsafe { &mut *ptr })
+    }
+
+    /// Free callback registered with Suricata flow storage that drops the
+    /// `Box<T>` backing a stored value.
+    unsafe extern "C" fn free(ptr: *mut c_void) {
+        if !ptr.is_null() {
+            // The drop runs across an FFI boundary, so guard against unwinding
+            // into C if `T`'s `Drop` panics.
+            let _ = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
+                drop(Box::from_raw(ptr as *mut T));
+            }));
+        }
+    }
+}
+
 /// Register a flow initialization callback.
 ///
 /// The callback is invoked whenever Suricata initializes a flow. It receives: