feat: workload attestation

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
 

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(

src/com/google/android/as/oss/privateinference/library/oakutil/PeerAttestedClientSessionConfigBuilder.kt

@@ -31,22 +31,34 @@ 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
   fun get(
     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
   private external fun nativeGet(
     publicKeyset: ByteArray,
     clock: AttestationVerificationClock,
     publisher: AttestationPublisher?,
+    workloadReferenceValues: ByteArray,
+    workloadAttestationId: String,
   ): OakSessionConfigBuilder
 }

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,
     )
   }
 }

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,
 )

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<dyn oak_time::Clock>,
+}
+
+impl AttestationVerifier for ReferenceValuesAttestationVerifier {
+    fn verify(
+        &self,
+        evidence: &Evidence,
+        endorsements: &Endorsements,
+    ) -> anyhow::Result<AttestationResults> {
+        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,19 +111,70 @@ pub fn new_java_session_config_builder(
         .map_err(|e| anyhow::anyhow!("Failed to create object: {e:?}"))
 }
 
+fn create_oak_containers_attestation_verifier<T: oak_time::Clock + 'static>(
+    clock: T,
+    reference_values: &ReferenceValues,
+    root_layer: &RootLayerReferenceValues,
+) -> anyhow::Result<Box<dyn AttestationVerifier>> {
+    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<T: oak_time::Clock + 'static>(
+    clock: T,
+    reference_values: &ReferenceValues,
+) -> anyhow::Result<(Box<dyn AttestationVerifier>, Box<dyn KeyExtractor>)> {
+    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,
     _class: JClass,
     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<jobject> {
     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<dyn AttestationPublisher> =

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;
 }