From 7628983543e4d5b7f4255fbca88ec7dbb9a41728 Mon Sep 17 00:00:00 2001 From: maceip Date: Tue, 5 May 2026 00:07:17 +0200 Subject: [PATCH] feat: workload attestation --- README.md | 10 ++ .../privateinference/library/oakutil/BUILD | 5 +- .../PeerAttestedClientSessionConfigBuilder.kt | 16 +- ...ttestedClientSessionConfigBuilderModule.kt | 18 ++- .../library/oakutil/jni/BUILD | 1 + .../jni/pi_client_session_config_jni.rs | 146 ++++++++++++++++-- .../library/oakutil/public_keys.proto | 22 ++- 7 files changed, 201 insertions(+), 17 deletions(-) diff --git a/README.md b/README.md index 1ed10772..3007c85d 100644 --- a/README.md +++ b/README.md @@ -52,6 +52,16 @@ You can learn more about Private Compute Services in this downloading of resources to the device with support for a binary transparency log based verification, ensuring these are the official resources provided by Google. +* Confidential inference: Enables Private Compute Core features to send + inference requests to a remote Private Aratea service through PCS. PCS + establishes an Oak encrypted session, authorizes the request, and forwards + opaque request bytes; the server-side inference workload and request schema + are not included in this repository. + +For confidential inference, workload attestation can be configured with Oak +reference values. When present, PCS requires the peer to present matching +workload evidence during Oak session setup and fails closed before sending +inference data if that proof is missing or does not verify. ## Note on dependencies diff --git a/src/com/google/android/as/oss/privateinference/library/oakutil/BUILD b/src/com/google/android/as/oss/privateinference/library/oakutil/BUILD index 3f7758e6..21b52f67 100644 --- a/src/com/google/android/as/oss/privateinference/library/oakutil/BUILD +++ b/src/com/google/android/as/oss/privateinference/library/oakutil/BUILD @@ -23,7 +23,10 @@ package(default_visibility = ["//visibility:public"]) proto_library( name = "public_keys_proto", srcs = ["public_keys.proto"], - deps = ["@com_google_protobuf//:timestamp_proto"], + deps = [ + "//third_party/oak/proto/attestation:reference_value_proto", + "@com_google_protobuf//:timestamp_proto", + ], ) cc_proto_library( diff --git a/src/com/google/android/as/oss/privateinference/library/oakutil/PeerAttestedClientSessionConfigBuilder.kt b/src/com/google/android/as/oss/privateinference/library/oakutil/PeerAttestedClientSessionConfigBuilder.kt index 48113500..c93e1893 100644 --- a/src/com/google/android/as/oss/privateinference/library/oakutil/PeerAttestedClientSessionConfigBuilder.kt +++ b/src/com/google/android/as/oss/privateinference/library/oakutil/PeerAttestedClientSessionConfigBuilder.kt @@ -31,7 +31,9 @@ object PeerAttestedClientSessionConfigBuilder { * * @param publicKeyset The public keyset for attestation verification. * @param clock The clock to use for attestation verification time validation. - * @param An optional publisher to publisher attestation evidence on successful session setups. + * @param publisher Optional publisher for attestation evidence on successful session setups. + * @param workloadReferenceValues Serialized Oak ReferenceValues for workload verification. + * @param workloadAttestationId Attestation ID expected for the workload evidence. */ @JvmStatic @JvmOverloads @@ -39,8 +41,16 @@ object PeerAttestedClientSessionConfigBuilder { publicKeyset: ByteArray, clock: AttestationVerificationClock, publisher: AttestationPublisher? = null, + workloadReferenceValues: ByteArray = ByteArray(0), + workloadAttestationId: String = "", ): OakSessionConfigBuilder { - return nativeGet(publicKeyset, clock, publisher) + return nativeGet( + publicKeyset, + clock, + publisher, + workloadReferenceValues, + workloadAttestationId, + ) } @JvmStatic @@ -48,5 +58,7 @@ object PeerAttestedClientSessionConfigBuilder { publicKeyset: ByteArray, clock: AttestationVerificationClock, publisher: AttestationPublisher?, + workloadReferenceValues: ByteArray, + workloadAttestationId: String, ): OakSessionConfigBuilder } diff --git a/src/com/google/android/as/oss/privateinference/library/oakutil/PeerAttestedClientSessionConfigBuilderModule.kt b/src/com/google/android/as/oss/privateinference/library/oakutil/PeerAttestedClientSessionConfigBuilderModule.kt index 2ad3ed64..0cd44a14 100644 --- a/src/com/google/android/as/oss/privateinference/library/oakutil/PeerAttestedClientSessionConfigBuilderModule.kt +++ b/src/com/google/android/as/oss/privateinference/library/oakutil/PeerAttestedClientSessionConfigBuilderModule.kt @@ -42,13 +42,29 @@ internal object PeerAttestedClientSessionConfigBuilderModule { ): OakSessionConfigBuilder { val validUntil = verificationKeys.validUntil.toJavaInstant() if (validUntil.isBefore(Instant.ofEpochMilli(clock.millisecondsSinceEpoch()))) { - throw AttestationVerificationException("Verification keys have expired since $validUntil") + throw AttestationVerificationException("Verification policy has expired since $validUntil") } + val hasWorkloadReferenceValues = verificationKeys.hasWorkloadReferenceValues() + if (verificationKeys.requireWorkloadReferenceValues && !hasWorkloadReferenceValues) { + throw AttestationVerificationException("Workload reference values are required but missing") + } + if (hasWorkloadReferenceValues && verificationKeys.workloadAttestationId.isBlank()) { + throw AttestationVerificationException( + "Workload attestation ID is required when workload reference values are present" + ) + } + + val workloadReferenceValues = + if (hasWorkloadReferenceValues) verificationKeys.workloadReferenceValues.toByteArray() + else ByteArray(0) + return PeerAttestedClientSessionConfigBuilder.get( verificationKeys.tinkSerializedPublicKeyset.toByteArray(), clock, publisher.getOrNull(), + workloadReferenceValues, + verificationKeys.workloadAttestationId, ) } } diff --git a/src/com/google/android/as/oss/privateinference/library/oakutil/jni/BUILD b/src/com/google/android/as/oss/privateinference/library/oakutil/jni/BUILD index 8b87ad2b..74270487 100644 --- a/src/com/google/android/as/oss/privateinference/library/oakutil/jni/BUILD +++ b/src/com/google/android/as/oss/privateinference/library/oakutil/jni/BUILD @@ -37,6 +37,7 @@ rust_library( "//third_party/oak/oak_time", "//third_party/rust/anyhow/v1:anyhow", # buildcleaner: keep "//third_party/rust/jni/v0_21:jni", # buildcleaner: keep + "@oak_crates_index//:prost", ], alwayslink = 1, ) diff --git a/src/com/google/android/as/oss/privateinference/library/oakutil/jni/pi_client_session_config_jni.rs b/src/com/google/android/as/oss/privateinference/library/oakutil/jni/pi_client_session_config_jni.rs index 3cb9f5ee..c447ccb5 100644 --- a/src/com/google/android/as/oss/privateinference/library/oakutil/jni/pi_client_session_config_jni.rs +++ b/src/com/google/android/as/oss/privateinference/library/oakutil/jni/pi_client_session_config_jni.rs @@ -18,26 +18,44 @@ extern crate alloc; -use alloc::{boxed::Box, format, string::ToString, sync::Arc, vec}; +use alloc::{ + boxed::Box, + format, + string::{String, ToString}, + sync::Arc, + vec, +}; use core::ptr::null_mut; use oak_time::Duration; +use anyhow::Context; use jni::{ - objects::{JByteArray, JClass, JObject, JValue}, + objects::{JByteArray, JClass, JObject, JString, JValue}, sys::{jlong, jobject}, JNIEnv, }; -use oak_attestation_verification::{EventLogVerifier, SessionBindingPublicKeyPolicy}; +use oak_attestation_verification::{ + verifier::{to_attestation_results, verify as verify_attestation}, + create_amd_verifier, create_insecure_verifier, create_intel_tdx_verifier, + EventLogVerifier, SessionBindingPublicKeyPolicy, +}; +use oak_attestation_verification_types::verifier::AttestationVerifier; use oak_crypto::certificate::certificate_verifier::CertificateVerifier; use oak_crypto_tink::signature_verifier::SignatureVerifier; use oak_jni_attestation_publisher::JNIAttestationPublisher; use oak_jni_attestation_verification_clock::JNIClock; use oak_proto_rust::attestation::CERTIFICATE_BASED_ATTESTATION_ID; +use oak_proto_rust::oak::attestation::v1::{ + reference_values, AttestationResults, Endorsements, Evidence, ReferenceValues, + RootLayerReferenceValues, +}; use oak_session::{ attestation::AttestationType, config::SessionConfig, config::SessionConfigBuilder, - handshake::HandshakeType, key_extractor::DefaultBindingKeyExtractor, + handshake::HandshakeType, + key_extractor::{DefaultBindingKeyExtractor, DefaultSigningKeyExtractor, KeyExtractor}, session::AttestationPublisher, }; +use prost::Message; /// Acceptable time period before the certificate validity starts and after it ends that allows /// devices with skewed clocks to validate certificates. @@ -50,6 +68,27 @@ const ALLOWED_CLOCK_SKEW: Duration = Duration::from_seconds(26 * 60 * 60); /// only valid for 216 days. const VALIDITY_LIMIT: Duration = Duration::from_seconds(216 * 24 * 60 * 60); +struct ReferenceValuesAttestationVerifier { + reference_values: ReferenceValues, + clock: Arc, +} + +impl AttestationVerifier for ReferenceValuesAttestationVerifier { + fn verify( + &self, + evidence: &Evidence, + endorsements: &Endorsements, + ) -> anyhow::Result { + let verification_result = verify_attestation( + self.clock.get_time().into_unix_millis(), + evidence, + endorsements, + &self.reference_values, + ); + Ok(to_attestation_results(&verification_result)) + } +} + macro_rules! runtime_exception { ($env:ident, $($msg:expr),*) => { $env.throw_new("java/lang/RuntimeException", format!($($msg,)*)) @@ -72,6 +111,53 @@ pub fn new_java_session_config_builder( .map_err(|e| anyhow::anyhow!("Failed to create object: {e:?}")) } +fn create_oak_containers_attestation_verifier( + clock: T, + reference_values: &ReferenceValues, + root_layer: &RootLayerReferenceValues, +) -> anyhow::Result> { + if root_layer.amd_sev.is_some() { + return Ok(Box::new(create_amd_verifier(clock, reference_values)?)); + } + if root_layer.intel_tdx.is_some() { + return Ok(Box::new(create_intel_tdx_verifier( + clock, + reference_values, + )?)); + } + if root_layer.insecure.is_some() { + return Ok(Box::new(create_insecure_verifier(clock, reference_values)?)); + } + + anyhow::bail!("No supported root layer reference values") +} + +fn create_workload_attestation_components( + clock: T, + reference_values: &ReferenceValues, +) -> anyhow::Result<(Box, Box)> { + match reference_values.r#type.as_ref() { + Some(reference_values::Type::OakContainers(rvs)) => { + let root_layer = rvs.root_layer.as_ref().context("No root layer reference values")?; + Ok(( + create_oak_containers_attestation_verifier(clock, reference_values, root_layer)?, + Box::new(DefaultBindingKeyExtractor {}), + )) + } + Some(reference_values::Type::OakRestrictedKernel(_)) => Ok(( + Box::new(ReferenceValuesAttestationVerifier { + reference_values: reference_values.clone(), + clock: Arc::new(clock), + }), + // Restricted Kernel verification exposes the session binding key + // through the extracted signing key path used by Oak's legacy + // default extractor. + Box::new(DefaultSigningKeyExtractor {}), + )), + _ => anyhow::bail!("Unsupported workload reference value type"), + } +} + #[no_mangle] extern "system" fn Java_com_google_android_as_oss_privateinference_library_oakutil_PeerAttestedClientSessionConfigBuilder_nativeGet( mut env: JNIEnv, @@ -79,12 +165,16 @@ extern "system" fn Java_com_google_android_as_oss_privateinference_library_oakut public_keyset_bytes: JByteArray, java_clock_object: JObject, nullable_attestation_publisher: JObject, + workload_reference_values_bytes: JByteArray, + workload_attestation_id: JString, ) -> jobject { let result = internal_native_get( &mut env, public_keyset_bytes, java_clock_object, nullable_attestation_publisher, + workload_reference_values_bytes, + workload_attestation_id, ); match result { Ok(result) => result, @@ -106,29 +196,61 @@ fn internal_native_get( public_keyset_bytes: JByteArray, java_clock_object: JObject, nullable_attestation_publisher: JObject, + workload_reference_values_bytes: JByteArray, + workload_attestation_id: JString, ) -> anyhow::Result { let jni_clock = JNIClock::new(env, &java_clock_object) .map_err(|e| anyhow::anyhow!("Failed to create JNIClock: {e:?}"))?; let public_keyset_vec = env .convert_byte_array(&public_keyset_bytes) .map_err(|e| anyhow::anyhow!("Failed to convert byte array: {e:?}"))?; + let workload_reference_values_vec = env + .convert_byte_array(&workload_reference_values_bytes) + .map_err(|e| anyhow::anyhow!("Failed to convert workload reference values: {e:?}"))?; + let workload_attestation_id: String = env + .get_string(&workload_attestation_id) + .map_err(|e| anyhow::anyhow!("Failed to convert workload attestation ID: {e:?}"))? + .into(); - let mut certificate_verifier = - CertificateVerifier::new(SignatureVerifier::new(&public_keyset_vec)); - certificate_verifier.set_allowed_clock_skew(ALLOWED_CLOCK_SKEW); - certificate_verifier.set_validity_limit(VALIDITY_LIMIT); + let mut session_config_builder = + SessionConfig::builder(AttestationType::PeerUnidirectional, HandshakeType::NoiseNN); - let policy = SessionBindingPublicKeyPolicy::new(certificate_verifier); + // Use workload verification as a fail-closed replacement for the legacy + // certificate path. Registering both verifiers would allow a peer to omit + // workload evidence and still pass via the default legacy aggregator. + if workload_reference_values_vec.is_empty() { + let mut certificate_verifier = + CertificateVerifier::new(SignatureVerifier::new(&public_keyset_vec)); + certificate_verifier.set_allowed_clock_skew(ALLOWED_CLOCK_SKEW); + certificate_verifier.set_validity_limit(VALIDITY_LIMIT); - let attestation_verifier = EventLogVerifier::new(vec![Box::new(policy)], Arc::new(jni_clock)); + let policy = SessionBindingPublicKeyPolicy::new(certificate_verifier); - let mut session_config_builder = - SessionConfig::builder(AttestationType::PeerUnidirectional, HandshakeType::NoiseNN) + let attestation_verifier = + EventLogVerifier::new(vec![Box::new(policy)], Arc::new(jni_clock)); + + session_config_builder = session_config_builder .add_peer_verifier_with_key_extractor( CERTIFICATE_BASED_ATTESTATION_ID.to_string(), Box::new(attestation_verifier), Box::new(DefaultBindingKeyExtractor {}), ); + } else { + if workload_attestation_id.is_empty() { + anyhow::bail!("Workload attestation ID is required for workload verification"); + } + let reference_values = ReferenceValues::decode(workload_reference_values_vec.as_slice()) + .map_err(|e| anyhow::anyhow!("Failed to decode workload reference values: {e:?}"))?; + let (attestation_verifier, key_extractor) = + create_workload_attestation_components(jni_clock, &reference_values)?; + + session_config_builder = session_config_builder + .add_peer_verifier_with_key_extractor( + workload_attestation_id, + attestation_verifier, + key_extractor, + ); + } if !nullable_attestation_publisher.is_null() { let attestation_publisher: Arc = diff --git a/src/com/google/android/as/oss/privateinference/library/oakutil/public_keys.proto b/src/com/google/android/as/oss/privateinference/library/oakutil/public_keys.proto index 8a9bd9fc..42657197 100644 --- a/src/com/google/android/as/oss/privateinference/library/oakutil/public_keys.proto +++ b/src/com/google/android/as/oss/privateinference/library/oakutil/public_keys.proto @@ -17,6 +17,7 @@ edition = "2024"; package com.google.android.as.oss.privateinference.library.oakutil; import "google/protobuf/timestamp.proto"; +import "third_party/oak/proto/attestation/reference_value.proto"; option java_outer_classname = "PublicKeys"; option java_package = "com.google.android.as.oss.privateinference.library.oakutil.proto"; @@ -25,6 +26,25 @@ message VerificationKeys { // Serialized com.google.crypto.tink.TinkProtoKeysetFormat keyset. bytes tink_serialized_public_keyset = 1; - // Time until keyset is guaranteed to be valid for verification. + // Time until this verification policy is guaranteed to be valid. google.protobuf.Timestamp valid_until = 2; + + // Optional Oak reference values for verifying the peer workload identity. + // + // When set, the client verifies the server's Oak evidence and endorsements + // against these values before sending inference request bytes into the + // encrypted session. + oak.attestation.v1.ReferenceValues workload_reference_values = 3; + + // Attestation ID that the peer must use for workload evidence. + // + // This must match the ID used by the server-side Oak attester, for example a + // Restricted Kernel or Oak Containers attestation ID. + string workload_attestation_id = 4; + + // If true, fail closed when workload_reference_values is not present. + // + // This lets clients distinguish legacy keyset-only verification from a + // deployment that requires workload identity verification. + bool require_workload_reference_values = 5; }