Added checks for AS and KBS policy setting

Signed-off-by: Jiale Zhang <[email protected]>

attestation-service/src/policy_engine/mod.rs

@@ -35,6 +35,8 @@ pub enum PolicyError {
     Base64DecodeFailed(#[from] base64::DecodeError),
     #[error("Illegal policy id. Only support alphabet, numeric, `-` or `_`")]
     InvalidPolicyId,
+    #[error("Illegal policy: {0}")]
+    InvalidPolicy(#[source] anyhow::Error),
     #[error("Failed to load reference data: {0}")]
     LoadReferenceDataFailed(#[source] anyhow::Error),
     #[error("Failed to set input data: {0}")]

attestation-service/src/policy_engine/opa/mod.rs

@@ -122,6 +122,16 @@ impl PolicyEngine for OPA {
             return Err(PolicyError::InvalidPolicyId);
         }
 
+        // Check if the policy is valid
+        {
+            let policy_content = String::from_utf8(policy_bytes.clone())
+                .map_err(|e| PolicyError::InvalidPolicy(e.into()))?;
+            let mut engine = regorus::Engine::new();
+            engine
+                .add_policy(policy_id.clone(), policy_content)
+                .map_err(PolicyError::InvalidPolicy)?;
+        }
+
         let mut policy_file_path = PathBuf::from(
             &self
                 .policy_dir_path

deps/verifier/src/snp/mod.rs

@@ -37,7 +37,7 @@ const LOADER_SPL_OID: Oid<'static> = oid!(1.3.6 .1 .4 .1 .3704 .1 .3 .1);
 const KDS_CERT_SITE: &str = "https://kdsintf.amd.com";
 const KDS_VCEK: &str = "/vcek/v1";
 
-/// Attestation report versions supported 
+/// Attestation report versions supported
 const REPORT_VERSION_MIN: u32 = 2;
 const REPORT_VERSION_MAX: u32 = 3;
 
@@ -110,7 +110,9 @@ impl Verifier for Snp {
 
         // See Trustee Issue#589 https://github.com/confidential-containers/trustee/issues/589
         if report.version < REPORT_VERSION_MIN || report.version > REPORT_VERSION_MAX {
-            return Err(anyhow!("Unexpected attestation report version. Check SNP Firmware ABI specification"));
+            return Err(anyhow!(
+                "Unexpected attestation report version. Check SNP Firmware ABI specification"
+            ));
         }
 
         if report.vmpl != 0 {

kbs/src/policy_engine/error.rs

@@ -33,4 +33,7 @@ pub enum KbsPolicyEngineError {
 
     #[error("Set Policy request is illegal for {0}")]
     IllegalSetPolicyRequest(&'static str),
+
+    #[error("Failed to set policy, illegal policy: {0}")]
+    InvalidPolicy(#[source] anyhow::Error),
 }

kbs/src/policy_engine/opa/mod.rs

@@ -61,6 +61,16 @@ impl PolicyEngineInterface for Opa {
     async fn set_policy(&mut self, policy: &str) -> Result<(), KbsPolicyEngineError> {
         let policy_bytes = base64::engine::general_purpose::URL_SAFE_NO_PAD.decode(policy)?;
 
+        // Check if the policy is valid
+        {
+            let policy_content = String::from_utf8(policy_bytes.clone())
+                .map_err(|e| KbsPolicyEngineError::InvalidPolicy(e.into()))?;
+            let mut engine = regorus::Engine::new();
+            engine
+                .add_policy(String::from("default"), policy_content)
+                .map_err(KbsPolicyEngineError::InvalidPolicy)?;
+        }
+
         tokio::fs::write(&self.policy_path, policy_bytes).await?;
 
         Ok(())
@@ -153,6 +163,13 @@ mod tests {
             res.err().unwrap(),
             KbsPolicyEngineError::IOError(_)
         ));
+
+        // Illegal policy
+        let res = set_policy_from_file(&mut opa, "test/data/policy_invalid_1.rego").await;
+        assert!(matches!(
+            res.err().unwrap(),
+            KbsPolicyEngineError::InvalidPolicy(_)
+        ));
     }
 
     #[rstest]
@@ -167,13 +184,6 @@ mod tests {
         1,
         Err(KbsPolicyEngineError::ResourcePathError)
     )]
-    #[case(
-        "test/data/policy_invalid_1.rego",
-        "my_repo/Alice/key",
-        "Alice",
-        1,
-        Err(KbsPolicyEngineError::PolicyLoadError)
-    )]
     #[case(
         "test/data/policy_invalid_2.rego",
         "my_repo/Alice/key",