genpolicy: Add state

Use regorous engine's add_data method to add state to the policy.
This data can later be accessed inside rego context through the data namespace.

Support state modifications (json-patches) that may be returned as a result from policy evaluation.

Also initialize a policy engine data slice "pstate" dedicated for storing state.

Fixes #10087

Signed-off-by: Saul Paredes <saulparedes@microsoft.com>

src/agent/Cargo.lock

@@ -1962,6 +1962,15 @@ dependencies = [
  "miniz_oxide",
 ]
 
+[[package]]
+name = "fluent-uri"
+version = "0.1.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "17c704e9dbe1ddd863da1e6ff3567795087b1eb201ce80d8fa81162e1516500d"
+dependencies = [
+ "bitflags 1.3.2",
+]
+
 [[package]]
 name = "fnv"
 version = "1.0.7"
@@ -2807,15 +2816,6 @@ dependencies = [
  "either",
 ]
 
-[[package]]
-name = "itertools"
-version = "0.12.1"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "ba291022dbbd398a455acf126c1e341954079855bc60dfdda641363bd6922569"
-dependencies = [
- "either",
-]
-
 [[package]]
 name = "itoa"
 version = "1.0.11"
@@ -2874,6 +2874,18 @@ dependencies = [
  "smallvec",
 ]
 
+[[package]]
+name = "json-patch"
+version = "2.0.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5b1fb8864823fad91877e6caea0baca82e49e8db50f8e5c9f9a453e27d3330fc"
+dependencies = [
+ "jsonptr",
+ "serde",
+ "serde_json",
+ "thiserror",
+]
+
 [[package]]
 name = "json-syntax"
 version = "0.12.5"
@@ -2893,6 +2905,17 @@ dependencies = [
  "utf8-decode",
 ]
 
+[[package]]
+name = "jsonptr"
+version = "0.4.7"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1c6e529149475ca0b2820835d3dce8fcc41c6b943ca608d32f35b449255e4627"
+dependencies = [
+ "fluent-uri",
+ "serde",
+ "serde_json",
+]
+
 [[package]]
 name = "jsonschema"
 version = "0.18.0"
