feat: add key_get series built-in operator (#338)

* feat: add key_get series built-in operator * fix: pass test * fix: pass test * fix: cargo fmt
casbin · Sep 2, 2024 · b5d3765 · b5d3765
1 parent f4c3350
commit b5d3765
Show file tree

Hide file tree

Showing 4 changed files with 345 additions and 45 deletions.
diff --git a/src/cached_enforcer.rs b/src/cached_enforcer.rs
@@ -8,7 +8,7 @@ use crate::{
     emitter::{clear_cache, Event, EventData, EventEmitter},
     enforcer::EnforceContext,
     enforcer::Enforcer,
-    model::Model,
+    model::{Model, OperatorFunction},
     rbac::RoleManager,
     Result,
 };
@@ -27,7 +27,7 @@ use crate::{error::ModelError, get_or_err};
 
 use async_trait::async_trait;
 use parking_lot::RwLock;
-use rhai::{Dynamic, ImmutableString};
+use rhai::Dynamic;
 
 use std::{collections::HashMap, sync::Arc};
 
@@ -123,11 +123,7 @@ impl CoreApi for CachedEnforcer {
     }
 
     #[inline]
-    fn add_function(
-        &mut self,
-        fname: &str,
-        f: fn(ImmutableString, ImmutableString) -> bool,
-    ) {
+    fn add_function(&mut self, fname: &str, f: OperatorFunction) {
         self.enforcer.add_function(fname, f);
     }
 

diff --git a/src/core_api.rs b/src/core_api.rs
@@ -1,7 +1,7 @@
 use crate::{
-    enforcer::EnforceContext, Adapter, Effector, EnforceArgs, Event,
-    EventEmitter, Filter, Model, Result, RoleManager, TryIntoAdapter,
-    TryIntoModel,
+    enforcer::EnforceContext, model::OperatorFunction, Adapter, Effector,
+    EnforceArgs, Event, EventEmitter, Filter, Model, Result, RoleManager,
+    TryIntoAdapter, TryIntoModel,
 };
 
 #[cfg(feature = "watcher")]
@@ -15,7 +15,6 @@ use crate::emitter::EventData;
 
 use async_trait::async_trait;
 use parking_lot::RwLock;
-use rhai::ImmutableString;
 
 use std::sync::Arc;
 
@@ -33,11 +32,7 @@ pub trait CoreApi: Send + Sync {
     ) -> Result<Self>
     where
         Self: Sized;
-    fn add_function(
-        &mut self,
-        fname: &str,
-        f: fn(ImmutableString, ImmutableString) -> bool,
-    );
+    fn add_function(&mut self, fname: &str, f: OperatorFunction);
     fn get_model(&self) -> &dyn Model;
     fn get_mut_model(&mut self) -> &mut dyn Model;
     fn get_adapter(&self) -> &dyn Adapter;

diff --git 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},
@@ -351,6 +351,23 @@ 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") {
             for (fname, ast) in ast_map {
@@ -380,7 +397,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 {
@@ -425,13 +442,9 @@ impl CoreApi for Enforcer {
     }
 
     #[inline]
-    fn add_function(
-        &mut self,
-        fname: &str,
-        f: fn(ImmutableString, ImmutableString) -> bool,
-    ) {
+    fn add_function(&mut self, fname: &str, f: OperatorFunction) {
         self.fm.add_function(fname, f);
-        self.engine.register_fn(fname, f);
+        Self::register_function(&mut self.engine, fname, f);
     }
 
     #[inline]
@@ -1340,7 +1353,11 @@ mod tests {
 
         e.add_function(
             "keyMatchCustom",
-            |s1: ImmutableString, s2: ImmutableString| key_match(&s1, &s2),
+            OperatorFunction::Arg2(
+                |s1: ImmutableString, s2: ImmutableString| {
+                    key_match(&s1, &s2).into()
+                },
+            ),
         );
 
         assert_eq!(