Skip to content

Commit

Permalink
new(plugin): simplify extractor signature even more
Browse files Browse the repository at this point in the history
You no longer need to explicitly specify the argument type,
it's inferred from the extractor signature now.

Signed-off-by: Grzegorz Nosek <grzegorz.nosek@sysdig.com>
  • Loading branch information
gnosek authored and poiana committed Oct 21, 2024
1 parent 07b3340 commit 16c0a8f
Show file tree
Hide file tree
Showing 10 changed files with 265 additions and 288 deletions.
5 changes: 1 addition & 4 deletions falco_plugin/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -60,7 +60,6 @@ pub mod base {
/// use falco_plugin::extract::{
/// EventInput,
/// ExtractFieldInfo,
/// ExtractFieldRequestArg,
/// ExtractPlugin,
/// ExtractRequest,
/// field};
Expand Down Expand Up @@ -93,7 +92,6 @@ pub mod base {
/// fn extract_sample(
/// &mut self,
/// _req: ExtractRequest<Self>,
/// _arg: ExtractFieldRequestArg,
/// ) -> Result<CString, Error> {
/// Ok(c"hello".to_owned())
/// }
Expand All @@ -117,8 +115,7 @@ pub mod base {
pub mod extract {
pub use crate::plugin::event::EventInput;
pub use crate::plugin::extract::schema::field;
pub use crate::plugin::extract::schema::{ExtractArgType, ExtractFieldInfo};
pub use crate::plugin::extract::ExtractFieldRequestArg;
pub use crate::plugin::extract::schema::ExtractFieldInfo;
pub use crate::plugin::extract::ExtractPlugin;
pub use crate::plugin::extract::ExtractRequest;
}
Expand Down
191 changes: 191 additions & 0 deletions falco_plugin/src/plugin/extract/extractor_fn.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,191 @@
use crate::extract::{ExtractPlugin, ExtractRequest};
use crate::plugin::extract::fields::Extract;
use crate::plugin::extract::schema::ExtractArgType;
use crate::plugin::extract::{ExtractField, ExtractFieldRequestArg};
use anyhow::Error;
use falco_plugin_api::ss_plugin_extract_field;
use std::ffi::CStr;

#[derive(Debug)]
pub struct ExtractLambda<P: ExtractPlugin> {
pub(in crate::plugin::extract) obj: *const (),

#[allow(clippy::type_complexity)]
pub(in crate::plugin::extract) func: fn(
obj: *const (),
plugin: &mut P,
field: &mut ss_plugin_extract_field,
request: ExtractRequest<'_, '_, '_, P>,
storage: &mut bumpalo::Bump,
) -> Result<(), Error>,
}

impl<P: ExtractPlugin> ExtractLambda<P> {
pub(in crate::plugin::extract) fn call(
&self,
plugin: &mut P,
field: &mut ss_plugin_extract_field,
request: ExtractRequest<'_, '_, '_, P>,
storage: &mut bumpalo::Bump,
) -> Result<(), Error> {
(self.func)(self.obj, plugin, field, request, storage)
}
}

#[derive(Debug)]
pub struct NoArg;

#[derive(Debug)]
pub struct IntArg;

#[derive(Debug)]
pub struct StringArg;

#[derive(Debug)]
pub struct OptIntArg;

#[derive(Debug)]
pub struct OptStringArg;

pub trait ExtractorFn<P, R, A>
where
P: ExtractPlugin,
R: Extract,
{
const ARG_TYPE: ExtractArgType;

fn call(
obj: *const (),
plugin: &mut P,
req: ExtractRequest<P>,
arg: ExtractFieldRequestArg,
) -> Result<R, Error>;

fn extract<'a>(
obj: *const (),
plugin: &'a mut P,
field: &mut ss_plugin_extract_field,
request: ExtractRequest<'a, '_, '_, P>,
storage: &mut bumpalo::Bump,
) -> Result<(), Error> {
let result = Self::call(obj, plugin, request, unsafe { field.key_unchecked() })?;
Ok(result.extract_to(field, storage)?)
}
}

impl<P, R, F> ExtractorFn<P, R, NoArg> for F
where
P: ExtractPlugin,
R: Extract,
F: Fn(&mut P, ExtractRequest<P>) -> Result<R, Error> + 'static,
{
const ARG_TYPE: ExtractArgType = ExtractArgType::None;

fn call(
obj: *const (),
plugin: &mut P,
req: ExtractRequest<P>,
arg: ExtractFieldRequestArg,
) -> Result<R, Error> {
anyhow::ensure!(matches!(arg, ExtractFieldRequestArg::None));

let func = obj as *const F;
unsafe { (*func)(plugin, req) }
}
}

impl<P, R, F> ExtractorFn<P, R, IntArg> for F
where
P: ExtractPlugin,
R: Extract,
F: Fn(&mut P, ExtractRequest<P>, u64) -> Result<R, Error> + 'static,
{
const ARG_TYPE: ExtractArgType = ExtractArgType::RequiredIndex;

fn call(
obj: *const (),
plugin: &mut P,
req: ExtractRequest<P>,
arg: ExtractFieldRequestArg,
) -> Result<R, Error> {
let ExtractFieldRequestArg::Int(arg) = arg else {
anyhow::bail!("Expected index argument, got {:?}", arg);
};

let func = obj as *const F;
unsafe { (*func)(plugin, req, arg) }
}
}

impl<P, R, F> ExtractorFn<P, R, OptIntArg> for F
where
P: ExtractPlugin,
R: Extract,
F: Fn(&mut P, ExtractRequest<P>, Option<u64>) -> Result<R, Error> + 'static,
{
const ARG_TYPE: ExtractArgType = ExtractArgType::OptionalIndex;

fn call(
obj: *const (),
plugin: &mut P,
req: ExtractRequest<P>,
arg: ExtractFieldRequestArg,
) -> Result<R, Error> {
let arg = match arg {
ExtractFieldRequestArg::Int(arg) => Some(arg),
ExtractFieldRequestArg::None => None,
_ => anyhow::bail!("Expected index argument, got {:?}", arg),
};

let func = obj as *const F;
unsafe { (*func)(plugin, req, arg) }
}
}

impl<P, R, F> ExtractorFn<P, R, StringArg> for F
where
P: ExtractPlugin,
R: Extract,
F: Fn(&mut P, ExtractRequest<P>, &CStr) -> Result<R, Error> + 'static,
{
const ARG_TYPE: ExtractArgType = ExtractArgType::RequiredKey;

fn call(
obj: *const (),
plugin: &mut P,
req: ExtractRequest<P>,
arg: ExtractFieldRequestArg,
) -> Result<R, Error> {
let ExtractFieldRequestArg::String(arg) = arg else {
anyhow::bail!("Expected key argument, got {:?}", arg);
};

let func = obj as *const F;
unsafe { (*func)(plugin, req, arg) }
}
}

impl<P, R, F> ExtractorFn<P, R, OptStringArg> for F
where
P: ExtractPlugin,
R: Extract,
F: Fn(&mut P, ExtractRequest<P>, Option<&CStr>) -> Result<R, Error> + 'static,
{
const ARG_TYPE: ExtractArgType = ExtractArgType::OptionalKey;

fn call(
obj: *const (),
plugin: &mut P,
req: ExtractRequest<P>,
arg: ExtractFieldRequestArg,
) -> Result<R, Error> {
let arg = match arg {
ExtractFieldRequestArg::String(arg) => Some(arg),
ExtractFieldRequestArg::None => None,
_ => anyhow::bail!("Expected key argument, got {:?}", arg),
};

let func = obj as *const F;
unsafe { (*func)(plugin, req, arg) }
}
}
81 changes: 18 additions & 63 deletions falco_plugin/src/plugin/extract/mod.rs
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
use crate::extract::{EventInput, ExtractArgType};
use crate::extract::EventInput;
use crate::plugin::base::Plugin;
use crate::plugin::extract::schema::ExtractFieldInfo;
use crate::tables::TableReader;
Expand All @@ -8,8 +8,8 @@ use std::any::TypeId;
use std::collections::BTreeMap;
use std::ffi::{CStr, CString};
use std::sync::Mutex;
use thiserror::Error;

mod extractor_fn;
pub mod fields;
pub mod schema;
#[doc(hidden)]
Expand All @@ -32,50 +32,8 @@ pub enum ExtractFieldRequestArg<'a> {
String(&'a CStr),
}

#[derive(Debug, Error)]
pub enum ArgError {
#[error("required argument missing")]
Missing,

#[error("unexpected argument")]
Unexpected,

#[error("expected string argument")]
ExpectedString,

#[error("expected int argument")]
ExpectedInt,
}

pub trait ExtractField {
unsafe fn key_unchecked(&self) -> ExtractFieldRequestArg;

unsafe fn key(&self, arg_type: ExtractArgType) -> Result<ExtractFieldRequestArg, ArgError> {
let key = unsafe { self.key_unchecked() };
match key {
k @ ExtractFieldRequestArg::None => match arg_type {
ExtractArgType::None => Ok(k),
ExtractArgType::OptionalIndex => Ok(k),
ExtractArgType::OptionalKey => Ok(k),
ExtractArgType::RequiredIndex => Err(ArgError::Missing),
ExtractArgType::RequiredKey => Err(ArgError::Missing),
},
k @ ExtractFieldRequestArg::Int(_) => match arg_type {
ExtractArgType::None => Err(ArgError::Unexpected),
ExtractArgType::OptionalIndex => Ok(k),
ExtractArgType::OptionalKey => Err(ArgError::ExpectedString),
ExtractArgType::RequiredIndex => Ok(k),
ExtractArgType::RequiredKey => Err(ArgError::ExpectedString),
},
k @ ExtractFieldRequestArg::String(_) => match arg_type {
ExtractArgType::None => Err(ArgError::Unexpected),
ExtractArgType::OptionalIndex => Err(ArgError::ExpectedInt),
ExtractArgType::OptionalKey => Ok(k),
ExtractArgType::RequiredIndex => Err(ArgError::ExpectedInt),
ExtractArgType::RequiredKey => Ok(k),
},
}
}
}

impl ExtractField for ss_plugin_extract_field {
Expand Down Expand Up @@ -152,8 +110,7 @@ where
///
/// fn extract_field_one(
/// &mut self,
/// req: ExtractContext<Self>,
/// arg: ExtractRequestArg) -> ... {
/// req: ExtractContext<Self>) -> ... {
/// let context = req.context.get_or_insert_with(|| self.make_context(...));
///
/// // use context
Expand All @@ -173,7 +130,7 @@ where
/// fn extract_sample(
/// &mut self,
/// req: ExtractRequest<Self>,
/// arg: ExtractFieldRequestArg,
/// arg: A, // optional
/// ) -> Result<R, Error>;
///
/// ```
Expand All @@ -186,11 +143,19 @@ where
/// - [`std::net::IpAddr`]
/// - [`falco_event::fields::types::PT_IPNET`]
///
/// and `A` is the argument to the field extraction:
///
/// | Argument declaration | `field` lookup | `field[5]` lookup | `field[foo]` lookup |
/// |----------------------|----------------|-------------------|---------------------|
/// | _missing_ | valid | - | - |
/// | `arg: u64` | - | valid | - |
/// | `arg: Option<u64>` | valid | valid | - |
/// | `arg: &CStr` | - | - | valid |
/// | `arg: Option<&CStr>` | valid | - | valid |
///
/// `req` is the extraction request ([`ExtractRequest`]), containing the context in which
/// the plugin is doing the work.
///
/// `arg` is the actual argument passed along with the field (see [`ExtractFieldRequestArg`])
///
/// To register extracted fields, add them to the [`ExtractPlugin::EXTRACT_FIELDS`] array, wrapped via [`crate::extract::field`]:
/// ```
/// use std::ffi::CStr;
Expand All @@ -200,9 +165,7 @@ where
/// use falco_plugin::base::Plugin;
/// use falco_plugin::extract::{
/// field,
/// ExtractArgType,
/// ExtractFieldInfo,
/// ExtractFieldRequestArg,
/// ExtractPlugin,
/// ExtractRequest};
/// use falco_plugin::tables::TablesInput;
Expand All @@ -225,20 +188,16 @@ where
/// fn extract_sample(
/// &mut self,
/// _req: ExtractRequest<Self>,
/// _arg: ExtractFieldRequestArg,
/// ) -> Result<u64, Error> {
/// Ok(10u64)
/// }
///
/// fn extract_arg(
/// &mut self,
/// _req: ExtractRequest<Self>,
/// arg: ExtractFieldRequestArg,
/// arg: u64,
/// ) -> Result<u64, Error> {
/// match arg {
/// ExtractFieldRequestArg::Int(i) => Ok(i),
/// _ => anyhow::bail!("wanted an int argument, got {:?}", arg)
/// }
/// Ok(arg)
/// }
/// }
///
Expand All @@ -248,15 +207,11 @@ where
/// type ExtractContext = ();
/// const EXTRACT_FIELDS: &'static [ExtractFieldInfo<Self>] = &[
/// field("sample.always_10", &Self::extract_sample),
/// field("sample.arg", &Self::extract_arg).with_arg(ExtractArgType::RequiredIndex),
/// field("sample.arg", &Self::extract_arg)
/// ];
/// }
///
/// ```
///
/// **Note**: while the returned field type is automatically determined based on the return type
/// of the function, the argument type defaults to [`ExtractArgType::None`] and must be explicitly specified
/// using [`ExtractFieldInfo::with_arg`] if the function expects an argument.
const EXTRACT_FIELDS: &'static [ExtractFieldInfo<Self>];

/// Generate the field schema for the Falco plugin framework
Expand Down Expand Up @@ -317,7 +272,7 @@ where
table_reader,
};

info.func.extract(self, req, request, info.arg, storage)?;
info.func.call(self, req, request, storage)?;
}
Ok(())
}
Expand Down
Loading

0 comments on commit 16c0a8f

Please sign in to comment.