Skip to content

Commit

Permalink
fix: pass test
Browse files Browse the repository at this point in the history
  • Loading branch information
@MuZhou233
MuZhou233 committed Sep 2, 2024
1 parent e7e0a91 commit b847143
Show file tree
Hide file tree
Showing 3 changed files with 46 additions and 87 deletions.
4 changes: 2 additions & 2 deletions src/core_api.rs
View file
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
use crate::{
model::OperatorFunction,
enforcer::EnforceContext, Adapter, Effector, EnforceArgs, Event,
EventEmitter, Filter, Model, Result, RoleManager, TryIntoAdapter,
TryIntoModel,
Expand All @@ -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;

Expand All @@ -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;
Expand Down
32 changes: 21 additions & 11 deletions src/enforcer.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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},
Expand Down Expand Up @@ -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") {
Expand Down Expand Up @@ -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 {
Expand Down Expand Up @@ -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]
Expand Down Expand Up @@ -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,
Expand Down
97 changes: 23 additions & 74 deletions src/model/function_map.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,86 +15,37 @@ use rhai::{Dynamic, ImmutableString};
static MAT_B: Lazy<Regex> = Lazy::new(|| Regex::new(r":[^/]*").unwrap());
static MAT_P: Lazy<Regex> = 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<String, fn(&[ImmutableString]) -> Dynamic>,
pub(crate) fm: HashMap<String, OperatorFunction>,
}

impl Default for FunctionMap {
fn default() -> FunctionMap {
let mut fm: HashMap<String, fn(&[ImmutableString]) -> Dynamic> =
let mut fm: HashMap<String, OperatorFunction> =
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 }
}
Expand All @@ -105,15 +56,15 @@ impl FunctionMap {
pub fn add_function(
&mut self,
fname: &str,
f: fn(&[ImmutableString]) -> Dynamic,
f: OperatorFunction,
) {
self.fm.insert(fname.to_owned(), f);
}

#[inline]
pub fn get_functions(
&self,
) -> impl Iterator<Item = (&String, &fn(&[ImmutableString]) -> Dynamic)>
) -> impl Iterator<Item = (&String, &OperatorFunction)>
{
self.fm.iter()
}
Expand All @@ -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()
Expand Down

0 comments on commit b847143

Please sign in to comment.