Enable WASM Workloads with Encrypted Layer Support

attestation-agent/coco_keyprovider/Cargo.toml

@@ -18,6 +18,8 @@ tracing.workspace = true
 tracing-subscriber.workspace = true
 futures = "0.3.32"
 jwt-simple.workspace = true
+kbs_protocol = { path = "../kbs_protocol" }
+# log.workspace = true
 protos = { path = "../../protos", default-features = false, features = [
     "grpc",
 ] }

attestation-agent/coco_keyprovider/src/enc_mods/crypto.rs

@@ -36,7 +36,36 @@ impl Display for Algorithm {
     }
 }
 
+fn validate_aes256_inputs(key: &[u8], iv: &[u8], algorithm: &Algorithm) -> Result<()> {
+    let (expected_key_len, expected_iv_len, alg_name) = match algorithm {
+        Algorithm::A256GCM => (32, 12, "A256GCM"),
+        Algorithm::A256CTR => (32, 16, "A256CTR"),
+    };
+
+    if key.len() != expected_key_len {
+        return Err(anyhow!(
+            "Invalid key length for {}: {} bytes (expected {} bytes)",
+            alg_name,
+            key.len(),
+            expected_key_len
+        ));
+    }
+
+    if iv.len() != expected_iv_len {
+        return Err(anyhow!(
+            "Invalid IV length for {}: {} bytes (expected {} bytes)",
+            alg_name,
+            iv.len(),
+            expected_iv_len
+        ));
+    }
+
+    Ok(())
+}
+
 pub fn encrypt(data: &[u8], key: &[u8], iv: &[u8], algorithm: &Algorithm) -> Result<Vec<u8>> {
+    validate_aes256_inputs(key, iv, algorithm)?;
+
     match algorithm {
         Algorithm::A256GCM => {
             use aes_gcm::KeyInit;
@@ -45,6 +74,29 @@ pub fn encrypt(data: &[u8], key: &[u8], iv: &[u8], algorithm: &Algorithm) -> Res
             let nonce = Nonce::from_slice(iv);
             cipher
                 .encrypt(nonce, data.as_ref())
+                .map_err(|e| anyhow!("Encrypt failed: {:?}", e))
+        }
+        Algorithm::A256CTR => {
+            use ctr::cipher::{KeyIvInit, StreamCipher};
+            let mut buf = data.to_vec();
+            let mut cipher = ctr::Ctr128BE::<Aes256>::new(key.into(), iv.into());
+            cipher.apply_keystream(&mut buf);
+            Ok(buf)
+        }
+    }
+}
+
+pub fn decrypt(data: &[u8], key: &[u8], iv: &[u8], algorithm: &Algorithm) -> Result<Vec<u8>> {
+    validate_aes256_inputs(key, iv, algorithm)?;
+
+    match algorithm {
+        Algorithm::A256GCM => {
+            use aes_gcm::{aead::Aead, KeyInit};
+            let decryption_key = Key::<Aes256Gcm>::from_slice(key);
+            let cipher = Aes256Gcm::new(decryption_key);
+            let nonce = Nonce::from_slice(iv);
+            cipher
+                .decrypt(nonce, data.as_ref())
                 .map_err(|e| anyhow!("Decrypt failed: {:?}", e))
         }
         Algorithm::A256CTR => {

attestation-agent/coco_keyprovider/src/enc_mods/kbs.rs

@@ -4,12 +4,58 @@
 //
 
 use anyhow::*;
+use base64::Engine;
 use jwt_simple::prelude::{Claims, Duration, Ed25519KeyPair, EdDSAKeyPairLike};
+use kbs_protocol::{
+    evidence_provider::NativeEvidenceProvider, KbsClientBuilder, KbsClientCapabilities, ResourceUri,
+};
+use log::debug;
 use reqwest::Url;
 use tracing::debug;
 
 const KBS_URL_PATH_PREFIX: &str = "kbs/v0/resource";
 
+pub(crate) async fn get_kek(kbs_addr: &Url, kid: &str) -> Result<Vec<u8>> {
+    let kid = kid.strip_prefix('/').unwrap_or(kid);
+
+    // Construct the resource URI in the format: kbs:///<repository>/<type>/<tag>
+    let resource_uri_str = format!("kbs:///{}", kid);
+    let resource_uri: ResourceUri = resource_uri_str
+        .as_str()
+        .try_into()
+        .map_err(|e| anyhow!("Failed to parse resource URI: {}", e))?;
+
+    let evidence_provider = NativeEvidenceProvider::new()
+        .context("Failed to create evidence provider for attestation")?;
+    let mut kbs_client =
+        KbsClientBuilder::with_evidence_provider(Box::new(evidence_provider), kbs_addr.as_str())
+            .build()
+            .context("Failed to build KBS client")?;
+
+    let mut key = kbs_client.get_resource(resource_uri).await?;
+
+    // If the key is not 32 bytes (256-bit AES key), try base64 decoding
+    // Some KBS implementations return base64-encoded keys or text with newlines
+    if key.len() != 32 {
+        // First, try to interpret as UTF-8 string and trim whitespace
+        let key_str = String::from_utf8_lossy(&key);
+        let trimmed = key_str.trim();
+        let engine = base64::engine::general_purpose::STANDARD;
+        let decoded = engine.decode(trimmed.as_bytes())?;
+
+        if decoded.len() == 32 {
+            key = decoded;
+        } else {
+            bail!(
+                "KBS returned key with invalid decoded length: {} bytes (expected 32 bytes)",
+                decoded.len()
+            );
+        }
+    }
+
+    Ok(key)
+}
+
 /// Register the given key with kid into the kbs. This request will be authorized with a
 /// JWT token, which will be signed by the private_key.
 pub(crate) async fn register_kek(

attestation-agent/coco_keyprovider/src/enc_mods/mod.rs

@@ -16,8 +16,8 @@ use tracing::{debug, info};
 
 use self::{crypto::Algorithm, kbs::register_kek};
 
-mod crypto;
-mod kbs;
+pub mod crypto;
+pub mod kbs;
 
 /// `AnnotationPacket` is what a encrypted image layer's
 /// `org.opencontainers.image.enc.keys.provider.attestation-agent`
@@ -83,7 +83,7 @@ fn parse_input_params(input: &str) -> Result<InputParams> {
     let keyid = map.get("keyid").map(|id| id.to_string());
     let keypath = map.get("keypath").map(|p| p.to_string());
     let algorithm = map
-        .get("keypath")
+        .get("algorithm")
         .map(|alg| (*alg).try_into().unwrap_or_default())
         .unwrap_or_default();
     Ok(InputParams {
@@ -150,7 +150,7 @@ async fn generate_key_parameters(input_params: &InputParams) -> Result<(Vec<u8>,
 
 /// Normalize the given keyid into (kbs addr, key path), s.t.
 /// converting `kbs://...` or `../..` to `(<kbs-addr>, <repository>/<type>/<tag>)`.
-fn normalize_path(keyid: &str) -> Result<(String, String)> {
+pub fn normalize_path(keyid: &str) -> Result<(String, String)> {
     debug!("normalize key id {keyid}");
     let path = keyid.strip_prefix(KBS_RESOURCE_URL_PREFIX).unwrap_or(keyid);
     let values: Vec<&str> = path.split('/').collect();

attestation-agent/coco_keyprovider/src/grpc/mod.rs

@@ -3,7 +3,7 @@
 // SPDX-License-Identifier: Apache-2.0
 //
 
-use crate::enc_mods;
+use crate::enc_mods::{crypto::Algorithm, enc_optsdata_gen_anno, kbs::get_kek, AnnotationPacket};
 use anyhow::*;
 use base64::Engine;
 use jwt_simple::prelude::Ed25519KeyPair;
@@ -89,7 +89,7 @@
             })
             .collect();
 
-        let annotation: String = enc_mods::enc_optsdata_gen_anno(
+        let annotation: String = enc_optsdata_gen_anno(
             (&self.kbs, &self.auth_private_key),
             &engine
                 .decode(optsdata)
@@ -123,13 +123,94 @@
 
     async fn un_wrap_key(
         &self,
-        _request: Request<KeyProviderKeyWrapProtocolInput>,
+        request: Request<KeyProviderKeyWrapProtocolInput>,
     ) -> Result<Response<KeyProviderKeyWrapProtocolOutput>, Status> {
-        debug!("The UnWrapKey API is called...");
-        debug!("UnWrapKey API is unimplemented!");
-        Err(Status::unimplemented(
-            "UnWrapKey API of sample-kbs is unimplemented!",
-        ))
+        let input_string = String::from_utf8(
+            request.into_inner().key_provider_key_wrap_protocol_input,
+        )
+        .map_err(|e| {
+            Status::invalid_argument(format!(
+                "key_provider_key_wrap_protocol_input is not legal utf8 string: {e:?}"
+            ))
+        })?;
+
+        let input: KeyProviderInput = serde_json::from_str::<KeyProviderInput>(&input_string)
+            .map_err(|e| {
+                Status::invalid_argument(format!("parse key provider input failed: {e:?}"))
+            })?;
+
+        let annotation_base64 = input.keyunwrapparams.annotation.ok_or_else(|| {
+            Status::invalid_argument("illegal keyunwrapparams without annotation")
+        })?;
+
+        let engine = base64::engine::general_purpose::STANDARD;
+        let annotation_bytes = engine.decode(annotation_base64).map_err(|e| {
+            Status::invalid_argument(format!("base64 decode annotation failed: {e:?}"))
+        })?;
+
+        let annotation: AnnotationPacket =
+            serde_json::from_slice(&annotation_bytes).map_err(|e| {
+                Status::invalid_argument(format!("parse annotation packet failed: {e:?}"))
+            })?;
+
+        let kbs_url = self.kbs.as_ref().ok_or_else(|| {
+            Status::internal("KBS URL not configured. Please provide KBS URL to keyprovider.")
+        })?;
+
+        let (kbs_addr, kbs_path) = crate::enc_mods::normalize_path(&annotation.kid)
+            .map_err(|e| Status::internal(format!("Failed to normalize key path: {:?}", e)))?;
+
+        let kbs_url_with_addr = if kbs_addr.is_empty() {
+            kbs_url.clone()
+        } else {
+            kbs_addr
+                .parse::<Url>()
+                .or_else(|_| format!("{}://{}", kbs_url.scheme(), kbs_addr).parse::<Url>())
+                .map_err(|_| {
+                    Status::internal(format!("Failed to parse KBS address: {}", kbs_addr))
+                })?
+        };
+
+        let kek = get_kek(&kbs_url_with_addr, &kbs_path).await.map_err(|e| {
+            error!(
+                "Failed to get KEK from KBS for kid={}: {:?}",
+                annotation.kid, e
+            );
+            Status::internal("Failed to get KEK from KBS")
+        })?;
+
+        let wrapped_data = engine
+            .decode(&annotation.wrapped_data)
+            .map_err(|e| Status::internal(format!("base64 decode wrapped_data failed: {e:?}")))?;
+
+        let iv = engine
+            .decode(&annotation.iv)
+            .map_err(|e| Status::internal(format!("base64 decode iv failed: {e:?}")))?;
+
+        let wrap_type: Algorithm = annotation
+            .wrap_type
+            .parse()
+            .map_err(|e| Status::internal(format!("Failed to parse wrap_type: {:?}", e)))?;
+
+        let optsdata = crate::enc_mods::crypto::decrypt(&wrapped_data, &kek, &iv, &wrap_type)
+            .map_err(|e| {
+                error!("Failed to decrypt key for kid={}: {:?}", annotation.kid, e);
+                Status::internal("Failed to decrypt key")
+            })?;
+
+        let output_struct = KeyUnwrapOutput {
+            keyunwrapresults: KeyUnwrapResults { optsdata },
+        };
+        let output = serde_json::to_string(&output_struct)
+            .map_err(|e| Status::internal(format!("serde json failed: {e:?}")))?
+            .as_bytes()
+            .to_vec();
+
+        let reply = KeyProviderKeyWrapProtocolOutput {
+            key_provider_key_wrap_protocol_output: output,
+        };
+
+        Result::Ok(Response::new(reply))
     }
 }
 

attestation-agent/docker/Dockerfile.keyprovider

@@ -24,7 +24,7 @@ RUN apt-get update && apt-get install -y \
 	libbtrfs-dev \
 	libdevmapper-dev \
 	pkg-config
-RUN git clone https://github.com/containers/skopeo $GOPATH/src/github.com/containers/skopeo 
+RUN git clone https://github.com/containers/skopeo $GOPATH/src/github.com/containers/skopeo
 WORKDIR $GOPATH/src/github.com/containers/skopeo
 # The dependency on skopeo is quite fragile as there are several versions of
 # the project that would generate an encrypted image with a gzip header that

attestation-agent/kbs_protocol/src/client/rcar_client.rs

@@ -414,7 +414,7 @@ mod test {
     use tokio::io::AsyncBufReadExt;
 
     use crate::{
-        evidence_provider::NativeEvidenceProvider, Error, KbsClientBuilder, KbsClientCapabilities,
+        evidence_provider::MockedEvidenceProvider, Error, KbsClientBuilder, KbsClientCapabilities,
     };
 
     use crate::client::rcar_client::{
@@ -512,7 +512,7 @@ mod test {
         let port = kbs.get_host_port_ipv4(8085).await.expect("get port");
         let kbs_host_url = format!("http://127.0.0.1:{port}");
 
-        let evidence_provider = Box::new(NativeEvidenceProvider::new().unwrap());
+        let evidence_provider = Box::new(MockedEvidenceProvider::default());
         let mut client = KbsClientBuilder::with_evidence_provider(evidence_provider, &kbs_host_url)
             .build()
             .expect("client create");

attestation-agent/kbs_protocol/src/evidence_provider/mock.rs

@@ -5,19 +5,34 @@
 
 use async_trait::async_trait;
 use attester::TeeEvidence;
+use base64::Engine;
 use kbs_types::Tee;
+use serde::{Deserialize, Serialize};
 
 use super::EvidenceProvider;
 
 use crate::Result;
 
+#[derive(Serialize, Deserialize, Debug)]
+struct SampleQuote {
+    svn: String,
+    report_data: String,
+}
+
 #[derive(Default)]
 pub struct MockedEvidenceProvider {}
 
 #[async_trait]
 impl EvidenceProvider for MockedEvidenceProvider {
-    async fn primary_evidence(&self, _runtime_data: Vec<u8>) -> Result<TeeEvidence> {
-        Ok("test evidence".into())
+    async fn primary_evidence(&self, runtime_data: Vec<u8>) -> Result<TeeEvidence> {
+        let evidence = SampleQuote {
+            svn: "1".to_string(),
+            report_data: base64::engine::general_purpose::STANDARD.encode(runtime_data),
+        };
+
+        serde_json::to_value(&evidence).map_err(|e| {
+            crate::Error::GetEvidence(format!("Serialize sample evidence failed: {e}"))
+        })
     }
 
     async fn get_additional_evidence(&self, _runtime_data: Vec<u8>) -> Result<String> {

image-rs/src/decoder/mod.rs

@@ -119,6 +119,11 @@ impl TryFrom<&str> for Compression {
             media_type_str = manifest::IMAGE_LAYER_GZIP_MEDIA_TYPE;
         }
 
+        // Handle WASM media types - WASM layers are typically uncompressed
+        if media_type_str == oci_client::manifest::WASM_LAYER_MEDIA_TYPE {
+            return Ok(Compression::Uncompressed);
+        }
+
         let media_type = MediaType::from(media_type_str);
 
         let decoder = match media_type {