Add runtime attestation support in image rs

.github/workflows/image_rs_build.yml

@@ -121,21 +121,21 @@ jobs:
         run: |
           sudo -E PATH=$PATH -s cargo test -p image-rs --features default
 
-      - name: Run cargo test - kata-cc (rust-tls version) with keywrap-grpc + keywrap-jwe
+      - name: Run cargo test - kata-cc (rust-tls version) with keywrap-grpc + keywrap-jwe + runtime-attestation-grpc
         run: |
-          sudo -E PATH=$PATH -s cargo test -p image-rs --no-default-features --features=encryption-ring,keywrap-grpc,snapshot-overlayfs,signature-cosign-rustls,signature-simple,getresource,oci-distribution/rustls-tls,keywrap-jwe
+          sudo -E PATH=$PATH -s cargo test -p image-rs --no-default-features --features=encryption-ring,keywrap-grpc,snapshot-overlayfs,signature-cosign-rustls,signature-simple,getresource,oci-distribution/rustls-tls,keywrap-jwe,runtime-attestation-grpc
 
-      - name: Run cargo test - kata-cc (native-tls version) with keywrap-grpc + keywrap-jwe
+      - name: Run cargo test - kata-cc (native-tls version) with keywrap-grpc + keywrap-jwe + runtime-attestation-grpc
         run: |
-          sudo -E PATH=$PATH -s cargo test -p image-rs --no-default-features --features=encryption-openssl,keywrap-grpc,snapshot-overlayfs,signature-cosign-native,signature-simple,getresource,oci-distribution/native-tls,keywrap-jwe
+          sudo -E PATH=$PATH -s cargo test -p image-rs --no-default-features --features=encryption-openssl,keywrap-grpc,snapshot-overlayfs,signature-cosign-native,signature-simple,getresource,oci-distribution/native-tls,keywrap-jwe,runtime-attestation-grpc
 
-      - name: Run cargo test - kata-cc (rust-tls version) with keywrap-ttrpc (default) + keywrap-jwe
+      - name: Run cargo test - kata-cc (rust-tls version) with keywrap-ttrpc (default) + keywrap-jwe + runtime-attestation-ttrpc
         run: |
-          sudo -E PATH=$PATH -s cargo test -p image-rs --no-default-features --features=kata-cc-rustls-tls,keywrap-jwe
+          sudo -E PATH=$PATH -s cargo test -p image-rs --no-default-features --features=kata-cc-rustls-tls,keywrap-jwe,runtime-attestation-ttrpc
 
-      - name: Run cargo test - kata-cc (native-tls version) with keywrap-ttrpc (default) + keywrap-jwe
+      - name: Run cargo test - kata-cc (native-tls version) with keywrap-ttrpc (default) + keywrap-jwe + runtime-attestation-ttrpc
         run: |
-          sudo -E PATH=$PATH -s cargo test -p image-rs --no-default-features --features=kata-cc-native-tls,keywrap-jwe
+          sudo -E PATH=$PATH -s cargo test -p image-rs --no-default-features --features=kata-cc-native-tls,keywrap-jwe,runtime-attestation-ttrpc
 
       - name: Clean test cache
         run: |

image-rs/Cargo.toml

@@ -197,3 +197,7 @@ getresource = ["lazy_static", "cfg-if"]
 nydus = ["lazy_static", "nydus-api", "nydus-service"]
 
 verity = ["devicemapper"]
+
+runtime-attestation = ["lazy_static", "cfg-if"]
+runtime-attestation-ttrpc = ["dep:ttrpc", "dep:protobuf", "ttrpc-codegen", "runtime-attestation"]
+runtime-attestation-grpc = ["prost", "tonic", "tonic-build", "runtime-attestation"]

image-rs/build.rs

