fix(v0.5): verify RA-TLS key possession and key-handover peer

kms/src/onboard_service.rs

@@ -249,11 +249,17 @@ impl Keys {
         domain: &str,
         pccs_url: Option<String>,
     ) -> Result<Self> {
+        let bootstrap = RaClient::new(other_kms_url.to_string(), true)?;
+        let tmp_ca = KmsClient::new(bootstrap).get_temp_ca_cert().await?;
+        let (ra_cert, ra_key) = gen_ra_cert(tmp_ca.temp_ca_cert, tmp_ca.temp_ca_key).await?;
         let attestation_slot = Arc::new(Mutex::new(None::<VerifiedAttestation>));
         let attestation_slot_out = attestation_slot.clone();
         let client = RaClientConfig::builder()
             .tls_no_check(true)
             .remote_uri(other_kms_url.to_string())
+            .tls_built_in_root_certs(false)
+            .tls_client_cert(ra_cert)
+            .tls_client_key(ra_key)
             .cert_validator(Box::new(move |info: Option<CertInfo>| {
                 let Some(info) = info else {
                     bail!("Source KMS did not present a TLS certificate");
@@ -270,13 +276,13 @@ impl Keys {
             .maybe_pccs_url(pccs_url.clone())
             .build()
             .into_client()?;
-        let mut kms_client = KmsClient::new(client);
-
-        let tmp_ca = kms_client.get_temp_ca_cert().await?;
-        let (ra_cert, ra_key) = gen_ra_cert(tmp_ca.temp_ca_cert, tmp_ca.temp_ca_key).await?;
-        let ra_client = RaClient::new_mtls(other_kms_url.into(), ra_cert, ra_key, pccs_url)
-            .context("Failed to create client")?;
-        kms_client = KmsClient::new(ra_client);
+        let info = dstack_client().info().await.context("Failed to get info")?;
+        let keys_res = KmsClient::new(client)
+            .get_kms_key(GetKmsKeyRequest {
+                vm_config: info.vm_config,
+            })
+            .await?;
+        // Only trust attestation from the connection that returned these keys.
         let source_attestation = attestation_slot
             .lock()
             .map_err(|_| anyhow::anyhow!("source attestation mutex poisoned"))?
@@ -285,13 +291,6 @@ impl Keys {
         ensure_kms_allowed(cfg, &source_attestation)
             .await
             .context("Source KMS is not allowed for onboarding")?;
-
-        let info = dstack_client().info().await.context("Failed to get info")?;
-        let keys_res = kms_client
-            .get_kms_key(GetKmsKeyRequest {
-                vm_config: info.vm_config,
-            })
-            .await?;
         if keys_res.keys.len() != 1 {
             return Err(anyhow::anyhow!("Invalid keys"));
         }

ra-rpc/Cargo.toml

@@ -13,6 +13,7 @@ license.workspace = true
 anyhow.workspace = true
 prpc.workspace = true
 rocket = { workspace = true, features = ["mtls"], optional = true }
+rustls = { workspace = true, optional = true }
 serde_json.workspace = true
 tracing.workspace = true
 reqwest = { workspace = true, default-features = false, features = ["rustls-tls", "charset"], optional = true }
@@ -28,5 +29,9 @@ or-panic.workspace = true
 [features]
 default = ["rocket", "client"]
 rocket = ["dep:rocket", "dep:rocket-vsock-listener"]
-client = ["reqwest"]
+client = ["reqwest", "dep:rustls"]
 openapi = ["dep:prost-types"]
+
+[dev-dependencies]
+tokio = { workspace = true, features = ["macros", "rt-multi-thread", "net", "io-util"] }
+tokio-rustls.workspace = true

ra-rpc/src/client.rs

@@ -2,7 +2,7 @@
 //
 // SPDX-License-Identifier: Apache-2.0
 
-use std::time::Duration;
+use std::{sync::Arc, time::Duration};
 
 use anyhow::{bail, Context, Result};
 use prpc::{
@@ -46,15 +46,20 @@ impl RaClientConfig {
     pub fn into_client(self) -> Result<RaClient> {
         let mut builder = Client::builder()
             .tls_sni(true)
-            .danger_accept_invalid_certs(self.tls_no_check)
             .danger_accept_invalid_hostnames(self.tls_no_check_hostname)
             .tls_built_in_root_certs(self.tls_built_in_root_certs)
             .connect_timeout(Duration::from_secs(5))
             .timeout(Duration::from_secs(60));
         if self.cert_validator.is_some() {
             builder = builder.tls_info(true);
         }
-        if let (Some(cert_pem), Some(key_pem)) = (self.tls_client_cert, self.tls_client_key) {
+        if self.tls_no_check {
+            builder = builder.use_preconfigured_tls(attestation_tls_config(
+                self.tls_client_cert.as_deref(),
+                self.tls_client_key.as_deref(),
+            )?);
+        } else if let (Some(cert_pem), Some(key_pem)) = (self.tls_client_cert, self.tls_client_key)
+        {
             let identity_pem = format!("{cert_pem}\n{key_pem}");
             let identity =
                 Identity::from_pem(identity_pem.as_bytes()).context("Failed to parse identity")?;
@@ -75,6 +80,77 @@ impl RaClientConfig {
     }
 }
 
+/// Attestation is checked on the response, but TLS must still prove possession
+/// of the attested certificate key. Reqwest's no-verifier also skips that proof.
+#[derive(Debug)]
+struct AttestationServerVerifier(Arc<rustls::crypto::CryptoProvider>);
+
+impl rustls::client::danger::ServerCertVerifier for AttestationServerVerifier {
+    fn verify_server_cert(
+        &self,
+        _end_entity: &rustls::pki_types::CertificateDer<'_>,
+        _intermediates: &[rustls::pki_types::CertificateDer<'_>],
+        _server_name: &rustls::pki_types::ServerName<'_>,
+        _ocsp_response: &[u8],
+        _now: rustls::pki_types::UnixTime,
+    ) -> Result<rustls::client::danger::ServerCertVerified, rustls::Error> {
+        Ok(rustls::client::danger::ServerCertVerified::assertion())
+    }
+
+    fn verify_tls12_signature(
+        &self,
+        message: &[u8],
+        cert: &rustls::pki_types::CertificateDer<'_>,
+        signature: &rustls::DigitallySignedStruct,
+    ) -> Result<rustls::client::danger::HandshakeSignatureValid, rustls::Error> {
+        rustls::crypto::verify_tls12_signature(
+            message,
+            cert,
+            signature,
+            &self.0.signature_verification_algorithms,
+        )
+    }
+
+    fn verify_tls13_signature(
+        &self,
+        message: &[u8],
+        cert: &rustls::pki_types::CertificateDer<'_>,
+        signature: &rustls::DigitallySignedStruct,
+    ) -> Result<rustls::client::danger::HandshakeSignatureValid, rustls::Error> {
+        rustls::crypto::verify_tls13_signature(
+            message,
+            cert,
+            signature,
+            &self.0.signature_verification_algorithms,
+        )
+    }
+
+    fn supported_verify_schemes(&self) -> Vec<rustls::SignatureScheme> {
+        self.0.signature_verification_algorithms.supported_schemes()
+    }
+}
+
+fn attestation_tls_config(cert: Option<&str>, key: Option<&str>) -> Result<rustls::ClientConfig> {
+    use rustls::pki_types::{pem::PemObject, CertificateDer, PrivateKeyDer};
+    let provider = rustls::crypto::CryptoProvider::get_default()
+        .cloned()
+        .unwrap_or_else(|| Arc::new(rustls::crypto::aws_lc_rs::default_provider()));
+    let builder = rustls::ClientConfig::builder_with_provider(provider.clone())
+        .with_safe_default_protocol_versions()?
+        .dangerous()
+        .with_custom_certificate_verifier(Arc::new(AttestationServerVerifier(provider)));
+    match (cert, key) {
+        (Some(cert), Some(key)) => builder
+            .with_client_auth_cert(
+                CertificateDer::pem_slice_iter(cert.as_bytes()).collect::<Result<Vec<_>, _>>()?,
+                PrivateKeyDer::from_pem_slice(key.as_bytes())?,
+            )
+            .context("invalid RA-TLS client identity"),
+        (None, None) => Ok(builder.with_no_client_auth()),
+        _ => bail!("RA-TLS client identity requires both certificate and key"),
+    }
+}
+
 pub struct RaClient {
     remote_uri: String,
     pccs_url: Option<String>,

ra-rpc/tests/server_key_possession.rs

@@ -0,0 +1,132 @@
+// SPDX-FileCopyrightText: © 2026 Phala Network <dstack@phala.network>
+//
+// SPDX-License-Identifier: Apache-2.0
+
+#![cfg(feature = "client")]
+
+use std::sync::Arc;
+
+use prpc::client::RequestClient;
+use ra_rpc::client::RaClientConfig;
+use ra_tls::{
+    cert::CertRequest,
+    rcgen::{KeyPair, PKCS_ECDSA_P256_SHA256},
+};
+use rustls::{
+    pki_types::PrivatePkcs8KeyDer,
+    server::{ClientHello, ResolvesServerCert},
+    sign::CertifiedKey,
+};
+use tokio::io::{AsyncReadExt, AsyncWriteExt};
+
+#[derive(Debug)]
+struct FixedCertificate(Arc<CertifiedKey>);
+
+impl ResolvesServerCert for FixedCertificate {
+    fn resolve(&self, _: ClientHello<'_>) -> Option<Arc<CertifiedKey>> {
+        Some(self.0.clone())
+    }
+}
+
+#[tokio::test]
+async fn tls12_requires_the_server_certificate_key_even_without_ca_validation() {
+    check_key_possession(&rustls::version::TLS12).await;
+}
+
+#[tokio::test]
+async fn tls13_requires_the_server_certificate_key_even_without_ca_validation() {
+    check_key_possession(&rustls::version::TLS13).await;
+}
+
+async fn check_key_possession(version: &'static rustls::SupportedProtocolVersion) {
+    for correct_key in [true, false] {
+        let key = KeyPair::generate_for(&PKCS_ECDSA_P256_SHA256).unwrap();
+        let cert = CertRequest::builder()
+            .subject("test server")
+            .key(&key)
+            .usage_server_auth(true)
+            .build()
+            .self_signed()
+            .unwrap();
+        let signing_key = if correct_key {
+            key
+        } else {
+            KeyPair::generate_for(&PKCS_ECDSA_P256_SHA256).unwrap()
+        };
+        let ca_key = KeyPair::generate_for(&PKCS_ECDSA_P256_SHA256).unwrap();
+        let ca = CertRequest::builder()
+            .subject("test client CA")
+            .key(&ca_key)
+            .ca_level(0)
+            .build()
+            .self_signed()
+            .unwrap();
+        let client_key = KeyPair::generate_for(&PKCS_ECDSA_P256_SHA256).unwrap();
+        let client_cert = CertRequest::builder()
+            .subject("test client")
+            .key(&client_key)
+            .usage_client_auth(true)
+            .build()
+            .signed_by(&ca, &ca_key)
+            .unwrap();
+        let mut roots = rustls::RootCertStore::empty();
+        roots.add(ca.der().clone()).unwrap();
+        let provider = Arc::new(rustls::crypto::aws_lc_rs::default_provider());
+        let client_verifier = rustls::server::WebPkiClientVerifier::builder_with_provider(
+            Arc::new(roots),
+            provider.clone(),
+        )
+        .build()
+        .unwrap();
+        let signing_key = provider
+            .key_provider
+            .load_private_key(PrivatePkcs8KeyDer::from(signing_key.serialize_der()).into())
+            .unwrap();
+        // A resolver permits deliberately pairing a stolen certificate with
+        // another key, unlike with_single_cert which rejects it up front.
+        let pinned_cert = cert.der().to_vec();
+        let certified = CertifiedKey::new(vec![cert.der().clone()], signing_key);
+        let config = rustls::ServerConfig::builder_with_provider(provider)
+            .with_protocol_versions(&[version])
+            .unwrap()
+            .with_client_cert_verifier(client_verifier)
+            .with_cert_resolver(Arc::new(FixedCertificate(Arc::new(certified))));
+        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
+        let address = listener.local_addr().unwrap();
+        let server = tokio::spawn(async move {
+            let (stream, _) = listener.accept().await.unwrap();
+            let Ok(mut stream) = tokio_rustls::TlsAcceptor::from(Arc::new(config))
+                .accept(stream)
+                .await
+            else {
+                return false;
+            };
+            let mut request = [0; 4096];
+            if !matches!(stream.read(&mut request).await, Ok(n) if n > 0) {
+                return false;
+            }
+            stream
+                .write_all(b"HTTP/1.1 200 OK\r\nContent-Length: 4\r\nConnection: close\r\n\r\nnull")
+                .await
+                .unwrap();
+            true
+        });
+        let client = RaClientConfig::builder()
+            .remote_uri(format!("https://{address}"))
+            .tls_no_check(true)
+            .tls_client_cert(client_cert.pem())
+            .tls_client_key(client_key.serialize_pem())
+            .verify_server_attestation(false)
+            .cert_validator(Box::new(move |info| {
+                let info = info.ok_or_else(|| anyhow::anyhow!("missing server certificate"))?;
+                anyhow::ensure!(info.cert_der == pinned_cert, "server certificate changed");
+                Ok(())
+            }))
+            .build()
+            .into_client()
+            .unwrap();
+        let result = client.request::<(), ()>("Probe", ()).await;
+        assert_eq!(result.is_ok(), correct_key, "{version:?}: {result:?}");
+        assert_eq!(server.await.unwrap(), correct_key);
+    }
+}