Merge pull request #229 from biscuit-auth/fix-param-substitution-in-c…

…losures fix: recursively collect and apply parameters in closures
biscuit-auth · Oct 22, 2024 · c2dd7a4 · c2dd7a4
2 parents 911ebe4 + 35edbd8
commit c2dd7a4
Show file tree

Hide file tree

Showing 2 changed files with 70 additions and 14 deletions.
diff --git a/biscuit-auth/src/token/builder.rs b/biscuit-auth/src/token/builder.rs
@@ -913,6 +913,43 @@ pub enum Op {
     Closure(Vec<String>, Vec<Op>),
 }
 
+impl Op {
+    fn collect_parameters(&self, parameters: &mut HashMap<String, Option<Term>>) {
+        match self {
+            Op::Value(Term::Parameter(ref name)) => {
+                parameters.insert(name.to_owned(), None);
+            }
+            Op::Closure(_, ops) => {
+                for op in ops {
+                    op.collect_parameters(parameters);
+                }
+            }
+            _ => {}
+        }
+    }
+
+    fn apply_parameters(self, parameters: &HashMap<String, Option<Term>>) -> Self {
+        match self {
+            Op::Value(Term::Parameter(ref name)) => {
+                if let Some(Some(t)) = parameters.get(name) {
+                    Op::Value(t.clone())
+                } else {
+                    self
+                }
+            }
+            Op::Value(_) => self,
+            Op::Unary(_) => self,
+            Op::Binary(_) => self,
+            Op::Closure(args, mut ops) => Op::Closure(
+                args,
+                ops.drain(..)
+                    .map(|op| op.apply_parameters(parameters))
+                    .collect(),
+            ),
+        }
+    }
+}
+
 impl Convert<datalog::Op> for Op {
     fn convert(&self, symbols: &mut SymbolTable) -> datalog::Op {
         match self {
@@ -1036,9 +1073,7 @@ impl Rule {
 
         for expression in &expressions {
             for op in &expression.ops {
-                if let Op::Value(Term::Parameter(name)) = &op {
-                    parameters.insert(name.to_string(), None);
-                }
+                op.collect_parameters(&mut parameters);
             }
         }
 
@@ -1282,14 +1317,7 @@ impl Rule {
                 expression.ops = expression
                     .ops
                     .drain(..)
-                    .map(|op| {
-                        if let Op::Value(Term::Parameter(name)) = &op {
-                            if let Some(Some(term)) = parameters.get(name) {
-                                return Op::Value(term.clone());
-                            }
-                        }
-                        op
-                    })
+                    .map(|op| op.apply_parameters(&parameters))
                     .collect();
             }
         }
@@ -2382,6 +2410,20 @@ mod tests {
         assert_eq!(s, "fact($var1, \"hello\", [0]) <- f1($var1, $var3), f2(\"hello\", $var3, 1), $var3.starts_with(\"hello\")");
     }
 
+    #[test]
+    fn set_closure_parameters() {
+        let mut rule = Rule::try_from("fact(true) <- false || {p1}").unwrap();
+        rule.set_lenient("p1", true).unwrap();
+        println!("{rule:?}");
+        let s = rule.to_string();
+        assert_eq!(s, "fact(true) <- false || true");
+
+        let mut rule = Rule::try_from("fact(true) <- false || {p1}").unwrap();
+        rule.set("p1", true).unwrap();
+        let s = rule.to_string();
+        assert_eq!(s, "fact(true) <- false || true");
+    }
+
     #[test]
     fn set_rule_scope_parameters() {
         let pubkey = PublicKey::from_bytes(

diff --git a/biscuit-parser/src/builder.rs b/biscuit-parser/src/builder.rs
@@ -188,6 +188,22 @@ pub enum Op {
     Closure(Vec<String>, Vec<Op>),
 }
 
+impl Op {
+    fn collect_parameters(&self, parameters: &mut HashMap<String, Option<Term>>) {
+        match self {
+            Op::Value(Term::Parameter(ref name)) => {
+                parameters.insert(name.to_owned(), None);
+            }
+            Op::Closure(_, ops) => {
+                for op in ops {
+                    op.collect_parameters(parameters);
+                }
+            }
+            _ => {}
+        }
+    }
+}
+
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub enum Unary {
     Negate,
@@ -332,9 +348,7 @@ impl Rule {
 
         for expression in &expressions {
             for op in &expression.ops {
-                if let Op::Value(Term::Parameter(name)) = &op {
-                    parameters.insert(name.to_string(), None);
-                }
+                op.collect_parameters(&mut parameters);
             }
         }