@@ -9,6 +9,12 @@ fn main() -> Result<()> {
     #[cfg(feature = "tonic-build")]
     tonic_build::compile_protos("./protos/getresource.proto").context("tonic build")?;
 
+    #[cfg(feature = "tonic-build")]
+    tonic_build::configure()
+        .build_server(true)
+        .protoc_arg("--experimental_allow_proto3_optional")
+        .compile(&["./protos/attestation-agent.proto"], &["./protos"])?;
+
     #[cfg(feature = "ttrpc-codegen")]
     ttrpc_codegen::Codegen::new()
         .out_dir("./src/resource/kbs/ttrpc_proto")
@@ -23,5 +29,19 @@ fn main() -> Result<()> {
         .run()
         .context("ttrpc build")?;
 
+    #[cfg(feature = "ttrpc-codegen")]
+    ttrpc_codegen::Codegen::new()
+        .out_dir("./src/runtime_attestation/ttrpc_proto")
+        .input("./protos/attestation-agent.proto")
+        .include("./protos")
+        .rust_protobuf()
+        .customize(ttrpc_codegen::Customize {
+            async_all: true,
+            ..Default::default()
+        })
+        .rust_protobuf_customize(ttrpc_codegen::ProtobufCustomize::default().gen_mod_rs(false))
+        .run()
+        .context("ttrpc build")?;
+
     Ok(())
 }

image-rs/protos/attestation-agent.proto

@@ -0,0 +1,25 @@
+syntax = "proto3";
+
+package attestation_agent;
+
+// Extend the dynamic/runtime measurement with given materials. This would change the state
+// of current TEE's status, e.g. TDX's RTMR, (v)TPM's PCR, by adding a record in eventlog.
+message ExtendRuntimeMeasurementRequest {
+    // The domain to which this event entry belongs. This domain is used to distinguish the semantics of log entries in different contexts.
+    string Domain = 1;
+
+    // Concrete operation type that this event entry records.
+    string Operation = 2;
+
+    // Detailed content of the operation that this event entry records.
+    string Content = 3;
+
+    // Which PCR will be extended with the hash of this entry.
+    optional uint64 RegisterIndex = 4;
+}
+
+message ExtendRuntimeMeasurementResponse {}
+
+service AttestationAgentService {
+    rpc ExtendRuntimeMeasurement(ExtendRuntimeMeasurementRequest) returns (ExtendRuntimeMeasurementResponse) {};
+}

image-rs/src/image.rs

@@ -29,6 +29,9 @@ use crate::snapshots::overlay::OverlayFs;
 #[cfg(feature = "nydus")]
 use crate::nydus::{service, utils};
 
+#[cfg(feature = "runtime-attestation")]
+use crate::runtime_attestation;
+
 /// Image security config dir contains important information such as
 /// security policy configuration file and signature verification configuration file.
 /// Therefore, it is necessary to ensure that the directory is stored in a safe place.
@@ -241,13 +244,29 @@ impl ImageClient {
         };
 
         let mut client = PullClient::new(
-            reference,
+            reference.clone(),
             &self.config.work_dir.join("layers"),
             &auth,
             self.config.max_concurrent_download,
         )?;
         let (image_manifest, image_digest, image_config) = client.pull_manifest().await?;
 
+        #[cfg(feature = "runtime-attestation")]
+        if let Err(e) = runtime_attestation::extend_runtime_measurement(
+            "image-rs",
+            "pull_image",
+            &format!(
+                "{}/{}@{}",
+                reference.registry(),
+                reference.repository(),
+                image_digest
+            ),
+        )
+        .await
+        {
+            log::warn!("Failed to extend_runtime_measurement {:?}", e);
+        }
+
         let id = image_manifest.config.digest.clone();
 
         let snapshot = match self.snapshots.get_mut(&self.config.default_snapshot) {

image-rs/src/lib.rs

@@ -16,6 +16,8 @@ pub mod meta_store;
 pub mod nydus;
 pub mod pull;
 pub mod resource;
+#[cfg(feature = "runtime-attestation")]
+pub mod runtime_attestation;
 #[cfg(feature = "signature")]
 pub mod signature;
 pub mod snapshots;

image-rs/src/runtime_attestation/grpc.rs

@@ -0,0 +1,56 @@
+// Copyright (c) 2024 Intel
+//
+// SPDX-License-Identifier: Apache-2.0
+
+use anyhow::*;
+use async_trait::async_trait;
+
+use tonic::transport::Channel;
+
+use super::Client;
+
+use self::attestation_agent::{
+    attestation_agent_service_client::AttestationAgentServiceClient,
+    ExtendRuntimeMeasurementRequest,
+};
+
+mod attestation_agent {
+    #![allow(unknown_lints)]
+    #![allow(clippy::derive_partial_eq_without_eq)]
+    #![allow(clippy::redundant_async_block)]
+    tonic::include_proto!("attestation_agent");
+}
+
+pub const AA_ADDR: &str = "http://127.0.0.1:50002";
+
+pub struct Grpc {
+    inner: AttestationAgentServiceClient<Channel>,
+}
+
+impl Grpc {
+    pub async fn new() -> Result<Self> {
+        let inner = AttestationAgentServiceClient::connect(AA_ADDR).await?;
+        Ok(Self { inner })
+    }
+}
+
+#[async_trait]
+impl Client for Grpc {
+    async fn extend_runtime_measurement(
+        &mut self,
+        domain: &str,
+        operation: &str,
+        content: &str,
+    ) -> Result<()> {
+        let req = tonic::Request::new(ExtendRuntimeMeasurementRequest {
+            domain: domain.to_string(),
+            operation: operation.to_string(),
+            content: content.to_string(),
+            ..Default::default()
+        });
+
+        self.inner.extend_runtime_measurement(req).await?;
+
+        Ok(())
+    }
+}

image-rs/src/runtime_attestation/mod.rs

@@ -0,0 +1,84 @@
+// Copyright (c) 2024 Intel
+//
+// SPDX-License-Identifier: Apache-2.0
+
+use anyhow::*;
+use async_trait::async_trait;
+use log::{info, warn};
+use std::result::Result::Ok;
+use tokio::sync::{Mutex, OnceCell};
+
+#[cfg(feature = "runtime-attestation-grpc")]
+mod grpc;
+
+#[cfg(feature = "runtime-attestation-ttrpc")]
+mod ttrpc;
+
+#[cfg(feature = "runtime-attestation-ttrpc")]
+mod ttrpc_proto;
+
+static RUNTIME_MEASUREMENT: OnceCell<Mutex<Option<RuntimeMeasurement>>> = OnceCell::const_new();
+
+pub struct RuntimeMeasurement {
+    client: Box<dyn Client>,
+}
+
+#[async_trait]
+trait Client: Send + Sync {
+    async fn extend_runtime_measurement(
+        &mut self,
+        domain: &str,
+        operation: &str,
+        content: &str,
+    ) -> Result<()>;
+}
+
+impl RuntimeMeasurement {
+    pub async fn new() -> Result<Self> {
+        let client: Box<dyn Client> = {
+            cfg_if::cfg_if! {
+                if #[cfg(feature = "runtime-attestation-ttrpc")] {
+                    info!("runtime-attestation uses ttrpc");
+                    Box::new(ttrpc::Ttrpc::new().context("ttrpc client init failed")?)
+                } else if #[cfg(feature = "runtime-attestation-grpc")] {
+                    info!("runtime-attestation uses grpc");
+            Box::new(grpc::Grpc::new().await.context("grpc client init failed")?)
+                } else {
+                    compile_error!("`runtime-attestation-ttrpc` or `runtime-attestation-grpc` must be enabled.");
+                }
+            }
+        };
+
+        Ok(Self { client })
+    }
+}
+
+async fn get_runtime_measurement() -> Mutex<Option<RuntimeMeasurement>> {
+    match RuntimeMeasurement::new().await {
+        Ok(runtime_measurement) => Mutex::new(Some(runtime_measurement)),
+        Err(err) => {
+            warn!("Failed to initialize runtime measurement: {:?}", err);
+            Mutex::new(None)
+        }
+    }
+}
+
+pub async fn extend_runtime_measurement(
+    domain: &str,
+    operation: &str,
+    content: &str,
+) -> Result<()> {
+    RUNTIME_MEASUREMENT
+        .get_or_init(get_runtime_measurement)
+        .await
+        .lock()
+        .await
+        .as_mut()
+        .ok_or_else(|| anyhow!("Uninitialized runtime measurement"))?
+        .client
+        .extend_runtime_measurement(domain, operation, content)
+        .await
+        .map_err(|e| anyhow!("Failed to extend runtime measurement: {:?}", e))?;
+
+    Ok(())
+}

