From ac87c04eb3734b739cb7fbf53d923ea3d88a3710 Mon Sep 17 00:00:00 2001
From: Jiale Zhang <zhangjiale@linux.alibaba.com>
Date: Sat, 6 May 2023 11:15:40 +0800
Subject: [PATCH] AS: Add API to set policy

Signed-off-by: Jiale Zhang <zhangjiale@linux.alibaba.com>
---
 bin/grpc-as/proto/attestation.proto |  6 ++++++
 bin/grpc-as/src/server.rs           | 21 ++++++++++++++++++++-
 src/lib.rs                          |  5 +++++
 src/policy_engine/mod.rs            |  2 ++
 src/policy_engine/opa/mod.rs        |  7 +++++++
 5 files changed, 40 insertions(+), 1 deletion(-)

diff --git a/bin/grpc-as/proto/attestation.proto b/bin/grpc-as/proto/attestation.proto
index 6f2ff6a..6f44de7 100644
--- a/bin/grpc-as/proto/attestation.proto
+++ b/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)
 }
diff --git a/bin/grpc-as/src/server.rs b/bin/grpc-as/src/server.rs
index ad9a781..f56df85 100644
--- a/bin/grpc-as/src/server.rs
+++ b/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,23 @@ 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)
+            .map_err(|e| Status::aborted(format!("Set Attestation Policy Failed: {e}")))?;
+
+        Ok(Response::new(SetPolicyResponse {}))
+    }
+
     async fn attestation_evaluate(
         &self,
         request: Request<AttestationRequest>,
diff --git a/src/lib.rs b/src/lib.rs
index 3b595d7..0cfc90f 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -84,6 +84,11 @@ impl AttestationService {
         })
     }
 
+    /// Set Attestation Verification Policy.
+    pub fn set_policy(&mut self, policy_base64_string: String) -> Result<()> {
+        self.policy_engine.set_policy(policy_base64_string)
+    }
+
     /// Evaluate Attestation Evidence.
     pub async fn evaluate(
         &self,
diff --git a/src/policy_engine/mod.rs b/src/policy_engine/mod.rs
index f8f67ce..18f4f7b 100644
--- a/src/policy_engine/mod.rs
+++ b/src/policy_engine/mod.rs
@@ -27,4 +27,6 @@ pub trait PolicyEngine {
         reference_data_map: HashMap<String, Vec<String>>,
         input: String,
     ) -> Result<(bool, String)>;
+
+    fn set_policy(&mut self, policy_base64_string: String) -> Result<()>;
 }
diff --git a/src/policy_engine/opa/mod.rs b/src/policy_engine/opa/mod.rs
index b4c305a..cf8d8df 100644
--- a/src/policy_engine/opa/mod.rs
+++ b/src/policy_engine/opa/mod.rs
@@ -85,6 +85,13 @@ impl PolicyEngine for OPA {
 
         Ok((res_kv["allow"].as_bool().unwrap_or(false), res))
     }
+
+    fn set_policy(&mut self, policy: String) -> Result<()> {
+        let policy_bytes = base64::decode_config(policy, base64::URL_SAFE_NO_PAD)
+            .map_err(|_| anyhow!("Base64 decode OPA policy string failed"))?;
+        fs::write(&self.policy_file_path, policy_bytes)
+            .map_err(|e| anyhow!("Write OPA policy to file failed: {:?}", e))
+    }
 }
 
 #[cfg(test)]