new(plugin): validate key types at runtime (including nested tables)

Signed-off-by: Grzegorz Nosek <[email protected]>

falco_plugin/src/plugin/tables/table/mod.rs

@@ -61,7 +61,7 @@ where
 {
     /// Look up an entry in `table` corresponding to `key`
     pub fn get_entry(&self, reader_vtable: &TableReader, key: &K) -> Result<E, Error> {
-        let raw_entry = unsafe { self.raw_table.get_entry(reader_vtable, key)? };
+        let raw_entry = self.raw_table.get_entry(reader_vtable, key)?;
         Ok(E::new(
             raw_entry,
             self.raw_table.table,
@@ -186,13 +186,14 @@ where
     /// The call to `get_table_field` takes a closure, which is passed a reference to
     /// a sample entry in the nested table (enough to manage its fields) and returns a tuple
     /// of `(the table field, (whatever_the_closure_returns))`
-    pub fn get_table_field<V, U, F, R>(
+    pub fn get_table_field<NK, V, U, F, R>(
         &self,
         tables_input: &TablesInput,
         name: &CStr,
         func: F,
     ) -> Result<(Field<V, E>, R), Error>
     where
+        NK: Key,
         for<'a> V::AssocData: From<&'a M>,
         V: Value + ?Sized,
         U: Entry,
@@ -204,7 +205,7 @@ where
 
         let fields = unsafe {
             self.raw_table
-                .with_subtable(field.field, tables_input, |subtable| {
+                .with_subtable::<NK, _, _>(field.field, tables_input, |subtable| {
                     let owned = RawTable {
                         table: subtable.table,
                     };
@@ -302,7 +303,7 @@ where
         field: *mut ss_plugin_table_field_t,
         tables_input: &TablesInput,
     ) -> Result<Self::AssocData, Error> {
-        table.with_subtable(field, tables_input, |subtable| {
+        table.with_subtable::<K, _, _>(field, tables_input, |subtable| {
             M::new(subtable, tables_input)
         })?
     }

falco_plugin/src/plugin/tables/table/raw.rs

@@ -1,5 +1,5 @@
 use crate::plugin::error::as_result::{AsResult, WithLastError};
-use crate::plugin::tables::data::{Key, Value};
+use crate::plugin::tables::data::{FieldTypeId, Key, Value};
 use crate::plugin::tables::entry::raw::RawEntry;
 use crate::plugin::tables::field::raw::RawField;
 use crate::plugin::tables::traits::TableMetadata;
@@ -11,6 +11,7 @@ use falco_plugin_api::{
     ss_plugin_table_entry_t, ss_plugin_table_field_t, ss_plugin_table_fieldinfo,
     ss_plugin_table_iterator_func_t, ss_plugin_table_iterator_state_t, ss_plugin_table_t,
 };
+use num_traits::FromPrimitive;
 use std::ffi::CStr;
 use std::ops::ControlFlow;
 
@@ -105,15 +106,20 @@ impl RawTable {
     }
 
     /// # Look up an entry in `table` corresponding to `key`
-    ///
-    /// # Safety
-    /// The key type must be the same as actually used by the table. Using the wrong type
-    /// (especially using a number if the real key type is a string) will lead to UB.
-    pub unsafe fn get_entry<K: Key>(
+    pub fn get_entry<K: Key>(
         &self,
         reader_vtable: &TableReader,
         key: &K,
     ) -> Result<RawEntry, anyhow::Error> {
+        let input = unsafe { &*(self.table as *mut falco_plugin_api::ss_plugin_table_input) };
+        if input.key_type != K::TYPE_ID as ss_plugin_state_type {
+            anyhow::bail!(
+                "Bad key type, requested {:?}, table has {:?}",
+                K::TYPE_ID,
+                FieldTypeId::from_u32(input.key_type),
+            );
+        }
+
         let entry =
             unsafe { (reader_vtable.get_table_entry)(self.table, &key.to_data() as *const _) };
 
@@ -239,13 +245,14 @@ impl RawTable {
         unsafe { Ok((writer_vtable.clear_table)(self.table).as_result()?) }
     }
 
-    pub(in crate::plugin::tables) unsafe fn with_subtable<F, R>(
+    pub(in crate::plugin::tables) unsafe fn with_subtable<K, F, R>(
         &self,
         field: *mut ss_plugin_table_field_t,
         tables_input: &TablesInput,
         func: F,
     ) -> Result<R, anyhow::Error>
     where
+        K: Key,
         F: FnOnce(&RawTable) -> R,
     {
         let entry = unsafe { (tables_input.writer_ext.create_table_entry)(self.table) };
@@ -262,6 +269,16 @@ impl RawTable {
             anyhow::bail!("Failed to get field value for temporary table entry")
         }
 
+        let input = unsafe { &*(val.table as *mut falco_plugin_api::ss_plugin_table_input) };
+        if input.key_type != K::TYPE_ID as ss_plugin_state_type {
+            unsafe { (tables_input.writer_ext.destroy_table_entry)(self.table, entry) };
+            anyhow::bail!(
+                "Bad key type, requested {:?}, table has {:?}",
+                K::TYPE_ID,
+                FieldTypeId::from_u32(input.key_type),
+            );
+        }
+
         let raw_table = unsafe { RawTable { table: val.table } };
         let ret = func(&raw_table);
         unsafe { (tables_input.writer_ext.destroy_table_entry)(self.table, entry) };
@@ -270,13 +287,13 @@ impl RawTable {
 
     #[doc(hidden)]
     // this is not really intended to be called by the end user, it's just for the derive macros
-    pub fn get_metadata<M: TableMetadata, V: Value + ?Sized>(
+    pub fn get_metadata<K: Key, M: TableMetadata, V: Value + ?Sized>(
         &self,
         field: &RawField<V>,
         tables_input: &TablesInput,
     ) -> Result<M, anyhow::Error> {
         unsafe {
-            self.with_subtable(field.field, tables_input, |subtable| {
+            self.with_subtable::<K, _, _>(field.field, tables_input, |subtable| {
                 M::new(subtable, tables_input)
             })
         }?

falco_plugin_tests/tests/scap_import_table_bad_key_type.rs

@@ -0,0 +1,84 @@
+use falco_plugin::anyhow;
+use falco_plugin::anyhow::Error;
+use falco_plugin::base::Plugin;
+use falco_plugin::event::events::types::EventType;
+use falco_plugin::extract::EventInput;
+use falco_plugin::parse::{ParseInput, ParsePlugin};
+use falco_plugin::static_plugin;
+use falco_plugin::tables::import::{Entry, Field, Table, TableMetadata};
+use falco_plugin::tables::TablesInput;
+use std::ffi::CStr;
+use std::sync::Arc;
+
+type Fd = Entry<Arc<FdMetadata>>;
+type FdTable = Table<i32, Fd>;
+
+#[derive(TableMetadata)]
+#[entry_type(Fd)]
+struct FdMetadata {
+    fd: Field<i64, Fd>,
+
+    #[name(c"type")]
+    fd_type: Field<u8, Fd>,
+}
+
+type Thread = Entry<Arc<ThreadMetadata>>;
+type ThreadTable = Table<i64, Thread>;
+
+#[derive(TableMetadata)]
+#[entry_type(Thread)]
+struct ThreadMetadata {
+    comm: Field<CStr, Thread>,
+    file_descriptors: Field<FdTable, Thread>,
+
+    #[custom]
+    num_events: Field<u64, Thread>,
+}
+
+struct DummyPlugin {
+    #[allow(dead_code)]
+    threads: ThreadTable,
+}
+
+impl Plugin for DummyPlugin {
+    const NAME: &'static CStr = c"dummy";
+    const PLUGIN_VERSION: &'static CStr = c"0.0.0";
+    const DESCRIPTION: &'static CStr = c"test plugin";
+    const CONTACT: &'static CStr = c"[email protected]";
+    type ConfigType = ();
+
+    fn new(input: Option<&TablesInput>, _config: Self::ConfigType) -> Result<Self, Error> {
+        let Some(input) = input else {
+            anyhow::bail!("Did not get tables input")
+        };
+
+        let threads = input.get_table(c"threads")?;
+
+        Ok(Self { threads })
+    }
+}
+
+impl ParsePlugin for DummyPlugin {
+    const EVENT_TYPES: &'static [EventType] = &[];
+    const EVENT_SOURCES: &'static [&'static str] = &["syscall"];
+
+    fn parse_event(
+        &mut self,
+        _event: &EventInput,
+        _parse_input: &ParseInput,
+    ) -> anyhow::Result<()> {
+        Ok(())
+    }
+}
+
+static_plugin!(PARSE_API = DummyPlugin);
+
+#[cfg(test)]
+mod tests {
+    use falco_plugin_tests::init_plugin;
+
+    #[test]
+    fn test_with_plugin() {
+        init_plugin(super::PARSE_API, c"").unwrap_err();
+    }
+}