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"] }
 

bin/grpc-as/proto/attestation.proto

@@ -19,7 +19,13 @@ message AttestationResponse {
     string attestation_results = 1;
 }
 
+message SetPolicyRequest {
+    string policy = 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,24 @@ 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!("Policy: {}", &request.policy);
+
+        self.write()
+            .await
+            .attestation_service
+            .set_policy(request.policy)
+            .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

@@ -84,6 +84,14 @@ impl AttestationService {
         })
     }
 
+    /// Set Attestation Verification Policy.
+    pub async fn set_policy(&mut self, policy_base64_string: String) -> Result<()> {
+        self.policy_engine
+            .set_policy(policy_base64_string)
+            .await
+            .map_err(|e| anyhow!("Cannot Set Policy: {:?}", e))
+    }
+
     /// Evaluate Attestation Evidence.
     pub async fn evaluate(
         &self,
@@ -117,7 +125,8 @@ impl AttestationService {
 
         let (result, policy_engine_output) = self
             .policy_engine
-            .evaluate(reference_data_map, tcb.clone())?;
+            .evaluate(reference_data_map, tcb.clone())
+            .await?;
 
         let attestation_results =
             AttestationResults::new(tee, result, None, Some(policy_engine_output), Some(tcb));

src/policy_engine/mod.rs

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

src/policy_engine/opa/mod.rs

@@ -1,5 +1,6 @@
 use crate::policy_engine::PolicyEngine;
 use anyhow::{anyhow, Result};
+use async_trait::async_trait;
 use serde_json::Value;
 use std::collections::HashMap;
 use std::ffi::CStr;
@@ -46,8 +47,9 @@ impl OPA {
     }
 }
 
+#[async_trait]
 impl PolicyEngine for OPA {
-    fn evaluate(
+    async fn evaluate(
         &self,
         reference_data_map: HashMap<String, Vec<String>>,
         input: String,
@@ -85,6 +87,14 @@ impl PolicyEngine for OPA {
 
         Ok((res_kv["allow"].as_bool().unwrap_or(false), res))
     }
+
+    async fn set_policy(&mut self, policy: String) -> Result<()> {
+        let policy_bytes = base64::decode_config(policy, base64::URL_SAFE_NO_PAD)
+            .map_err(|e| anyhow!("Base64 decode OPA policy string failed: {:?}", e))?;
+        tokio::fs::write(&self.policy_file_path, policy_bytes)
+            .await
+            .map_err(|e| anyhow!("Write OPA policy to file failed: {:?}", e))
+    }
 }
 
 #[cfg(test)]
@@ -108,7 +118,7 @@ mod tests {
         .to_string()
     }
 
-    #[test]
+    #[tokio::test]
     fn test_evaluate() {
         let opa = OPA {
             policy_file_path: PathBuf::from("./src/policy_engine/opa/default_policy.rego"),
@@ -117,12 +127,22 @@ mod tests {
         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))
+            .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));
         assert!(res.is_ok(), "OPA execution() should be success");
         assert!(res.unwrap().0 == false, "allow should be false");
     }
+
+    #[tokio::test]
+    fn test_set_policy() {
+        let opa = OPA::new(PathBuf::from("./test_data")).unwrap();
+        let policy_base64_string = "cGFja2FnZSBwb2xpY3kKCmltcG9ydCBmdXR1cmUua2V5d29yZHMuZXZlcnkKCmRlZmF1bHQgYWxsb3cgPSBmYWxzZQo=";
+
+        assert!(opa.set_policy(policy_base64_string).await);
+    }
 }