image-rs/src/runtime_attestation/ttrpc.rs

@@ -0,0 +1,48 @@
+// Copyright (c) 2024 Intel
+//
+// SPDX-License-Identifier: Apache-2.0
+
+use anyhow::*;
+use async_trait::async_trait;
+
+use super::ttrpc_proto::attestation_agent::ExtendRuntimeMeasurementRequest;
+use super::ttrpc_proto::attestation_agent_ttrpc::AttestationAgentServiceClient;
+use super::Client;
+use ttrpc::context;
+
+const SOCKET_ADDR: &str =
+    "unix:///run/confidential-containers/attestation-agent/attestation-agent.sock";
+
+pub struct Ttrpc {
+    client: AttestationAgentServiceClient,
+}
+
+impl Ttrpc {
+    pub fn new() -> Result<Self> {
+        let inner = ttrpc::asynchronous::Client::connect(SOCKET_ADDR)?;
+        let client = AttestationAgentServiceClient::new(inner);
+        Ok(Self { client })
+    }
+}
+
+#[async_trait]
+impl Client for Ttrpc {
+    async fn extend_runtime_measurement(
+        &mut self,
+        domain: &str,
+        operation: &str,
+        content: &str,
+    ) -> Result<()> {
+        let req = ExtendRuntimeMeasurementRequest {
+            Domain: domain.to_string(),
+            Operation: operation.to_string(),
+            Content: content.to_string(),
+            ..Default::default()
+        };
+        self.client
+            .extend_runtime_measurement(context::with_timeout(50 * 1000 * 1000 * 1000), &req)
+            .await
+            .map_err(|e| anyhow!("extend runtime measurement ttrpc error: {:?}", e))?;
+        Ok(())
+    }
+}