Improvements

1. Skip recording undefined variables 2. Parse `in` correctly if it is not imported. 3. base64.decode Signed-off-by: Anand Krishnamoorthi <[email protected]>
anakrish · Nov 1, 2023 · d982807 · d982807
1 parent 6228eaa
commit d982807
Show file tree

Hide file tree

Showing 7 changed files with 61 additions and 33 deletions.
diff --git a/Cargo.toml b/Cargo.toml
@@ -14,6 +14,7 @@ log = "0.4.17"
 env_logger="0.10.0"
 lazy_static = "1.4.0"
 rand = "0.8.5"
+data-encoding = "2.4.0"
 
 [dev-dependencies]
 clap = { version = "4.4.7", features = ["derive"] }

diff --git a/src/builtins/encoding.rs b/src/builtins/encoding.rs
@@ -0,0 +1,28 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+use crate::ast::Expr;
+use crate::builtins;
+use crate::builtins::utils::{ensure_args_count, ensure_string};
+use crate::lexer::Span;
+use crate::value::Value;
+
+use std::collections::HashMap;
+
+use anyhow::Result;
+use data_encoding::BASE64;
+
+pub fn register(m: &mut HashMap<&'static str, builtins::BuiltinFcn>) {
+    m.insert("base64.decode", (base64_decode, 1));
+}
+
+fn base64_decode(span: &Span, params: &[Expr], args: &[Value]) -> Result<Value> {
+    let name = "base64.decode";
+    ensure_args_count(span, name, params, args, 1)?;
+
+    let encoded_str = ensure_string(name, &params[0], &args[0])?;
+    let decoded_bytes = BASE64.decode(encoded_str.as_bytes())?;
+    Ok(Value::String(
+        String::from_utf8_lossy(&decoded_bytes).to_string(),
+    ))
+}
diff --git a/src/builtins/mod.rs b/src/builtins/mod.rs
@@ -7,6 +7,7 @@ mod bitwise;
 pub mod comparison;
 mod conversions;
 mod debugging;
+mod encoding;
 pub mod numbers;
 mod objects;
 pub mod sets;
@@ -44,7 +45,7 @@ lazy_static! {
 	conversions::register(&mut m);
 	//units::register(&mut m);
 	types::register(&mut m);
-	//encoding::register(&mut m);
+	encoding::register(&mut m);
 	//token_signing::register(&mut m);
 	//token_verification::register(&mut m);
 	//time::register(&mut m);

diff --git a/src/interpreter.rs b/src/interpreter.rs
@@ -478,6 +478,11 @@ impl<'source> Interpreter<'source> {
             }
         };
 
+        // Omit recording undefined values.
+        if value == Value::Undefined {
+            return Ok(Value::Bool(false));
+        }
+
         self.add_variable_or(name)?;
 
         // TODO: optimize this

diff --git a/src/parser.rs b/src/parser.rs
@@ -714,7 +714,7 @@ impl<'source> Parser<'source> {
         let start = self.tok.1.start;
         let mut expr = self.parse_bool_expr()?;
 
-        while self.tok.1.text() == "in" {
+        while self.tok.1.text() == "in" && self.future_keywords.get("in").is_some() {
             expr = self.parse_membership_tail(start, expr, None)?;
         }
 
@@ -731,7 +731,7 @@ impl<'source> Parser<'source> {
             expr = self.parse_membership_tail(start, expr, Some(value))?;
         }
 
-        while self.tok.1.text() == "in" {
+        while self.tok.1.text() == "in" && self.is_imported_future_keyword("in") {
             expr = self.parse_membership_tail(start, expr, None)?;
         }
 
@@ -834,7 +834,7 @@ impl<'source> Parser<'source> {
             vars.push(span);
         }
 
-        if self.tok.1.text() != "in" || self.future_keywords.get("in").is_none() {
+        if self.tok.1.text() != "in" || !self.is_imported_future_keyword("in") {
             if self.tok.1.text() == "in" {
                 self.warn_future_keyword();
             }

diff --git a/tests/interpreter/mod.rs b/tests/interpreter/mod.rs
@@ -181,18 +181,14 @@ pub fn check_output(computed_results: &[Value], expected_results: &[Value]) -> R
     Ok(())
 }
 
-fn query_results_to_value(query_results: QueryResults) -> Result<Value> {
+fn push_query_results(query_results: QueryResults, results: &mut Vec<Value>) {
     if let Some(query_result) = query_results.result.last() {
         if !query_result.bindings.is_empty_object() {
-            return Ok(query_result.bindings.clone());
-        } else {
-            return match query_result.expressions.last() {
-                Some(v) => Ok(v["value"].clone()),
-                _ => bail!("no expressions in query results"),
-            };
+            results.push(query_result.bindings.clone());
+        } else if let Some(v) = query_result.expressions.last() {
+            results.push(v["value"].clone());
         }
     }
-    bail!("query result incomplete")
 }
 
 pub fn eval_file_first_rule(
@@ -268,22 +264,20 @@ pub fn eval_file_first_rule(
             }
 
             // Now eval the query.
-            results.push(query_results_to_value(interpreter.eval_user_query(
-                &query_node,
-                &query_schedule,
-                enable_tracing,
-            )?)?);
+            push_query_results(
+                interpreter.eval_user_query(&query_node, &query_schedule, enable_tracing)?,
+                &mut results,
+            );
         }
     } else {
         // it no input is defined then one evaluation of all modules is performed
         interpreter.eval(&data_opt, &None, enable_tracing, Some(schedule))?;
 
-        // Now eval the query.
-        results.push(query_results_to_value(interpreter.eval_user_query(
-            &query_node,
-            &query_schedule,
-            enable_tracing,
-        )?)?);
+        // Now eval the query
+        push_query_results(
+            interpreter.eval_user_query(&query_node, &query_schedule, enable_tracing)?,
+            &mut results,
+        );
     }
 
     Ok(results)
@@ -359,22 +353,20 @@ pub fn eval_file(
             interpreter.eval_modules(&Some(input), enable_tracing)?;
 
             // Now eval the query.
-            results.push(query_results_to_value(interpreter.eval_user_query(
-                &query_node,
-                &query_schedule,
-                enable_tracing,
-            )?)?);
+            push_query_results(
+                interpreter.eval_user_query(&query_node, &query_schedule, enable_tracing)?,
+                &mut results,
+            );
         }
     } else {
         // it no input is defined then one evaluation of all modules is performed
         interpreter.eval(&data_opt, &None, enable_tracing, Some(schedule))?;
 
         // Now eval the query.
-        results.push(query_results_to_value(interpreter.eval_user_query(
-            &query_node,
-            &query_schedule,
-            enable_tracing,
-        )?)?);
+        push_query_results(
+            interpreter.eval_user_query(&query_node, &query_schedule, enable_tracing)?,
+            &mut results,
+        );
     }
 
     Ok(results)

diff --git a/tests/parser/cases/some/some.in.yaml b/tests/parser/cases/some/some.in.yaml
@@ -134,6 +134,7 @@ cases:
   - note: same-line
     rego: |
       package test
+      import future.keywords.in
       x = y{
           some a b in {4, 5}
           [1, 2, 3][a] == 3