feat(oak_proxy): verify container image digests in Confidential Space

oak_attestation_gcp/src/policy.rs

@@ -212,9 +212,20 @@ impl ConfidentialSpacePolicy {
                             workload_endorsement,
                             ref_value,
                         ),
-                        None => {
-                            Err(ConfidentialSpaceVerificationError::MissingWorkloadEndorsementError)
-                        }
+                        None => match ref_value.r#type {
+                            Some(binary_reference_value::Type::Digests(_))
+                            | Some(binary_reference_value::Type::Skip(_)) => {
+                                verify_endorsement_wrapper(
+                                    verification_time,
+                                    &image_reference,
+                                    &SignedEndorsement::default(),
+                                    ref_value,
+                                )
+                            }
+                            _ => Err(
+                                ConfidentialSpaceVerificationError::MissingWorkloadEndorsementError,
+                            ),
+                        },
                     }
                 }
                 WorkloadReferenceValues::ContainerImageReferencePrefix(container_image_reference_prefix) => {

oak_proxy/README.md

@@ -207,6 +207,12 @@ configuration file passed via the `--config` command-line argument.
 
 - `server_proxy_url`: The WebSocket `Url` of the server proxy.
 
+### Client CLI Flags
+
+- `--http-error-on-fail`: (Optional, Boolean) If set, the client proxy returns an **HTTP 502 Bad Gateway** response to the connected application if the handshake or attestation with the server proxy fails.
+  - **Default**: `false` (Connection is dropped/reset silently on failure).
+  - **Use Case**: Useful when the client application is an HTTP user agent (like a browser or `curl`) that can render a helpful error message instead of a generic connection error.
+
 ### Server-Specific Options
 
 - `backend_address`: The `SocketAddr` of the final backend application where
@@ -238,6 +244,7 @@ Confidential Space attestations.
 
 **`server.toml`**
 
+
 ```toml
 listen_address = "127.0.0.1:8081"
 backend_address = "127.0.0.1:8080"
@@ -249,9 +256,57 @@ type = "confidential_space"
 # Verify the client's Confidential Space attestation.
 [[attestation_verifiers]]
 type = "confidential_space"
-root_certificate_pem_path = "/path/to/gcp_root.pem"
+# Download from https://confidentialcomputing.googleapis.com/.well-known/confidential_space_root.crt
+root_certificate_pem_path = "/path/to/confidential_space_root.crt"
+# Optional: Restrict allowed container images
+# expected_image_digests = ["sha256:e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855"]
+# Optional: Use a signed policy for allowed images
+# signed_policy_path = "/path/to/policy.json"
+# policy_signature_path = "/path/to/policy.sig"
+# policy_public_key_pem_path = "/path/to/key.pem"
+```
+
+### Signed Policy Format
+
+If using `signed_policy_path`, the policy file must be a valid **in-toto Statement** (v1).
+
+Example `policy.json`:
+
+```json
+{
+  "_type": "https://in-toto.io/Statement/v1",
+  "subject": [
+    {
+      "name": "gcr.io/my-project/my-image",
+      "digest": {
+        "sha256": "cc564c5f64a18fc6e53dd737ec19774b68b609eb59cac4f13dcfd25cb61f3f68"
+      }
+    },
+    {
+      "name": "gcr.io/my-project/other-image",
+      "digest": {
+        "sha256": "33e3d4a3a7b79ec43c89803c1dbb1e15317b72f78e9e2ed86d8ff9c1d1431380"
+      }
+    }
+  ],
+  "predicateType": "https://example.com/custom/v1",
+  "predicate": {}
+}
 ```
 
+**Schema Details:**
+
+- **`_type`**: Required. Must be `https://in-toto.io/Statement/v1`.
+- **`subject`**: Required. A list of artifacts covered by this policy.
+  - **`name`**: (Optional) Human-readable name or URI of the artifact (e.g., image URL).
+  - **`digest`**: Required. A map of cryptographic digests.
+    - **`sha256`**: Required. The SHA-256 digest of the container image (hex-encoded).
+- **`predicateType`**: Required. URI identifying the policy type (e.g., `https://example.com/custom/v1`).
+- **`predicate`**: (Optional) Additional policy data (ignored by `oak_proxy` for digest verification).
+
+This format follows the [in-toto Statement Specification](https://github.com/in-toto/attestation/blob/main/spec/v1/statement.md).
+
+
 **`client.toml`**
 
 ```toml
@@ -265,7 +320,8 @@ type = "confidential_space"
 # Verify the server's Confidential Space attestation.
 [[attestation_verifiers]]
 type = "confidential_space"
-root_certificate_pem_path = "/path/to/gcp_root.pem"
+# Download from https://confidentialcomputing.googleapis.com/.well-known/confidential_space_root.crt
+root_certificate_pem_path = "/path/to/confidential_space_root.crt"
 ```
 
 ## Extending Attestation

oak_proxy/client/src/main.rs

@@ -24,7 +24,10 @@ use oak_proxy_lib::{
     websocket::{read_message, write_message},
 };
 use oak_session::{ClientSession, ProtocolEngine, Session};
-use tokio::net::{TcpListener, TcpStream};
+use tokio::{
+    io::AsyncWriteExt,
+    net::{TcpListener, TcpStream},
+};
 use url::Url;
 
 #[derive(Parser, Debug)]
@@ -43,13 +46,19 @@ struct Args {
     /// This will override the value in the config file.
     #[arg(long, env = "OAK_PROXY_SERVER_URL")]
     server_proxy_url: Option<Url>,
+
+    /// If set, returns an HTTP 502 Bad Gateway response with error details to the client when the
+    /// upstream handshake or attestation fails, instead of silently dropping the connection.
+    /// Useful when the client is an HTTP user agent.
+    #[arg(long)]
+    http_error_on_fail: bool,
 }
 
 #[tokio::main]
 async fn main() -> anyhow::Result<()> {
     env_logger::init();
 
-    let Args { mut config, listen_address, server_proxy_url } = Args::parse();
+    let Args { mut config, listen_address, server_proxy_url, http_error_on_fail } = Args::parse();
 
     // The command-line arguments override the values from the config file.
     if let Some(listen_address) = listen_address {
@@ -72,43 +81,78 @@ async fn main() -> anyhow::Result<()> {
         log::info!("[Client] Accepted connection from {}", peer_address);
         let config = config.clone();
         tokio::spawn(async move {
-            if let Err(err) = handle_connection(stream, &config).await {
+            if let Err(err) = handle_connection(stream, &config, http_error_on_fail).await {
                 log::error!("[Client] Error handling connection: {:?}", err);
             }
         });
     }
 }
 
-async fn handle_connection(app_stream: TcpStream, config: &ClientConfig) -> anyhow::Result<()> {
-    let server_proxy_url =
-        config.server_proxy_url.as_ref().context("server_proxy_url wasn't set")?;
-    let (mut server_proxy_stream, _) = tokio_tungstenite::connect_async(server_proxy_url).await?;
-    log::info!("[Client] Connected to server proxy at {}", server_proxy_url);
-
-    let client_config = config::build_session_config(
-        &config.attestation_generators,
-        &config.attestation_verifiers,
-    )?;
-    let mut session = ClientSession::create(client_config)?;
-
-    // Handshake
-    while !session.is_open() {
-        if let Some(request) = session.get_outgoing_message()? {
-            write_message(&mut server_proxy_stream, &request).await?;
-        }
+async fn handle_connection(
+    mut app_stream: TcpStream,
+    config: &ClientConfig,
+    http_error_on_fail: bool,
+) -> anyhow::Result<()> {
+    // Connection and Handshake
+    let setup_result = async {
+        let server_proxy_url =
+            config.server_proxy_url.as_ref().context("server_proxy_url wasn't set")?;
+        let (mut server_proxy_stream, _) =
+            tokio_tungstenite::connect_async(server_proxy_url).await?;
+        log::info!("[Client] Connected to server proxy at {}", server_proxy_url);
+
+        let client_config = config::build_session_config(
+            &config.attestation_generators,
+            &config.attestation_verifiers,
+        )?;
+        let mut session = ClientSession::create(client_config)?;
+
+        while !session.is_open() {
+            if let Some(request) = session.get_outgoing_message()? {
+                write_message(&mut server_proxy_stream, &request).await?;
+            }
 
-        if !session.is_open() {
-            let response = read_message(&mut server_proxy_stream).await?;
-            session.put_incoming_message(response)?;
+            if !session.is_open() {
+                let response = read_message(&mut server_proxy_stream).await?;
+                session.put_incoming_message(response)?;
+            }
+        }
+        Ok((session, server_proxy_stream))
+    }
+    .await;
+
+    match setup_result {
+        Ok((session, server_proxy_stream)) => {
+            log::info!("[Client] Oak Session established with server proxy.");
+            proxy::<
+                ClientSession,
+                oak_proto_rust::oak::session::v1::SessionResponse,
+                oak_proto_rust::oak::session::v1::SessionRequest,
+            >(
+                PeerRole::Client,
+                session,
+                app_stream,
+                server_proxy_stream,
+                config.keep_alive_interval,
+            )
+            .await
+        }
+        Err(err) => {
+            if http_error_on_fail {
+                let error_msg = format!("Attestation/Handshake Failed: {:#}", err);
+                let body = format!("[Oak-Proxy] {}", error_msg);
+                let response = format!(
+                "HTTP/1.1 502 Bad Gateway\r\nContent-Type: text/plain\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}",
+                body.len(),
+                body
+            );
+                if let Err(write_err) = app_stream.write_all(response.as_bytes()).await {
+                    log::warn!("Failed to write HTTP error response: {:?}", write_err);
+                }
+                let _ = app_stream.flush().await;
+                let _ = app_stream.shutdown().await;
+            }
+            Err(err)
         }
     }
-
-    log::info!("[Client] Oak Session established with server proxy.");
-
-    proxy::<
-        ClientSession,
-        oak_proto_rust::oak::session::v1::SessionResponse,
-        oak_proto_rust::oak::session::v1::SessionRequest,
-    >(PeerRole::Client, session, app_stream, server_proxy_stream, config.keep_alive_interval)
-    .await
 }

oak_proxy/lib/BUILD

@@ -23,7 +23,7 @@ package(
 
 rust_library(
     name = "oak_proxy_lib",
-    srcs = glob(["src/**/*.rs"]),
+    srcs = glob(["src/**/*.rs"], exclude = ["src/tests.rs"]),
     deps = [
         "//oak_attestation",
         "//oak_attestation_gcp",
@@ -45,6 +45,7 @@ rust_library(
         "@oak_crates_index//:sha2",
         "@oak_crates_index//:tokio",
         "@oak_crates_index//:tokio-tungstenite",
+        "@oak_crates_index//:serde_json",
         "@oak_crates_index//:toml",
         "@oak_crates_index//:tungstenite",
         "@oak_crates_index//:url",
@@ -63,3 +64,12 @@ rust_test(
         "@oak_crates_index//:toml",
     ],
 )
+
+rust_test(
+    name = "oak_proxy_lib_unit_tests",
+    crate = ":oak_proxy_lib",
+    deps = [
+        "@oak_crates_index//:rand", 
+        "@oak_crates_index//:tempfile",
+    ],
+)

oak_proxy/lib/src/config/confidential_space.rs

@@ -16,6 +16,7 @@
 
 use std::sync::Arc;
 
+use anyhow::Context;
 use oak_attestation::public_key::{PublicKeyAttester, PublicKeyEndorser};
 use oak_attestation_gcp::{
     attestation::request_attestation_token,
@@ -25,7 +26,13 @@ use oak_attestation_gcp::{
 use oak_attestation_verification::EventLogVerifier;
 use oak_proto_rust::{
     attestation::CONFIDENTIAL_SPACE_ATTESTATION_ID,
-    oak::attestation::v1::{ConfidentialSpaceEndorsement, ConfidentialSpaceReferenceValues},
+    oak::{
+        attestation::v1::{
+            binary_reference_value, confidential_space_reference_values, BinaryReferenceValue,
+            ConfidentialSpaceEndorsement, ConfidentialSpaceReferenceValues, Digests,
+        },
+        RawDigest as CommonRawDigest,
+    },
 };
 use oak_session::{
     config::SessionConfigBuilder, key_extractor::DefaultBindingKeyExtractor,
@@ -75,17 +82,67 @@ impl ConfidentialSpaceGeneratorParams {
 #[derive(Deserialize, Serialize, Debug, Clone)]
 pub struct ConfidentialSpaceVerifierParams {
     pub root_certificate_pem_path: String,
+    pub expected_image_digests: Option<Vec<String>>,
+    pub signed_policy_path: Option<String>,
+    pub policy_signature_path: Option<String>,
+    pub policy_public_key_pem_path: Option<String>,
 }
 
 impl ConfidentialSpaceVerifierParams {
+    fn get_allowed_digests(&self) -> anyhow::Result<Vec<Vec<u8>>> {
+        let mut allowed_digests = Vec::new();
+
+        if let Some(digests) = &self.expected_image_digests {
+            allowed_digests.extend(digests.clone());
+        }
+
+        if let (Some(policy_path), Some(sig_path), Some(key_path)) = (
+            &self.signed_policy_path,
+            &self.policy_signature_path,
+            &self.policy_public_key_pem_path,
+        ) {
+            let policy = crate::config::signed_policy::SignedPolicy::load(
+                std::path::Path::new(policy_path),
+                std::path::Path::new(sig_path),
+                std::path::Path::new(key_path),
+            )?;
+            allowed_digests.extend(policy.allowed_digests);
+        }
+
+        allowed_digests
+            .into_iter()
+            .map(|hex_digest| hex::decode(hex_digest).context("decoding hex digest"))
+            .collect()
+    }
+
     pub fn apply(&self, builder: SessionConfigBuilder) -> anyhow::Result<SessionConfigBuilder> {
         let root_pem = std::fs::read_to_string(&self.root_certificate_pem_path)
             .expect("could not read root certificate");
 
-        let reference_values = ConfidentialSpaceReferenceValues {
-            root_certificate_pem: root_pem,
-            r#container_image: None,
+        let allowed_digests = self.get_allowed_digests()?;
+
+        let container_image = if allowed_digests.is_empty() {
+            None
+        } else {
+            let raw_digests: Vec<CommonRawDigest> = allowed_digests
+                .into_iter()
+                .map(|digest| CommonRawDigest { sha2_256: digest, ..Default::default() })
+                .collect();
+
+            let digests = Digests {
+                #[allow(deprecated)]
+                digests: raw_digests,
+            };
+            let binary_reference_value = BinaryReferenceValue {
+                r#type: Some(binary_reference_value::Type::Digests(digests)),
+            };
+            Some(confidential_space_reference_values::ContainerImage::ImageReferenceValue(
+                binary_reference_value,
+            ))
         };
+
+        let reference_values =
+            ConfidentialSpaceReferenceValues { root_certificate_pem: root_pem, container_image };
         let policy = confidential_space_policy_from_reference_values(&reference_values)?;
         let attestation_verifier = EventLogVerifier::new(
             vec![Box::new(policy)],
@@ -100,3 +157,7 @@ impl ConfidentialSpaceVerifierParams {
         ))
     }
 }
+
+#[cfg(test)]
+#[path = "confidential_space_tests.rs"]
+mod tests;