AS: Add API to set policy

Signed-off-by: Jiale Zhang <zhangjiale@linux.alibaba.com>

.gitignore

@@ -8,5 +8,6 @@ target
 reference_values
 
 test_data/*_output.txt
+test_data/opa/
 
 tools/

Cargo.toml

@@ -57,7 +57,7 @@ sled = "0.34.7"
 strum = "0.24.0"
 strum_macros = "0.24.0"
 tempfile = "3.3.0"
-tokio = { version = "1.0", features = ["rt-multi-thread", "macros"] }
+tokio = { version = "1.0", features = ["rt-multi-thread", "macros", "fs"] }
 tonic = { version = "0.8.1", optional = true }
 uuid = { version = "1.1.2", features = ["v4"] }
 

as-types/src/lib.rs

@@ -48,3 +48,10 @@ pub struct ResultOutput {
     pub verifier_output: Option<String>,
     pub policy_engine_output: Option<String>,
 }
+
+#[derive(Serialize, Deserialize, Debug, Clone)]
+pub struct SetPolicyInput {
+    pub r#type: String,
+    pub policy_id: String,
+    pub policy: String,
+}

bin/grpc-as/proto/attestation.proto

@@ -19,7 +19,13 @@ message AttestationResponse {
     string attestation_results = 1;
 }
 
+message SetPolicyRequest {
+    string input = 1;
+}
+message SetPolicyResponse {}
+
 service AttestationService {
     rpc AttestationEvaluate(AttestationRequest) returns (AttestationResponse) {};
+    rpc SetAttestationPolicy(SetPolicyRequest) returns (SetPolicyResponse) {};
     // Get the GetPolicyRequest.user and GetPolicyRequest.tee specified Policy(.rego)
 }

bin/grpc-as/src/server.rs

@@ -7,7 +7,9 @@ use tonic::transport::Server;
 use tonic::{Request, Response, Status};
 
 use crate::as_api::attestation_service_server::{AttestationService, AttestationServiceServer};
-use crate::as_api::{AttestationRequest, AttestationResponse, Tee as GrpcTee};
+use crate::as_api::{
+    AttestationRequest, AttestationResponse, SetPolicyRequest, SetPolicyResponse, Tee as GrpcTee,
+};
 
 use crate::rvps_api::reference_value_provider_service_server::{
     ReferenceValueProviderService, ReferenceValueProviderServiceServer,
@@ -61,6 +63,27 @@ impl AttestationServer {
 
 #[tonic::async_trait]
 impl AttestationService for Arc<RwLock<AttestationServer>> {
+    async fn set_attestation_policy(
+        &self,
+        request: Request<SetPolicyRequest>,
+    ) -> Result<Response<SetPolicyResponse>, Status> {
+        let request: SetPolicyRequest = request.into_inner();
+
+        debug!("SetPolicyInput: {}", &request.input);
+
+        let set_policy_input: as_types::SetPolicyInput = serde_json::from_str(&request.input)
+            .map_err(|_| Status::aborted("Bad SetPolicyInput"))?;
+
+        self.write()
+            .await
+            .attestation_service
+            .set_policy(set_policy_input)
+            .await
+            .map_err(|e| Status::aborted(format!("Set Attestation Policy Failed: {e}")))?;
+
+        Ok(Response::new(SetPolicyResponse {}))
+    }
+
     async fn attestation_evaluate(
         &self,
         request: Request<AttestationRequest>,

src/lib.rs

@@ -20,6 +20,7 @@ pub mod rvps;
 pub mod verifier;
 
 use anyhow::{anyhow, Context, Result};
+use as_types::SetPolicyInput;
 use config::Config;
 pub use kbs_types::{Attestation, Tee};
 use policy_engine::PolicyEngine;
@@ -84,6 +85,14 @@ impl AttestationService {
         })
     }
 
+    /// Set Attestation Verification Policy.
+    pub async fn set_policy(&mut self, input: SetPolicyInput) -> Result<()> {
+        self.policy_engine
+            .set_policy(input)
+            .await
+            .map_err(|e| anyhow!("Cannot Set Policy: {:?}", e))
+    }
+
     /// Evaluate Attestation Evidence.
     pub async fn evaluate(
         &self,
@@ -115,9 +124,11 @@ impl AttestationService {
             .await
             .map_err(|e| anyhow!("Generate reference data failed{:?}", e))?;
 
+        // Now only support using default policy to evaluate
         let (result, policy_engine_output) = self
             .policy_engine
-            .evaluate(reference_data_map, tcb.clone())?;
+            .evaluate(reference_data_map, tcb.clone(), None)
+            .await?;
 
         let attestation_results =
             AttestationResults::new(tee, result, None, Some(policy_engine_output), Some(tcb));

src/policy_engine/mod.rs

@@ -1,4 +1,6 @@
 use anyhow::Result;
+use as_types::SetPolicyInput;
+use async_trait::async_trait;
 use serde::Deserialize;
 use std::collections::HashMap;
 use std::path::Path;
@@ -21,10 +23,14 @@ impl PolicyEngineType {
     }
 }
 
+#[async_trait]
 pub trait PolicyEngine {
-    fn evaluate(
+    async fn evaluate(
         &self,
         reference_data_map: HashMap<String, Vec<String>>,
         input: String,
+        policy_id: Option<String>,
     ) -> Result<(bool, String)>;
+
+    async fn set_policy(&mut self, input: SetPolicyInput) -> Result<()>;
 }

src/policy_engine/opa/mod.rs

@@ -1,5 +1,7 @@
 use crate::policy_engine::PolicyEngine;
-use anyhow::{anyhow, Result};
+use anyhow::{anyhow, bail, Result};
+use as_types::SetPolicyInput;
+use async_trait::async_trait;
 use serde_json::Value;
 use std::collections::HashMap;
 use std::ffi::CStr;
@@ -23,36 +25,51 @@ pub struct GoString {
 
 #[derive(Debug)]
 pub struct OPA {
-    policy_file_path: PathBuf,
+    policy_dir_path: PathBuf,
 }
 
 impl OPA {
     pub fn new(work_dir: PathBuf) -> Result<Self> {
-        let mut policy_file_path = work_dir;
+        let mut policy_dir_path = work_dir;
 
-        policy_file_path.push("opa");
-        if !policy_file_path.as_path().exists() {
-            fs::create_dir_all(&policy_file_path)
+        policy_dir_path.push("opa");
+        if !policy_dir_path.as_path().exists() {
+            fs::create_dir_all(&policy_dir_path)
                 .map_err(|e| anyhow!("Create policy dir failed: {:?}", e))?;
         }
 
-        policy_file_path.push("policy.rego");
-        if !policy_file_path.as_path().exists() {
+        let mut default_policy_path = PathBuf::from(
+            &policy_dir_path
+                .to_str()
+                .ok_or_else(|| anyhow!("Policy DirPath to string failed"))?,
+        );
+        default_policy_path.push("default.rego");
+        if !default_policy_path.as_path().exists() {
             let policy = std::include_str!("default_policy.rego").to_string();
-            fs::write(&policy_file_path, policy)?;
+            fs::write(&default_policy_path, policy)?;
         }
 
-        Ok(Self { policy_file_path })
+        Ok(Self { policy_dir_path })
     }
 }
 
+#[async_trait]
 impl PolicyEngine for OPA {
-    fn evaluate(
+    async fn evaluate(
         &self,
         reference_data_map: HashMap<String, Vec<String>>,
         input: String,
+        policy_id: Option<String>,
     ) -> Result<(bool, String)> {
-        let policy = fs::read_to_string(&self.policy_file_path)
+        let policy_file_path = format!(
+            "{}/{}.rego",
+            self.policy_dir_path
+                .to_str()
+                .ok_or_else(|| anyhow!("Miss Policy DirPath"))?,
+            policy_id.unwrap_or("default".to_string())
+        );
+        let policy = tokio::fs::read_to_string(policy_file_path)
+            .await
             .map_err(|e| anyhow!("Read OPA policy file failed: {:?}", e))?;
 
         let policy_go = GoString {
@@ -85,6 +102,26 @@ impl PolicyEngine for OPA {
 
         Ok((res_kv["allow"].as_bool().unwrap_or(false), res))
     }
+
+    async fn set_policy(&mut self, input: SetPolicyInput) -> Result<()> {
+        if input.r#type != "rego" {
+            bail!("OPA Policy Engine only support .rego policy");
+        }
+
+        let policy_bytes = base64::decode_config(input.policy, base64::URL_SAFE_NO_PAD)
+            .map_err(|e| anyhow!("Base64 decode OPA policy string failed: {:?}", e))?;
+        let mut policy_file_path = PathBuf::from(
+            &self
+                .policy_dir_path
+                .to_str()
+                .ok_or_else(|| anyhow!("Policy DirPath to string failed"))?,
+        );
+        policy_file_path.push(format!("{}.rego", input.policy_id));
+
+        tokio::fs::write(&policy_file_path, policy_bytes)
+            .await
+            .map_err(|e| anyhow!("Write OPA policy to file failed: {:?}", e))
+    }
 }
 
 #[cfg(test)]
@@ -108,21 +145,46 @@ mod tests {
         .to_string()
     }
 
-    #[test]
-    fn test_evaluate() {
+    #[tokio::test]
+    async fn test_evaluate() {
         let opa = OPA {
-            policy_file_path: PathBuf::from("./src/policy_engine/opa/default_policy.rego"),
+            policy_dir_path: PathBuf::from("./src/policy_engine/opa"),
         };
+        let default_policy_id = "default_policy".to_string();
 
         let reference_data: HashMap<String, Vec<String>> =
             serde_json::from_str(&dummy_reference(5)).unwrap();
 
-        let res = opa.evaluate(reference_data.clone(), dummy_input(5, 5));
+        let res = opa
+            .evaluate(
+                reference_data.clone(),
+                dummy_input(5, 5),
+                Some(default_policy_id.clone()),
+            )
+            .await;
         assert!(res.is_ok(), "OPA execution() should be success");
         assert!(res.unwrap().0 == true, "allow should be true");
 
-        let res = opa.evaluate(reference_data, dummy_input(0, 0));
+        let res = opa
+            .evaluate(reference_data, dummy_input(0, 0), Some(default_policy_id))
+            .await;
         assert!(res.is_ok(), "OPA execution() should be success");
         assert!(res.unwrap().0 == false, "allow should be false");
     }
+
+    #[tokio::test]
+    async fn test_set_policy() {
+        let mut opa = OPA::new(PathBuf::from("./test_data")).unwrap();
+        let policy = "package policy
+default allow = true"
+            .to_string();
+
+        let input = SetPolicyInput {
+            r#type: "rego".to_string(),
+            policy_id: "test".to_string(),
+            policy: base64::encode_config(policy, base64::URL_SAFE_NO_PAD),
+        };
+
+        assert!(opa.set_policy(input).await.is_ok());
+    }
 }