@@ -2996,6 +3019,7 @@ dependencies = [
  "futures",
  "image-rs",
  "ipnetwork",
+ "json-patch",
  "kata-sys-util",
  "kata-types",
  "lazy_static",
@@ -4902,14 +4926,12 @@ checksum = "7a66a03ae7c801facd77a29370b4faec201768915ac14a721ba36f20bc9c209b"
 
 [[package]]
 name = "regorus"
-version = "0.1.5"
+version = "0.2.8"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "77dd872918e5c172bd42ac49716f89a15e35be513bba3d902e355a531529a87f"
+checksum = "843c3d97f07e3b5ac0955d53ad0af4c91fe4a4f8525843ece5bf014f27829b73"
 dependencies = [
  "anyhow",
- "itertools 0.12.1",
  "lazy_static",
- "num",
  "rand",
  "regex",
  "scientific",
@@ -5409,7 +5431,7 @@ dependencies = [
  "aes",
  "aes-gcm",
  "anyhow",
- "base64 0.21.7",
+ "base64 0.22.1",
  "block-padding",
  "blowfish",
  "buffered-reader",

src/agent/Cargo.toml

@@ -80,11 +80,13 @@ strum_macros = "0.26.2"
 image-rs = { git = "https://github.com/confidential-containers/guest-components", rev = "v0.10.0", default-features = false, optional = true }
 
 # Agent Policy
-regorus = { version = "0.1.4", default-features = false, features = [
+regorus = { version = "0.2.6", default-features = false, features = [
     "arc",
     "regex",
+    "std",
 ], optional = true }
 cdi = { git = "https://github.com/cncf-tags/container-device-interface-rs", rev = "fba5677a8e7cc962fc6e495fcec98d7d765e332a" }
+json-patch = "2.0.0"
 
 [dev-dependencies]
 tempfile = "3.1.0"

src/agent/src/policy.rs

@@ -3,7 +3,7 @@
 // SPDX-License-Identifier: Apache-2.0
 //
 
-use anyhow::Result;
+use anyhow::{bail, Result};
 use protobuf::MessageDyn;
 use tokio::io::AsyncWriteExt;
 
@@ -68,6 +68,12 @@ pub struct AgentPolicy {
     engine: regorus::Engine,
 }
 
+#[derive(serde::Deserialize, Debug)]
+struct MetadataResponse {
+    allowed: bool,
+    ops: Option<json_patch::Patch>,
+}
+
 impl AgentPolicy {
     /// Create AgentPolicy object.
     pub fn new() -> Self {
@@ -82,6 +88,17 @@ impl AgentPolicy {
         let mut engine = regorus::Engine::new();
         engine.set_strict_builtin_errors(false);
         engine.set_gather_prints(true);
+        // assign a slice of the engine data "pstate" to be used as policy state
+        engine
+            .add_data(
+                regorus::Value::from_json_str(
+                    r#"{
+                        "pstate": {}
+                    }"#,
+                )
+                .unwrap(),
+            )
+            .unwrap();
         engine
     }
 
@@ -112,6 +129,23 @@ impl AgentPolicy {
         Ok(())
     }
 
+    async fn apply_patch_to_state(&mut self, patch: json_patch::Patch) -> Result<()> {
+        // Convert the current engine data to a JSON value
+        let mut state = serde_json::to_value(self.engine.get_data())?;
+
+        // Apply the patch to the state
+        json_patch::patch(&mut state, &patch)?;
+
+        // Clear the existing data in the engine
+        self.engine.clear_data();
+
+        // Add the patched state back to the engine
+        self.engine
+            .add_data(regorus::Value::from_json_str(&state.to_string())?)?;
+
+        Ok(())
+    }
+
     /// Ask regorus if an API call should be allowed or not.
     async fn allow_request(&mut self, ep: &str, ep_input: &str) -> Result<(bool, String)> {
         debug!(sl!(), "policy check: {ep}");
@@ -120,13 +154,56 @@ impl AgentPolicy {
         let query = format!("data.agent_policy.{ep}");
         self.engine.set_input_json(ep_input)?;
 
-        let mut allow = match self.engine.eval_bool_query(query, false) {
-            Ok(a) => a,
-            Err(e) => {
-                if !self.allow_failures {
-                    return Err(e);
+        let results = self.engine.eval_query(query, false)?;
+
+        let prints = match self.engine.take_prints() {
+            Ok(p) => p.join(" "),
+            Err(e) => format!("Failed to get policy log: {e}"),
+        };
+
+        if results.result.len() != 1 {
+            // Results are empty when AllowRequestsFailingPolicy is used to allow a Request that hasn't been defined in the policy
+            if self.allow_failures {
+                return Ok((true, prints));
+            }
+            bail!(
+                "policy check: unexpected eval_query result len {:?}",
+                results
+            );
+        }
+
+        if results.result[0].expressions.len() != 1 {
+            bail!(
+                "policy check: unexpected eval_query result expressions {:?}",
+                results
+            );
+        }
+
+        let mut allow = match &results.result[0].expressions[0].value {
+            regorus::Value::Bool(b) => *b,
+
+            // Match against a specific variant that could be interpreted as MetadataResponse
+            regorus::Value::Object(obj) => {
+                let json_str = serde_json::to_string(obj)?;
+
+                self.log_eval_input(ep, &json_str).await;
+
+                let metadata_response: MetadataResponse = serde_json::from_str(&json_str)?;
+
+                if metadata_response.allowed {
+                    if let Some(ops) = metadata_response.ops {
+                        self.apply_patch_to_state(ops).await?;
+                    }
                 }
-                false
+                metadata_response.allowed
+            }
+
+            _ => {
+                error!(sl!(), "allow_request: unexpected eval_query result type");
+                bail!(
+                    "policy check: unexpected eval_query result type {:?}",
+                    results
+                );
             }
         };
 
@@ -135,11 +212,6 @@ impl AgentPolicy {
             allow = true;
         }
 
-        let prints = match self.engine.take_prints() {
-            Ok(p) => p.join(" "),
-            Err(e) => format!("Failed to get policy log: {e}"),
-        };
-
         Ok((allow, prints))
     }