use std::sync::Arc;
use async_trait::async_trait;
use cpex_core::context::PluginContext;
use cpex_core::error::PluginError;
use cpex_core::extensions::{SubjectExtension, WorkloadIdentity};
use cpex_core::hooks::payload::Extensions;
use cpex_core::hooks::trait_def::{HookHandler, PluginResult};
use cpex_core::identity::{IdentityHook, IdentityPayload, TokenSource, HOOK_IDENTITY_RESOLVE};
use cpex_core::manager::PluginManager;
use cpex_core::plugin::{OnError, Plugin, PluginConfig, PluginMode};
struct SubjectResolver {
cfg: PluginConfig,
subject_id: String,
}
#[async_trait]
impl Plugin for SubjectResolver {
fn config(&self) -> &PluginConfig {
&self.cfg
}
}
impl HookHandler<IdentityHook> for SubjectResolver {
async fn handle(
&self,
payload: &IdentityPayload,
_ext: &Extensions,
_ctx: &mut PluginContext,
) -> PluginResult<IdentityPayload> {
assert!(
!payload.raw_token().is_empty(),
"subject resolver expected a non-empty token",
);
let mut updated = payload.clone();
updated.subject = Some(SubjectExtension {
id: Some(self.subject_id.clone()),
..Default::default()
});
PluginResult::modify_payload(updated)
}
}
struct WorkloadResolver {
cfg: PluginConfig,
require_prior_subject: bool,
}
#[async_trait]
impl Plugin for WorkloadResolver {
fn config(&self) -> &PluginConfig {
&self.cfg
}
}
impl HookHandler<IdentityHook> for WorkloadResolver {
async fn handle(
&self,
payload: &IdentityPayload,
_ext: &Extensions,
_ctx: &mut PluginContext,
) -> PluginResult<IdentityPayload> {
if self.require_prior_subject {
assert!(
payload.subject.is_some(),
"workload resolver expected prior subject in chained run",
);
}
let spiffe_id = payload
.headers()
.get("x-spiffe-id")
.cloned()
.unwrap_or_else(|| "spiffe://example.com/unknown".to_string());
let mut updated = payload.clone();
updated.caller_workload = Some(WorkloadIdentity {
spiffe_id: Some(spiffe_id),
trust_domain: Some("example.com".to_string()),
..Default::default()
});
PluginResult::modify_payload(updated)
}
}
struct RejectingResolver {
cfg: PluginConfig,
}
#[async_trait]
impl Plugin for RejectingResolver {
fn config(&self) -> &PluginConfig {
&self.cfg
}
}
impl HookHandler<IdentityHook> for RejectingResolver {
async fn handle(
&self,
_payload: &IdentityPayload,
_ext: &Extensions,
_ctx: &mut PluginContext,
) -> PluginResult<IdentityPayload> {
PluginResult::deny(cpex_core::error::PluginViolation::new(
"auth.expired",
"token expired",
))
}
}
fn config(name: &str, priority: i32) -> PluginConfig {
PluginConfig {
name: name.to_string(),
kind: "test".to_string(),
description: None,
author: None,
version: None,
hooks: vec![HOOK_IDENTITY_RESOLVE.to_string()],
mode: PluginMode::Sequential,
priority,
on_error: OnError::Fail,
capabilities: Default::default(),
tags: Vec::new(),
conditions: Vec::new(),
config: None,
}
}
fn build_payload(token: &str) -> IdentityPayload {
let mut headers = std::collections::HashMap::new();
headers.insert("authorization".to_string(), format!("Bearer {}", token));
headers.insert(
"x-spiffe-id".to_string(),
"spiffe://example.com/agent-1".to_string(),
);
IdentityPayload::new(token, TokenSource::Bearer)
.with_source_header("Authorization")
.with_headers(headers)
}
fn extract_identity(result: &cpex_core::executor::PipelineResult) -> IdentityPayload {
IdentityPayload::from_pipeline_result(result)
.expect("PipelineResult had no IdentityPayload — denied or wrong hook type")
}
#[tokio::test]
async fn single_resolver_populates_subject() {
let mgr = Arc::new(PluginManager::default());
let cfg = config("subject-resolver", 10);
let plugin = Arc::new(SubjectResolver {
cfg: cfg.clone(),
subject_id: "alice@corp.com".to_string(),
});
mgr.register_handler::<IdentityHook, _>(plugin, cfg)
.unwrap();
mgr.initialize().await.unwrap();
let (result, _bg) = mgr
.invoke_named::<IdentityHook>(
HOOK_IDENTITY_RESOLVE,
build_payload("eyJ.fake.jwt"),
Extensions::default(),
None,
)
.await;
assert!(result.continue_processing, "pipeline should allow");
let final_payload = extract_identity(&result);
assert_eq!(
final_payload.subject.as_ref().unwrap().id.as_deref(),
Some("alice@corp.com"),
);
assert_eq!(final_payload.raw_token(), "eyJ.fake.jwt");
assert_eq!(final_payload.source_header(), Some("Authorization"));
}
#[tokio::test]
async fn two_resolvers_chain_populates_both_slots() {
let mgr = Arc::new(PluginManager::default());
let subject_cfg = config("subject-resolver", 10);
let subject = Arc::new(SubjectResolver {
cfg: subject_cfg.clone(),
subject_id: "alice@corp.com".to_string(),
});
mgr.register_handler::<IdentityHook, _>(subject, subject_cfg)
.unwrap();
let workload_cfg = config("workload-resolver", 20); let workload = Arc::new(WorkloadResolver {
cfg: workload_cfg.clone(),
require_prior_subject: true,
});
mgr.register_handler::<IdentityHook, _>(workload, workload_cfg)
.unwrap();
mgr.initialize().await.unwrap();
let (result, _bg) = mgr
.invoke_named::<IdentityHook>(
HOOK_IDENTITY_RESOLVE,
build_payload("eyJ.fake.jwt"),
Extensions::default(),
None,
)
.await;
assert!(result.continue_processing, "pipeline should allow");
let final_payload = extract_identity(&result);
assert_eq!(
final_payload.subject.as_ref().unwrap().id.as_deref(),
Some("alice@corp.com"),
);
let workload = final_payload
.caller_workload
.as_ref()
.expect("workload resolver should have populated caller_workload");
assert_eq!(
workload.spiffe_id.as_deref(),
Some("spiffe://example.com/agent-1"),
);
assert_eq!(final_payload.raw_token(), "eyJ.fake.jwt");
}
#[tokio::test]
async fn rejecting_resolver_halts_pipeline() {
let mgr = Arc::new(PluginManager::default());
let cfg = config("rejecting-resolver", 10);
let plugin = Arc::new(RejectingResolver { cfg: cfg.clone() });
mgr.register_handler::<IdentityHook, _>(plugin, cfg)
.unwrap();
mgr.initialize().await.unwrap();
let (result, _bg) = mgr
.invoke_named::<IdentityHook>(
HOOK_IDENTITY_RESOLVE,
build_payload("eyJ.expired.jwt"),
Extensions::default(),
None,
)
.await;
assert!(!result.continue_processing, "rejection should halt");
let violation = result.violation.expect("rejected → violation present");
assert_eq!(violation.code, "auth.expired");
assert_eq!(violation.reason, "token expired");
}
#[tokio::test]
async fn apply_to_extensions_populates_security_and_preserves_existing_fields() {
use cpex_core::extensions::raw_credentials::{
RawCredentialsExtension, RawInboundToken, TokenKind, TokenRole,
};
use cpex_core::extensions::SecurityExtension;
struct FullResolver {
cfg: PluginConfig,
}
#[async_trait]
impl Plugin for FullResolver {
fn config(&self) -> &PluginConfig {
&self.cfg
}
}
impl HookHandler<IdentityHook> for FullResolver {
async fn handle(
&self,
payload: &IdentityPayload,
_ext: &Extensions,
_ctx: &mut PluginContext,
) -> PluginResult<IdentityPayload> {
let token_bytes = payload.raw_token().to_string();
let mut updated = payload.clone();
updated.subject = Some(SubjectExtension {
id: Some("alice@corp.com".into()),
..Default::default()
});
let mut raw = RawCredentialsExtension::default();
raw.inbound_tokens.insert(
TokenRole::User,
RawInboundToken::new(token_bytes, "Authorization", TokenKind::Jwt),
);
updated.raw_credentials = Some(raw);
PluginResult::modify_payload(updated)
}
}
let mgr = Arc::new(PluginManager::default());
let cfg = config("full-resolver", 10);
let plugin = Arc::new(FullResolver { cfg: cfg.clone() });
mgr.register_handler::<IdentityHook, _>(plugin, cfg)
.unwrap();
mgr.initialize().await.unwrap();
let mut initial_security = SecurityExtension::default();
initial_security.add_label("PII");
initial_security.classification = Some("internal".into());
let initial_ext = Extensions {
security: Some(Arc::new(initial_security)),
..Default::default()
};
let (result, _bg) = mgr
.invoke_named::<IdentityHook>(
HOOK_IDENTITY_RESOLVE,
build_payload("eyJ.fake.jwt"),
initial_ext.clone(),
None,
)
.await;
assert!(result.continue_processing);
let final_payload = extract_identity(&result);
let updated_ext = final_payload.apply_to_extensions(initial_ext);
let sec = updated_ext
.security
.as_ref()
.expect("security slot present");
assert_eq!(
sec.subject.as_ref().unwrap().id.as_deref(),
Some("alice@corp.com"),
);
assert!(sec.has_label("PII"), "pre-existing label survived apply");
assert_eq!(sec.classification.as_deref(), Some("internal"));
let raw = updated_ext
.raw_credentials
.as_ref()
.expect("raw_credentials slot present");
let user_token = raw
.inbound_tokens
.get(&TokenRole::User)
.expect("user token present");
assert_eq!(user_token.source_header, "Authorization");
assert_eq!(&*user_token.token, "eyJ.fake.jwt");
}
#[tokio::test]
async fn from_pipeline_result_returns_none_on_deny() {
let mgr = Arc::new(PluginManager::default());
let cfg = config("rejecter", 10);
let plugin = Arc::new(RejectingResolver { cfg: cfg.clone() });
mgr.register_handler::<IdentityHook, _>(plugin, cfg)
.unwrap();
mgr.initialize().await.unwrap();
let (result, _bg) = mgr
.invoke_named::<IdentityHook>(
HOOK_IDENTITY_RESOLVE,
build_payload("eyJ.tok"),
Extensions::default(),
None,
)
.await;
assert!(!result.continue_processing);
assert!(IdentityPayload::from_pipeline_result(&result).is_none());
}
#[tokio::test]
async fn cap_gating_post_apply_through_cmf_dispatch() {
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::sync::Mutex;
use cpex_core::cmf::enums::Role;
use cpex_core::cmf::{CmfHook, Message, MessagePayload};
use cpex_core::extensions::raw_credentials::{
RawCredentialsExtension, RawInboundToken, TokenKind, TokenRole,
};
use cpex_core::extensions::SecurityExtension;
struct FullResolver {
cfg: PluginConfig,
}
#[async_trait]
impl Plugin for FullResolver {
fn config(&self) -> &PluginConfig {
&self.cfg
}
}
impl HookHandler<IdentityHook> for FullResolver {
async fn handle(
&self,
payload: &IdentityPayload,
_ext: &Extensions,
_ctx: &mut PluginContext,
) -> PluginResult<IdentityPayload> {
let token = payload.raw_token().to_string();
let mut updated = payload.clone();
updated.subject = Some(SubjectExtension {
id: Some("alice@corp.com".into()),
..Default::default()
});
let mut raw = RawCredentialsExtension::default();
raw.inbound_tokens.insert(
TokenRole::User,
RawInboundToken::new(token, "Authorization", TokenKind::Jwt),
);
updated.raw_credentials = Some(raw);
PluginResult::modify_payload(updated)
}
}
struct InboundReader {
cfg: PluginConfig,
saw_token_count: Arc<AtomicUsize>,
saw_subject_id: Arc<Mutex<Option<String>>>,
}
#[async_trait]
impl Plugin for InboundReader {
fn config(&self) -> &PluginConfig {
&self.cfg
}
}
impl HookHandler<CmfHook> for InboundReader {
async fn handle(
&self,
_payload: &MessagePayload,
ext: &Extensions,
_ctx: &mut PluginContext,
) -> PluginResult<MessagePayload> {
let n = ext
.raw_credentials
.as_ref()
.map(|r| r.inbound_tokens.len())
.unwrap_or(0);
self.saw_token_count.store(n, Ordering::SeqCst);
let id = ext
.security
.as_ref()
.and_then(|s| s.subject.as_ref())
.and_then(|s| s.id.clone());
*self.saw_subject_id.lock().unwrap() = id;
PluginResult::allow()
}
}
struct InboundBlind {
cfg: PluginConfig,
saw_any_credentials: Arc<AtomicBool>,
}
#[async_trait]
impl Plugin for InboundBlind {
fn config(&self) -> &PluginConfig {
&self.cfg
}
}
impl HookHandler<CmfHook> for InboundBlind {
async fn handle(
&self,
_payload: &MessagePayload,
ext: &Extensions,
_ctx: &mut PluginContext,
) -> PluginResult<MessagePayload> {
self.saw_any_credentials
.store(ext.raw_credentials.is_some(), Ordering::SeqCst);
PluginResult::allow()
}
}
let mgr = Arc::new(PluginManager::default());
let id_cfg = config("full-resolver", 10);
mgr.register_handler::<IdentityHook, _>(
Arc::new(FullResolver {
cfg: id_cfg.clone(),
}),
id_cfg,
)
.unwrap();
let reader_saw_count = Arc::new(AtomicUsize::new(usize::MAX)); let reader_saw_subject = Arc::new(Mutex::new(None));
let reader_cfg = PluginConfig {
name: "inbound-reader".into(),
kind: "test".into(),
description: None,
author: None,
version: None,
hooks: vec!["cmf.tool_pre_invoke".into()],
mode: PluginMode::Sequential,
priority: 10,
on_error: OnError::Fail,
capabilities: ["read_inbound_credentials", "read_subject"]
.iter()
.map(|s| s.to_string())
.collect(),
tags: Vec::new(),
conditions: Vec::new(),
config: None,
};
mgr.register_handler_for_names::<CmfHook, _>(
Arc::new(InboundReader {
cfg: reader_cfg.clone(),
saw_token_count: Arc::clone(&reader_saw_count),
saw_subject_id: Arc::clone(&reader_saw_subject),
}),
reader_cfg,
&["cmf.tool_pre_invoke"],
)
.unwrap();
let blind_saw_creds = Arc::new(AtomicBool::new(false));
let blind_cfg = PluginConfig {
name: "inbound-blind".into(),
kind: "test".into(),
description: None,
author: None,
version: None,
hooks: vec!["cmf.tool_pre_invoke".into()],
mode: PluginMode::Sequential,
priority: 20,
on_error: OnError::Fail,
capabilities: Default::default(), tags: Vec::new(),
conditions: Vec::new(),
config: None,
};
mgr.register_handler_for_names::<CmfHook, _>(
Arc::new(InboundBlind {
cfg: blind_cfg.clone(),
saw_any_credentials: Arc::clone(&blind_saw_creds),
}),
blind_cfg,
&["cmf.tool_pre_invoke"],
)
.unwrap();
mgr.initialize().await.unwrap();
let mut initial_security = SecurityExtension::default();
initial_security.add_label("PII");
let initial_ext = Extensions {
security: Some(Arc::new(initial_security)),
..Default::default()
};
let (id_result, _bg) = mgr
.invoke_named::<IdentityHook>(
HOOK_IDENTITY_RESOLVE,
build_payload("eyJ.fake.jwt"),
initial_ext.clone(),
None,
)
.await;
assert!(id_result.continue_processing);
let identity =
IdentityPayload::from_pipeline_result(&id_result).expect("identity should have resolved");
let updated_ext = identity.apply_to_extensions(initial_ext);
let cmf_payload = MessagePayload {
message: Message::text(Role::User, "fetch sensitive data"),
};
let (cmf_result, _bg) = mgr
.invoke_named::<CmfHook>("cmf.tool_pre_invoke", cmf_payload, updated_ext, None)
.await;
assert!(
cmf_result.continue_processing,
"CMF dispatch should not be blocked: violation = {:?}",
cmf_result.violation,
);
assert_eq!(
reader_saw_count.load(Ordering::SeqCst),
1,
"InboundReader with read_inbound_credentials should see 1 token",
);
assert_eq!(
reader_saw_subject.lock().unwrap().as_deref(),
Some("alice@corp.com"),
);
assert!(
!blind_saw_creds.load(Ordering::SeqCst),
"InboundBlind without credential caps must NOT see raw_credentials",
);
}
#[allow(dead_code)]
fn _force_plugin_error_link(_e: PluginError) {}