From 88243e52b7dcb14a099c8448110cf59aeb54d859 Mon Sep 17 00:00:00 2001 From: Jiale Zhang Date: Sat, 6 May 2023 11:15:40 +0800 Subject: [PATCH] AS: Add API to set policy Signed-off-by: Jiale Zhang --- .gitignore | 1 + Cargo.toml | 2 +- as-types/src/lib.rs | 7 +++ bin/grpc-as/proto/attestation.proto | 6 ++ bin/grpc-as/src/server.rs | 25 +++++++- src/lib.rs | 13 +++- src/policy_engine/mod.rs | 8 ++- src/policy_engine/opa/mod.rs | 96 ++++++++++++++++++++++++----- 8 files changed, 137 insertions(+), 21 deletions(-) diff --git a/.gitignore b/.gitignore index b58d716..07b7477 100644 --- a/.gitignore +++ b/.gitignore @@ -8,5 +8,6 @@ target reference_values test_data/*_output.txt +test_data/opa/ tools/ \ No newline at end of file diff --git a/Cargo.toml b/Cargo.toml index a0d3e2b..3c41d3d 100644 --- a/Cargo.toml +++ b/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"] } diff --git a/as-types/src/lib.rs b/as-types/src/lib.rs index e6d38d1..f5e0fd8 100644 --- a/as-types/src/lib.rs +++ b/as-types/src/lib.rs @@ -48,3 +48,10 @@ pub struct ResultOutput { pub verifier_output: Option, pub policy_engine_output: Option, } + +#[derive(Serialize, Deserialize, Debug, Clone)] +pub struct SetPolicyInput { + pub r#type: String, + pub policy_id: String, + pub policy: String, +} diff --git a/bin/grpc-as/proto/attestation.proto b/bin/grpc-as/proto/attestation.proto index 6f2ff6a..f9def45 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 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) } diff --git a/bin/grpc-as/src/server.rs b/bin/grpc-as/src/server.rs index ad9a781..0bca654 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,27 @@ impl AttestationServer { #[tonic::async_trait] impl AttestationService for Arc> { + async fn set_attestation_policy( + &self, + request: Request, + ) -> Result, 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, diff --git a/src/lib.rs b/src/lib.rs index 3b595d7..20a5dff 100644 --- a/src/lib.rs +++ b/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)); diff --git a/src/policy_engine/mod.rs b/src/policy_engine/mod.rs index f8f67ce..17a5f91 100644 --- a/src/policy_engine/mod.rs +++ b/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>, input: String, + policy_id: Option, ) -> Result<(bool, String)>; + + async fn set_policy(&mut self, input: SetPolicyInput) -> Result<()>; } diff --git a/src/policy_engine/opa/mod.rs b/src/policy_engine/opa/mod.rs index b4c305a..54a707f 100644 --- a/src/policy_engine/opa/mod.rs +++ b/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 { - 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>, input: String, + policy_id: Option, ) -> 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> = 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()); + } }