From b847143bd9573b327a940cbb5c4997420a0d2a05 Mon Sep 17 00:00:00 2001 From: MuZhou233 Date: Mon, 2 Sep 2024 17:54:17 +0800 Subject: [PATCH] fix: pass test --- src/core_api.rs | 4 +- src/enforcer.rs | 32 ++++++++----- src/model/function_map.rs | 97 ++++++++++----------------------------- 3 files changed, 46 insertions(+), 87 deletions(-) diff --git a/src/core_api.rs b/src/core_api.rs index ca47c186..cd78040b 100644 --- a/src/core_api.rs +++ b/src/core_api.rs @@ -1,4 +1,5 @@ use crate::{ + model::OperatorFunction, enforcer::EnforceContext, Adapter, Effector, EnforceArgs, Event, EventEmitter, Filter, Model, Result, RoleManager, TryIntoAdapter, TryIntoModel, @@ -15,7 +16,6 @@ use crate::emitter::EventData; use async_trait::async_trait; use parking_lot::RwLock; -use rhai::{Dynamic, ImmutableString}; use std::sync::Arc; @@ -36,7 +36,7 @@ pub trait CoreApi: Send + Sync { fn add_function( &mut self, fname: &str, - f: fn(&[ImmutableString]) -> Dynamic, + f: OperatorFunction, ); fn get_model(&self) -> &dyn Model; fn get_mut_model(&mut self) -> &mut dyn Model; diff --git a/src/enforcer.rs b/src/enforcer.rs index 6dac1ce9..2e94137d 100644 --- a/src/enforcer.rs +++ b/src/enforcer.rs @@ -7,7 +7,7 @@ use crate::{ error::{ModelError, PolicyError, RequestError}, get_or_err, get_or_err_with_context, management_api::MgmtApi, - model::{FunctionMap, Model}, + model::{FunctionMap, Model, OperatorFunction}, rbac::{DefaultRoleManager, RoleManager}, register_g_function, util::{escape_assertion, escape_eval}, @@ -350,6 +350,19 @@ impl Enforcer { } })) } + + fn register_function( + engine: &mut Engine, + key: &str, + f: OperatorFunction, + ) { + match f { + OperatorFunction::Arg0(func) => {engine.register_fn(key, func);} + OperatorFunction::Arg1(func) => {engine.register_fn(key, func);} + OperatorFunction::Arg2(func) => {engine.register_fn(key, func);} + OperatorFunction::Arg3(func) => {engine.register_fn(key, func);} + } + } pub(crate) fn register_g_functions(&mut self) -> Result<()> { if let Some(ast_map) = self.model.get_model().get("g") { @@ -380,7 +393,7 @@ impl CoreApi for Enforcer { engine.register_global_module(CASBIN_PACKAGE.as_shared_module()); for (key, &func) in fm.get_functions() { - engine.register_fn(key, func); + Self::register_function(&mut engine, key, func); } let mut e = Self { @@ -428,10 +441,10 @@ impl CoreApi for Enforcer { fn add_function( &mut self, fname: &str, - f: fn(&[ImmutableString]) -> Dynamic, + f: OperatorFunction, ) { self.fm.add_function(fname, f); - self.engine.register_fn(fname, f); + Self::register_function(&mut self.engine, fname, f); } #[inline] @@ -1338,13 +1351,10 @@ mod tests { let adapter1 = FileAdapter::new("examples/keymatch_policy.csv"); let mut e = Enforcer::new(m1, adapter1).await.unwrap(); - e.add_function("keyMatchCustom", |args: &[ImmutableString]| { - key_match( - &args.get(0).unwrap_or(&ImmutableString::from("")), - &args.get(1).unwrap_or(&ImmutableString::from("")), - ) - .into() - }); + e.add_function( + "keyMatchCustom", + OperatorFunction::Arg2(|s1: ImmutableString, s2: ImmutableString| key_match(&s1, &s2).into()), + ); assert_eq!( true, diff --git a/src/model/function_map.rs b/src/model/function_map.rs index ef2a5654..3321f1b7 100644 --- a/src/model/function_map.rs +++ b/src/model/function_map.rs @@ -15,86 +15,37 @@ use rhai::{Dynamic, ImmutableString}; static MAT_B: Lazy = Lazy::new(|| Regex::new(r":[^/]*").unwrap()); static MAT_P: Lazy = Lazy::new(|| Regex::new(r"\{[^/]*\}").unwrap()); -use std::any::Any; use std::{borrow::Cow, collections::HashMap}; +#[derive(Clone, Copy)] +pub enum OperatorFunction { + Arg0(fn() -> Dynamic), + Arg1(fn(ImmutableString) -> Dynamic), + Arg2(fn(ImmutableString, ImmutableString) -> Dynamic), + Arg3(fn(ImmutableString, ImmutableString, ImmutableString) -> Dynamic), +} + pub struct FunctionMap { - pub(crate) fm: HashMap Dynamic>, + pub(crate) fm: HashMap, } impl Default for FunctionMap { fn default() -> FunctionMap { - let mut fm: HashMap Dynamic> = + let mut fm: HashMap = HashMap::new(); - fm.insert("keyMatch".to_owned(), |args: &[ImmutableString]| { - key_match( - args.get(0).unwrap_or(&ImmutableString::from("")), - args.get(1).unwrap_or(&ImmutableString::from("")), - ) - .into() - }); - fm.insert("keyGet".to_owned(), |args: &[ImmutableString]| { - key_get( - args.get(0).unwrap_or(&ImmutableString::from("")), - args.get(1).unwrap_or(&ImmutableString::from("")), - ) - .into() - }); - fm.insert("keyMatch2".to_owned(), |args: &[ImmutableString]| { - key_match2( - args.get(0).unwrap_or(&ImmutableString::from("")), - args.get(1).unwrap_or(&ImmutableString::from("")), - ) - .into() - }); - fm.insert("keyGet2".to_owned(), |args: &[ImmutableString]| { - key_get2( - args.get(0).unwrap_or(&ImmutableString::from("")), - args.get(1).unwrap_or(&ImmutableString::from("")), - args.get(2).unwrap_or(&ImmutableString::from("")), - ) - .into() - }); - fm.insert("keyMatch3".to_owned(), |args: &[ImmutableString]| { - key_match3( - args.get(0).unwrap_or(&ImmutableString::from("")), - args.get(1).unwrap_or(&ImmutableString::from("")), - ) - .into() - }); - fm.insert("keyGet3".to_owned(), |args: &[ImmutableString]| { - key_get3( - args.get(0).unwrap_or(&ImmutableString::from("")), - args.get(1).unwrap_or(&ImmutableString::from("")), - args.get(2).unwrap_or(&ImmutableString::from("")), - ) - .into() - }); - fm.insert("regexMatch".to_owned(), |args: &[ImmutableString]| { - regex_match( - args.get(0).unwrap_or(&ImmutableString::from("")), - args.get(1).unwrap_or(&ImmutableString::from("")), - ) - .into() - }); + fm.insert("keyMatch".to_owned(), OperatorFunction::Arg2(|s1: ImmutableString, s2: ImmutableString| { key_match(&s1, &s2).into() })); + fm.insert("keyGet".to_owned(), OperatorFunction::Arg2(|s1: ImmutableString, s2: ImmutableString| { key_get(&s1, &s2).into() })); + fm.insert("keyMatch2".to_owned(), OperatorFunction::Arg2(|s1: ImmutableString, s2: ImmutableString| { key_match2(&s1, &s2).into() })); + fm.insert("keyGet2".to_owned(), OperatorFunction::Arg3(|s1: ImmutableString, s2: ImmutableString, s3: ImmutableString| { key_get2(&s1, &s2, &s3).into() })); + fm.insert("keyMatch3".to_owned(), OperatorFunction::Arg2(|s1: ImmutableString, s2: ImmutableString| { key_match3(&s1, &s2).into() })); + fm.insert("keyGet3".to_owned(), OperatorFunction::Arg3(|s1: ImmutableString, s2: ImmutableString, s3: ImmutableString| { key_get3(&s1, &s2, &s3).into() })); + fm.insert("regexMatch".to_owned(), OperatorFunction::Arg2(|s1: ImmutableString, s2: ImmutableString| { regex_match(&s1, &s2).into() })); #[cfg(feature = "glob")] - fm.insert("globMatch".to_owned(), |args: &[ImmutableString]| { - glob_match( - args.get(0).unwrap_or(&ImmutableString::from("")), - args.get(1).unwrap_or(&ImmutableString::from("")), - ) - .into() - }); + fm.insert("globMatch".to_owned(), OperatorFunction::Arg2(|s1: ImmutableString, s2: ImmutableString| { glob_match(&s1, &s2).into() })); #[cfg(feature = "ip")] - fm.insert("ipMatch".to_owned(), |args: &[ImmutableString]| { - ip_match( - args.get(0).unwrap_or(&ImmutableString::from("")), - args.get(1).unwrap_or(&ImmutableString::from("")), - ) - .into() - }); + fm.insert("ipMatch".to_owned(), OperatorFunction::Arg2(|s1: ImmutableString, s2: ImmutableString| { ip_match(&s1, &s2).into() })); FunctionMap { fm } } @@ -105,7 +56,7 @@ impl FunctionMap { pub fn add_function( &mut self, fname: &str, - f: fn(&[ImmutableString]) -> Dynamic, + f: OperatorFunction, ) { self.fm.insert(fname.to_owned(), f); } @@ -113,7 +64,7 @@ impl FunctionMap { #[inline] pub fn get_functions( &self, - ) -> impl Iterator Dynamic)> + ) -> impl Iterator { self.fm.iter() } @@ -137,10 +88,8 @@ pub fn key_match(key1: &str, key2: &str) -> bool { // "bar/foo" will be returned. pub fn key_get(key1: &str, key2: &str) -> String { if let Some(i) = key2.find('*') { - if key1.len() > i { - if key1[..i] == key2[..i] { - return key1[i..].to_string(); - } + if key1.len() > i && key1[..i] == key2[..i] { + return key1[i..].to_string(); } } "".to_string()