diff --git a/crates/cpex-core/examples/cmf_capabilities_demo.rs b/crates/cpex-core/examples/cmf_capabilities_demo.rs index 230c5a36..8843a30e 100644 --- a/crates/cpex-core/examples/cmf_capabilities_demo.rs +++ b/crates/cpex-core/examples/cmf_capabilities_demo.rs @@ -42,7 +42,7 @@ impl Plugin for IdentityChecker { } impl HookHandler for IdentityChecker { - fn handle( + async fn handle( &self, payload: &MessagePayload, extensions: &Extensions, @@ -136,7 +136,7 @@ impl Plugin for HeaderInjector { } impl HookHandler for HeaderInjector { - fn handle( + async fn handle( &self, _payload: &MessagePayload, extensions: &Extensions, @@ -201,7 +201,7 @@ impl Plugin for AuditLogger { } impl HookHandler for AuditLogger { - fn handle( + async fn handle( &self, payload: &MessagePayload, extensions: &Extensions, diff --git a/crates/cpex-core/examples/plugin_demo.rs b/crates/cpex-core/examples/plugin_demo.rs index f0d28f6d..12cfd3f5 100644 --- a/crates/cpex-core/examples/plugin_demo.rs +++ b/crates/cpex-core/examples/plugin_demo.rs @@ -76,7 +76,7 @@ impl Plugin for IdentityResolver { } impl HookHandler for IdentityResolver { - fn handle( + async fn handle( &self, payload: &ToolInvokePayload, _extensions: &Extensions, @@ -98,7 +98,7 @@ impl HookHandler for IdentityResolver { } impl HookHandler for IdentityResolver { - fn handle( + async fn handle( &self, payload: &ToolInvokePayload, _extensions: &Extensions, @@ -126,7 +126,7 @@ impl Plugin for PiiGuard { } impl HookHandler for PiiGuard { - fn handle( + async fn handle( &self, payload: &ToolInvokePayload, _extensions: &Extensions, @@ -171,7 +171,7 @@ impl Plugin for AuditLogger { } impl HookHandler for AuditLogger { - fn handle( + async fn handle( &self, payload: &ToolInvokePayload, _extensions: &Extensions, @@ -186,7 +186,7 @@ impl HookHandler for AuditLogger { } impl HookHandler for AuditLogger { - fn handle( + async fn handle( &self, payload: &ToolInvokePayload, _extensions: &Extensions, @@ -200,6 +200,89 @@ impl HookHandler for AuditLogger { } } +// --------------------------------------------------------------------------- +// Awaiting plugin example — RemoteAuthz +// --------------------------------------------------------------------------- +// +// `HookHandler` is async by design — `handle` is `async fn`. +// Plugins that don't need to `.await` anything still write +// `async fn handle` and return synchronously; this plugin shows the +// other direction, where the body genuinely awaits per-invocation +// work. The realistic version would call a remote authz service +// (gRPC, HTTP, OPA, Cedarling, etc.); here we simulate the network +// round-trip with a small `tokio::time::sleep` so the demo runs +// offline. +// +// Key things this shows: +// 1. Per-request latency state is *cached at init* — the handler +// consults the in-memory ACL and only "calls out" on a miss. +// Hot-path I/O is the most common source of latency regressions +// in plugins, so prefer initialize-time loading wherever you can. +// 2. Registration uses the exact same factory pattern as any other +// plugin — `TypedHandlerAdapter::` and the same +// `register_factory` call. There is no separate async path. +struct RemoteAuthz { + cfg: PluginConfig, + /// ACL "fetched" at init. Populated in Plugin::initialize. + allowed_users: tokio::sync::RwLock>, +} + +#[async_trait] +impl Plugin for RemoteAuthz { + fn config(&self) -> &PluginConfig { + &self.cfg + } + /// Pretend we're loading the ACL from a remote service. In a real + /// plugin this would be `client.fetch_acl().await`; we simulate + /// the round-trip with a small sleep so the demo runs offline. + async fn initialize(&self) -> Result<(), Box> { + tokio::time::sleep(std::time::Duration::from_millis(2)).await; + let mut acl = self.allowed_users.write().await; + acl.extend(["alice", "bob"].iter().map(|s| s.to_string())); + println!( + " [remote-authz] initialized — ACL cached ({} users)", + acl.len() + ); + Ok(()) + } + async fn shutdown(&self) -> Result<(), Box> { + println!(" [remote-authz] shutdown"); + Ok(()) + } +} + +impl HookHandler for RemoteAuthz { + async fn handle( + &self, + payload: &ToolInvokePayload, + _extensions: &Extensions, + _ctx: &mut PluginContext, + ) -> PluginResult { + // Cache hit path — fast. + let acl = self.allowed_users.read().await; + if acl.contains(&payload.user) { + println!( + " [remote-authz] OK (cache hit): user '{}' allowed", + payload.user + ); + return PluginResult::allow(); + } + drop(acl); // release read lock before the fake remote call + // Cache miss path — simulate a remote authz check. In a real + // plugin this is where you'd `.await` a gRPC or HTTP call. + // The latency cost is real and shows up on the request path. + tokio::time::sleep(std::time::Duration::from_millis(1)).await; + println!( + " [remote-authz] DENIED (cache miss + remote check): user '{}'", + payload.user + ); + PluginResult::deny(PluginViolation::new( + "remote_authz_denied", + format!("User '{}' not in remote ACL", payload.user), + )) + } +} + // --------------------------------------------------------------------------- // Step 3: Create plugin factories // --------------------------------------------------------------------------- @@ -264,6 +347,27 @@ impl PluginFactory for AuditLoggerFactory { } } +/// Factory for the async plugin. Note the factory body is identical +/// in shape to the sync factories above — `TypedHandlerAdapter` and +/// the `register_factory` path don't care that the underlying handler +/// is async. The framework hides the choice. +struct RemoteAuthzFactory; +impl PluginFactory for RemoteAuthzFactory { + fn create(&self, config: &PluginConfig) -> Result> { + let plugin = Arc::new(RemoteAuthz { + cfg: config.clone(), + allowed_users: tokio::sync::RwLock::new(std::collections::HashSet::new()), + }); + Ok(PluginInstance { + plugin: plugin.clone(), + handlers: vec![( + "tool_pre_invoke", + Arc::new(TypedHandlerAdapter::::new(plugin)), + )], + }) + } +} + // --------------------------------------------------------------------------- // Step 4: Build extensions with MetaExtension for routing // --------------------------------------------------------------------------- @@ -318,6 +422,7 @@ async fn main() { mgr.register_factory("builtin/identity", Box::new(IdentityFactory)); mgr.register_factory("builtin/pii", Box::new(PiiGuardFactory)); mgr.register_factory("builtin/audit", Box::new(AuditLoggerFactory)); + mgr.register_factory("builtin/remote_authz", Box::new(RemoteAuthzFactory)); mgr.load_config(cpex_config).unwrap(); println!("\n--- Initializing plugins ---\n"); @@ -401,8 +506,38 @@ async fn main() { print_result("some_other_tool (wildcard)", &result); bg.wait_for_background_tasks().await; - // --- Scenario 5: No user identity --- - println!("=== Scenario 5: list_departments (no user identity) ===\n"); + // --- Scenario 5: Awaiting plugin — cache hit --- + // RemoteAuthz's `handle` is `async fn` and reads from a tokio + // RwLock. Its initialize() pre-loaded an ACL containing "alice" + // and "bob"; this call exercises the cache-hit fast path. + println!("=== Scenario 5: query_external_data (async plugin, cache hit) ===\n"); + let payload = ToolInvokePayload { + tool_name: "query_external_data".into(), + user: "alice".into(), + arguments: "dataset=sales".into(), + }; + let ext = make_tool_extensions("query_external_data", &[]); + let (result, bg) = mgr.invoke::(payload, ext, None).await; + print_result("query_external_data (alice — in ACL)", &result); + bg.wait_for_background_tasks().await; + + // --- Scenario 6: Awaiting plugin — cache miss path with .await --- + // "charlie" is not in the cached ACL, so RemoteAuthz takes the + // cache-miss branch and `.await`s a simulated remote call before + // denying. + println!("=== Scenario 6: query_external_data (async plugin, cache miss) ===\n"); + let payload = ToolInvokePayload { + tool_name: "query_external_data".into(), + user: "charlie".into(), + arguments: "dataset=sales".into(), + }; + let ext = make_tool_extensions("query_external_data", &[]); + let (result, bg) = mgr.invoke::(payload, ext, None).await; + print_result("query_external_data (charlie — not in ACL)", &result); + bg.wait_for_background_tasks().await; + + // --- Scenario 7: No user identity --- + println!("=== Scenario 7: list_departments (no user identity) ===\n"); let payload = ToolInvokePayload { tool_name: "list_departments".into(), user: "".into(), diff --git a/crates/cpex-core/examples/plugin_demo.yaml b/crates/cpex-core/examples/plugin_demo.yaml index 9e3dd610..07051e95 100644 --- a/crates/cpex-core/examples/plugin_demo.yaml +++ b/crates/cpex-core/examples/plugin_demo.yaml @@ -15,6 +15,10 @@ global: # "pii" group — activated when a route has the "pii" tag pii: plugins: [pii-guard] + # "external_authz" group — activated when a route has the + # "needs_remote_authz" tag. Fires the async RemoteAuthz plugin. + external_authz: + plugins: [remote-authz] plugins: - name: identity-resolver @@ -33,6 +37,17 @@ plugins: config: clearance_level: confidential + # Awaiting plugin — its `handle` body uses `.await` for a + # (simulated) remote authz call on cache miss. Wired in exactly + # the same way as plugins whose `handle` body has no `.await`; + # registration is identical either way. + - name: remote-authz + kind: builtin/remote_authz + hooks: [tool_pre_invoke] + mode: sequential + priority: 30 + on_error: fail + - name: audit-logger kind: builtin/audit hooks: [tool_pre_invoke, tool_post_invoke] @@ -53,6 +68,16 @@ routes: plugins: - audit-logger + # Tool that requires remote authz — triggers the async plugin + # on top of the standard "all" policy stack. The "external_authz" + # tag matches the policy group of the same name, which fires + # remote-authz. + - tool: query_external_data + meta: + tags: [external_authz] + plugins: + - audit-logger + # Wildcard — catch-all for unmatched tools - tool: "*" plugins: diff --git a/crates/cpex-core/src/hooks/adapter.rs b/crates/cpex-core/src/hooks/adapter.rs index 7acc7b12..985108ca 100644 --- a/crates/cpex-core/src/hooks/adapter.rs +++ b/crates/cpex-core/src/hooks/adapter.rs @@ -3,14 +3,20 @@ // SPDX-License-Identifier: Apache-2.0 // Authors: Teryl Taylor // -// TypedHandlerAdapter — bridges typed HookHandler to type-erased -// AnyHookHandler. +// TypedHandlerAdapter — bridges typed HookHandler to the +// type-erased AnyHookHandler. // // This is framework plumbing that plugin authors never see. When a // plugin is registered via `manager.register_handler::()`, the // manager creates a TypedHandlerAdapter internally. The adapter // translates between Box (what the executor passes) -// and the concrete payload type (what the handler expects). +// and the concrete payload type (what the handler expects), and awaits +// the typed handler's future before re-erasing the result. +// +// `HookHandler` is async-by-default (native AFIT). Plugins that +// don't await anything still write `async fn handle(...)`; the +// compiler emits a trivially-ready future that LLVM inlines, so the +// `.await` here is a no-op for sync-style plugins. use std::marker::PhantomData; use std::sync::Arc; @@ -33,6 +39,11 @@ use crate::registry::AnyHookHandler; /// Created automatically by `PluginManager::register_handler()`. Plugin /// authors never instantiate this directly. /// +/// `HookHandler` is async (native AFIT), so the adapter awaits the +/// returned future before re-erasing the result. Plugins that don't +/// `.await` anything compile to a ready future that LLVM inlines, so +/// they pay no observable cost over a plain function call. +/// /// # Type Parameters /// /// - `H` — the hook type (implements `HookTypeDef`). @@ -73,12 +84,18 @@ where P: Plugin + HookHandler + 'static, { /// Downcast the type-erased payload to the concrete type and call - /// the plugin's typed `handle()` method. + /// the plugin's typed `handle()` method, awaiting the returned + /// future. /// /// The framework retains ownership of the payload — the handler /// receives a borrow (`&H::Payload`) and clones only if it needs /// to modify. The result is erased back to `ErasedResultFields` /// for the executor. + /// + /// For plugins whose body contains no `.await`, the compiler emits + /// a trivially-ready future and LLVM inlines this `.await` to a + /// direct return — there is no observable runtime cost over a + /// plain function call. async fn invoke( &self, payload: &dyn PluginPayload, @@ -97,7 +114,7 @@ where ), })?; - let result = self.plugin.handle(typed_ref, extensions, ctx); + let result = self.plugin.handle(typed_ref, extensions, ctx).await; let plugin_result: PluginResult = result.into(); Ok(erase_result(plugin_result)) diff --git a/crates/cpex-core/src/hooks/trait_def.rs b/crates/cpex-core/src/hooks/trait_def.rs index e07c7ab0..b75d2b8a 100644 --- a/crates/cpex-core/src/hooks/trait_def.rs +++ b/crates/cpex-core/src/hooks/trait_def.rs @@ -78,28 +78,70 @@ pub trait HookTypeDef: Send + Sync + 'static { /// Plugin authors implement this trait (alongside [`Plugin`]) to handle /// a specific hook. The type parameter `H` ties the handler to a /// `HookTypeDef`, ensuring the correct payload and result types at -/// compile time. +/// compile time. The framework creates a type-erased adapter internally +/// when you register — you never touch `AnyHookHandler` directly. /// -/// The framework creates a type-erased adapter internally when you -/// register — you never touch `AnyHookHandler` directly. +/// # Async by design +/// +/// `handle` is an `async fn`. Plugins that don't need to `.await` +/// anything still write `async fn handle(...)` and return synchronously +/// — the compiler emits a trivially-ready future and LLVM inlines it +/// at the adapter site, so there's no observable runtime cost over a +/// plain function. Plugins that *do* need to `.await` (fresh JWKS +/// fetch, RPC to an authz service, dynamic policy lookup) just use +/// `.await` inside the body. +/// +/// **Best practice:** even when async is available, prefer pre-loading +/// state in [`Plugin::initialize`] and reading from cache in `handle`. +/// Hot-path I/O is the most common source of latency regressions. +/// +/// # Native AFIT, not `#[async_trait]` +/// +/// The trait uses native `async fn` (return-position `impl Future`) +/// rather than `#[async_trait]`. This avoids a per-call heap +/// allocation: the returned future is monomorphized into the +/// [`TypedHandlerAdapter`] rather than boxed. The trait is therefore +/// **not object-safe** — you cannot have `Box>`. +/// We don't need that; type erasure happens one layer up at +/// [`AnyHookHandler`]. /// /// # Examples /// /// ```rust,ignore -/// impl HookHandler for MyPlugin { -/// fn handle( +/// // Synchronous plugin — no .await, no extra cost +/// impl HookHandler for AllowPlugin { +/// async fn handle( /// &self, -/// payload: MessagePayload, -/// extensions: &Extensions, -/// ctx: &PluginContext, +/// _payload: &MessagePayload, +/// _extensions: &Extensions, +/// _ctx: &mut PluginContext, /// ) -> PluginResult { /// PluginResult::allow() /// } /// } /// -/// // Registration — no AnyHookHandler needed: -/// manager.register_handler::(plugin, config)?; +/// // Async plugin — calls .await inside the body +/// impl HookHandler for AuthzPlugin { +/// async fn handle( +/// &self, +/// payload: &MyPayload, +/// _extensions: &Extensions, +/// _ctx: &mut PluginContext, +/// ) -> PluginResult { +/// match self.client.check(&payload.user).await { +/// Ok(true) => PluginResult::allow(), +/// _ => PluginResult::deny(/* ... */), +/// } +/// } +/// } +/// +/// // Registration is the same for both: +/// manager.register_handler::(plugin, config)?; /// ``` +/// +/// [`PluginManager::register_handler`]: crate::manager::PluginManager::register_handler +/// [`AnyHookHandler`]: crate::registry::AnyHookHandler +/// [`TypedHandlerAdapter`]: crate::hooks::adapter::TypedHandlerAdapter pub trait HookHandler: Plugin + Send + Sync { /// Handle the hook invocation. /// @@ -112,12 +154,18 @@ pub trait HookHandler: Plugin + Send + Sync { /// the modified copy in `PluginResult::modify_payload()`. This /// pushes the clone cost to the plugin that actually needs it — /// read-only plugins (validators, auditors) never pay for a copy. + /// + /// Returns a `Send`-able future so the executor can drive it from + /// any worker thread (including the concurrent-phase `JoinSet`). + /// `H::Result` is already `Send + Sync` per the `HookTypeDef` + /// bound, so the `Send` constraint comes for free for typical + /// handlers. fn handle( &self, payload: &H::Payload, extensions: &Extensions, ctx: &mut PluginContext, - ) -> H::Result; + ) -> impl std::future::Future + Send; } // --------------------------------------------------------------------------- diff --git a/crates/cpex-core/src/manager.rs b/crates/cpex-core/src/manager.rs index 3ad5977a..16764d49 100644 --- a/crates/cpex-core/src/manager.rs +++ b/crates/cpex-core/src/manager.rs @@ -1280,7 +1280,7 @@ mod tests { } impl HookHandler for AllowPlugin { - fn handle( + async fn handle( &self, _payload: &TestPayload, _extensions: &Extensions, @@ -1309,7 +1309,7 @@ mod tests { } impl HookHandler for DenyPlugin { - fn handle( + async fn handle( &self, _payload: &TestPayload, _extensions: &Extensions, @@ -2068,7 +2068,7 @@ mod tests { } impl HookHandler for TransformPlugin { - fn handle( + async fn handle( &self, payload: &TestPayload, _extensions: &Extensions, @@ -3348,7 +3348,7 @@ routes: } } impl HookHandler for LifecyclePlugin { - fn handle( + async fn handle( &self, _payload: &TestPayload, _extensions: &Extensions, @@ -4089,7 +4089,7 @@ routes: } impl HookHandler for InitTrackingPlugin { - fn handle( + async fn handle( &self, _payload: &TestPayload, _extensions: &Extensions, @@ -4842,4 +4842,117 @@ routes: // With filter_extensions, security IS Some but with empty labels and no subject // So saw_security will be true, but the content is filtered } + + // ----------------------------------------------------------------------- + // Awaiting handler tests + // + // `HookHandler` is async by design. These tests cover handlers + // that genuinely `.await` inside the body — sleeps, yields, and + // co-registration with handlers whose body has no `.await` at all. + // ----------------------------------------------------------------------- + + /// Plugin that genuinely awaits inside its handler. Increments a + /// shared counter after the await resolves so the test can verify + /// the handler ran end-to-end and observed its async point. + struct AsyncCounterPlugin { + cfg: PluginConfig, + counter: Arc, + } + + #[async_trait] + impl Plugin for AsyncCounterPlugin { + fn config(&self) -> &PluginConfig { + &self.cfg + } + } + + impl HookHandler for AsyncCounterPlugin { + async fn handle( + &self, + _payload: &TestPayload, + _extensions: &Extensions, + _ctx: &mut PluginContext, + ) -> PluginResult { + tokio::task::yield_now().await; + tokio::time::sleep(std::time::Duration::from_micros(1)).await; + self.counter + .fetch_add(1, std::sync::atomic::Ordering::SeqCst); + PluginResult::allow() + } + } + + /// Verifies that a handler that genuinely `.await`s gets driven + /// to completion before its result is observed. + #[tokio::test] + async fn test_async_handler_registers_and_invokes() { + let mgr = PluginManager::default(); + let counter = Arc::new(std::sync::atomic::AtomicU64::new(0)); + let cfg = make_config("async-counter", 10, PluginMode::Sequential); + let plugin = Arc::new(AsyncCounterPlugin { + cfg: cfg.clone(), + counter: counter.clone(), + }); + + // Same call path as sync plugins — no `register_async_handler`. + mgr.register_handler::(plugin, cfg).unwrap(); + mgr.initialize().await.unwrap(); + + let payload: Box = Box::new(TestPayload { + value: "test".into(), + }); + let (result, _) = mgr + .invoke_by_name("test_hook", payload, Extensions::default(), None) + .await; + + assert!(result.continue_processing); + assert!(result.violation.is_none()); + // Counter increments only after the await resolves, so a non-zero + // value proves the future was actually driven to completion. + assert_eq!( + counter.load(std::sync::atomic::Ordering::SeqCst), + 1, + "async handler should have run once", + ); + } + + /// A handler with no `.await` (AllowPlugin) and a handler that + /// genuinely awaits (AsyncCounterPlugin) co-register on the same + /// hook via the same `register_handler` call. Both run in priority + /// order. + #[tokio::test] + async fn test_mixed_sync_and_async_handlers_in_same_hook() { + let mgr = PluginManager::default(); + let counter = Arc::new(std::sync::atomic::AtomicU64::new(0)); + + let sync_cfg = make_config("sync-allow", 10, PluginMode::Sequential); + let sync_plugin = Arc::new(AllowPlugin { + cfg: sync_cfg.clone(), + }); + mgr.register_handler::(sync_plugin, sync_cfg) + .unwrap(); + + let async_cfg = make_config("async-counter", 20, PluginMode::Sequential); + let async_plugin = Arc::new(AsyncCounterPlugin { + cfg: async_cfg.clone(), + counter: counter.clone(), + }); + mgr.register_handler::(async_plugin, async_cfg) + .unwrap(); + + mgr.initialize().await.unwrap(); + + let payload: Box = Box::new(TestPayload { + value: "test".into(), + }); + let (result, _) = mgr + .invoke_by_name("test_hook", payload, Extensions::default(), None) + .await; + + assert!(result.continue_processing); + assert_eq!( + counter.load(std::sync::atomic::Ordering::SeqCst), + 1, + "awaiting plugin should have run alongside the non-awaiting plugin", + ); + } } diff --git a/crates/cpex-ffi/src/lib.rs b/crates/cpex-ffi/src/lib.rs index f8bb615f..760f62d9 100644 --- a/crates/cpex-ffi/src/lib.rs +++ b/crates/cpex-ffi/src/lib.rs @@ -1032,7 +1032,7 @@ mod tests { } impl cpex_core::hooks::HookHandler for PanickingPlugin { - fn handle( + async fn handle( &self, _payload: &GenericPayload, _extensions: &Extensions, diff --git a/examples/go-demo/ffi/src/cmf_plugins.rs b/examples/go-demo/ffi/src/cmf_plugins.rs index f59d9e84..a033576f 100644 --- a/examples/go-demo/ffi/src/cmf_plugins.rs +++ b/examples/go-demo/ffi/src/cmf_plugins.rs @@ -68,7 +68,7 @@ impl Plugin for ToolPolicyPlugin { } impl HookHandler for ToolPolicyPlugin { - fn handle( + async fn handle( &self, payload: &MessagePayload, extensions: &Extensions, @@ -197,7 +197,7 @@ impl Plugin for HeaderInjectorPlugin { } impl HookHandler for HeaderInjectorPlugin { - fn handle( + async fn handle( &self, payload: &MessagePayload, extensions: &Extensions, diff --git a/examples/go-demo/ffi/src/demo_plugins.rs b/examples/go-demo/ffi/src/demo_plugins.rs index 84351929..f27125c9 100644 --- a/examples/go-demo/ffi/src/demo_plugins.rs +++ b/examples/go-demo/ffi/src/demo_plugins.rs @@ -61,7 +61,7 @@ impl Plugin for IdentityChecker { } impl HookHandler for IdentityChecker { - fn handle( + async fn handle( &self, payload: &GenericPayload, extensions: &Extensions, @@ -130,7 +130,7 @@ impl Plugin for PiiGuard { } impl HookHandler for PiiGuard { - fn handle( + async fn handle( &self, payload: &GenericPayload, extensions: &Extensions, @@ -212,7 +212,7 @@ impl Plugin for AuditLogger { } impl HookHandler for AuditLogger { - fn handle( + async fn handle( &self, payload: &GenericPayload, extensions: &Extensions,