Skip to content

Commit

Permalink
opa: Refactor opa module errors
Browse files Browse the repository at this point in the history
Fixes: #394
Signed-off-by: Kartik Joshi <[email protected]>
  • Loading branch information
kartikjoshi21 committed Jun 18, 2024
1 parent 3dfbfb3 commit cc55b51
Show file tree
Hide file tree
Showing 2 changed files with 72 additions and 29 deletions.
44 changes: 40 additions & 4 deletions attestation-service/attestation-service/src/policy_engine/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2,8 +2,10 @@ use anyhow::Result;
use async_trait::async_trait;
use serde::Deserialize;
use std::collections::HashMap;
use std::io;
use std::path::Path;
use strum::EnumString;
use thiserror::Error;

pub mod opa;

Expand All @@ -13,6 +15,40 @@ pub enum PolicyEngineType {
OPA,
}

#[derive(Error, Debug)]
pub enum RegoError {
#[error("Failed to create policy directory: {0}")]
CreatePolicyDirFailed(#[source] io::Error),
#[error("Failed to convert policy directory path to string")]
PolicyDirPathToStringFailed,
#[error("Failed to write default policy: {0}")]
WriteDefaultPolicyFailed(#[source] io::Error),
#[error("Failed to read OPA policy file: {0}")]
ReadPolicyFileFailed(#[source] io::Error),
#[error("Failed to write OPA policy to file: {0}")]
WritePolicyFileFailed(#[source] io::Error),
#[error("Failed to load policy: {0}")]
LoadPolicyFailed(#[source] anyhow::Error),
#[error("Policy evaluation denied for {policy_id}")]
PolicyDenied { policy_id: String },
#[error("Serde json error: {0}")]
SerdeJsonError(#[from] serde_json::Error),
#[error("IO error: {0}")]
IOError(#[from] std::io::Error),
#[error("Base64 decode OPA policy string failed: {0}")]
Base64DecodeFailed(#[source] base64::DecodeError),
#[error("Illegal policy id. Only support alphabet, numeric, `-` or `_`")]
InvalidPolicyId,
#[error("Failed to load reference data: {0}")]
LoadReferenceDataFailed(#[source] anyhow::Error),
#[error("Failed to set input data: {0}")]
SetInputDataFailed(#[source] anyhow::Error),
#[error("Failed to evaluate policy: {0}")]
EvalPolicyFailed(#[source] anyhow::Error),
#[error("json serialization failed: {0}")]
JsonSerializationFailed(#[source] anyhow::Error),
}

impl PolicyEngineType {
pub fn to_policy_engine(&self, work_dir: &Path) -> Result<Box<dyn PolicyEngine + Send + Sync>> {
match self {
Expand All @@ -37,13 +73,13 @@ pub trait PolicyEngine {
reference_data_map: HashMap<String, Vec<String>>,
input: String,
policy_ids: Vec<String>,
) -> Result<HashMap<String, PolicyDigest>>;
) -> Result<HashMap<String, PolicyDigest>, RegoError>;

async fn set_policy(&mut self, policy_id: String, policy: String) -> Result<()>;
async fn set_policy(&mut self, policy_id: String, policy: String) -> Result<(), RegoError>;

/// The result is a map. The key is the policy id, and the
/// value is the digest of the policy (using **Sha384**).
async fn list_policies(&self) -> Result<HashMap<String, PolicyDigest>>;
async fn list_policies(&self) -> Result<HashMap<String, PolicyDigest>, RegoError>;

async fn get_policy(&self, policy_id: String) -> Result<String>;
async fn get_policy(&self, policy_id: String) -> Result<String, RegoError>;
}
Original file line number Diff line number Diff line change
Expand Up @@ -2,40 +2,39 @@
// Licensed under the Apache License, Version 2.0, see LICENSE for details.
// SPDX-License-Identifier: Apache-2.0

use anyhow::{anyhow, bail, Context, Result};
use anyhow::{Context, Result};
use async_trait::async_trait;
use base64::Engine;
use sha2::{Digest, Sha384};
use std::collections::HashMap;
use std::fs;
use std::path::PathBuf;

use super::{PolicyDigest, PolicyEngine};
use super::{RegoError, PolicyDigest, PolicyEngine};

#[derive(Debug, Clone)]
pub struct OPA {
policy_dir_path: PathBuf,
}

impl OPA {
pub fn new(work_dir: PathBuf) -> Result<Self> {
pub fn new(work_dir: PathBuf) -> Result<Self, RegoError> {
let mut policy_dir_path = work_dir;

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))?;
fs::create_dir_all(&policy_dir_path).map_err(RegoError::CreatePolicyDirFailed)?;
}

let mut default_policy_path = PathBuf::from(
&policy_dir_path
.to_str()
.ok_or_else(|| anyhow!("Policy DirPath to string failed"))?,
.ok_or_else(|| RegoError::PolicyDirPathToStringFailed)?,
);
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(&default_policy_path, policy)?;
fs::write(&default_policy_path, policy).map_err(RegoError::WriteDefaultPolicyFailed)?;
}

Ok(Self { policy_dir_path })
Expand All @@ -55,21 +54,21 @@ impl PolicyEngine for OPA {
reference_data_map: HashMap<String, Vec<String>>,
input: String,
policy_ids: Vec<String>,
) -> Result<HashMap<String, PolicyDigest>> {
) -> Result<HashMap<String, PolicyDigest>, RegoError> {
let mut res = HashMap::new();

let policy_dir_path = self
.policy_dir_path
.to_str()
.ok_or_else(|| anyhow!("Miss Policy DirPath"))?;
.ok_or_else(|| RegoError::PolicyDirPathToStringFailed)?;

for policy_id in &policy_ids {
let input = input.clone();
let policy_file_path = format!("{policy_dir_path}/{policy_id}.rego");

let policy = tokio::fs::read_to_string(policy_file_path.clone())
.await
.context("Read OPA policy file failed")?;
.map_err(RegoError::ReadPolicyFileFailed)?;

let mut engine = regorus::Engine::new();

Expand All @@ -84,21 +83,29 @@ impl PolicyEngine for OPA {
// Add policy as data
engine
.add_policy(policy_id.clone(), policy)
.context("load policy")?;
.map_err(RegoError::LoadPolicyFailed)?;

let reference_data_map = serde_json::to_string(&reference_data_map)?;
let reference_data_map =
regorus::Value::from_json_str(&format!("{{\"reference\":{reference_data_map}}}"))?;
regorus::Value::from_json_str(&format!("{{\"reference\":{reference_data_map}}}"))
.map_err(RegoError::JsonSerializationFailed)?;
engine
.add_data(reference_data_map)
.context("load reference data")?;
.map_err(RegoError::LoadReferenceDataFailed)?;

// Add TCB claims as input
engine.set_input_json(&input).context("set input")?;
engine
.set_input_json(&input)
.context("set input")
.map_err(RegoError::SetInputDataFailed)?;

let allow = engine.eval_bool_query("data.policy.allow".to_string(), false)?;
let allow = engine
.eval_bool_query("data.policy.allow".to_string(), false)
.map_err(RegoError::EvalPolicyFailed)?;
if !allow {
bail!("TEE evidence does not pass policy {policy_id}");
return Err(RegoError::PolicyDenied {
policy_id: policy_id.clone(),
});
}

res.insert(policy_id.clone(), policy_hash);
Expand All @@ -107,30 +114,30 @@ impl PolicyEngine for OPA {
Ok(res)
}

async fn set_policy(&mut self, policy_id: String, policy: String) -> Result<()> {
async fn set_policy(&mut self, policy_id: String, policy: String) -> Result<(), RegoError> {
let policy_bytes = base64::engine::general_purpose::URL_SAFE_NO_PAD
.decode(policy)
.context("Base64 decode OPA policy string")?;
.map_err(RegoError::Base64DecodeFailed)?;

if !Self::is_valid_policy_id(&policy_id) {
bail!("illegal policy id. Only support alphabet, numeric, `-` or `_` ");
return Err(RegoError::InvalidPolicyId);
}

let mut policy_file_path = PathBuf::from(
&self
.policy_dir_path
.to_str()
.ok_or_else(|| anyhow!("Policy DirPath to string failed"))?,
.ok_or_else(|| RegoError::PolicyDirPathToStringFailed)?,
);

policy_file_path.push(format!("{}.rego", policy_id));

tokio::fs::write(&policy_file_path, policy_bytes)
.await
.map_err(|e| anyhow!("Write OPA policy to file failed: {:?}", e))
.map_err(RegoError::WritePolicyFileFailed)
}

async fn list_policies(&self) -> Result<HashMap<String, PolicyDigest>> {
async fn list_policies(&self) -> Result<HashMap<String, PolicyDigest>, RegoError> {
let mut policy_ids = Vec::new();
let mut entries = tokio::fs::read_dir(&self.policy_dir_path).await?;
while let Some(entry) = entries.next_entry().await? {
Expand All @@ -150,7 +157,7 @@ impl PolicyEngine for OPA {
let policy_file_path = self.policy_dir_path.join(format!("{id}.rego"));
let policy = tokio::fs::read(policy_file_path)
.await
.context("Read OPA policy file failed")?;
.map_err(RegoError::ReadPolicyFileFailed)?;

let mut hasher = Sha384::new();
hasher.update(policy);
Expand All @@ -164,11 +171,11 @@ impl PolicyEngine for OPA {
Ok(policy_list)
}

async fn get_policy(&self, policy_id: String) -> Result<String> {
async fn get_policy(&self, policy_id: String) -> Result<String, RegoError> {
let policy_file_path = self.policy_dir_path.join(format!("{policy_id}.rego"));
let policy = tokio::fs::read(policy_file_path)
.await
.context("Read OPA policy file failed")?;
.map_err(RegoError::ReadPolicyFileFailed)?;
let base64_policy = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(policy);
Ok(base64_policy)
}
Expand Down

0 comments on commit cc55b51

Please sign in to comment.