refactor(cert): move csr assembly out of cert-client

Author: kvinwang

cert-client/src/lib.rs

@@ -6,11 +6,7 @@ use anyhow::{Context, Result};
 use dstack_kms_rpc::{kms_client::KmsClient, SignCertRequest};
 use dstack_types::{AppKeys, KeyProvider};
 use ra_rpc::client::{RaClient, RaClientConfig};
-use ra_tls::{
-    attestation::{QuoteContentType, VersionedAttestation},
-    cert::{generate_ra_cert, CaCert, CertConfigV2, CertSigningRequestV2, Csr},
-    rcgen::KeyPair,
-};
+use ra_tls::cert::{generate_ra_cert, CaCert, CertSigningRequestV2};
 
 pub enum CertRequestClient {
     Local {
@@ -92,35 +88,4 @@ impl CertRequestClient {
             }
         }
     }
-
-    pub async fn request_cert(&self, key: &KeyPair, config: CertConfigV2) -> Result<Vec<String>> {
-        let pubkey = key.public_key_der();
-        let report_data = QuoteContentType::RaTlsCert.to_report_data(&pubkey);
-        let attestation = ra_rpc::Attestation::quote(&report_data)
-            .context("Failed to get quote for cert pubkey")?
-            .into_versioned();
-        self.sign_cert_with_attestation(key, config, attestation).await
-    }
-
-    pub async fn sign_cert_with_attestation(
-        &self,
-        key: &KeyPair,
-        config: CertConfigV2,
-        mut attestation: VersionedAttestation,
-    ) -> Result<Vec<String>> {
-        let pubkey = key.public_key_der();
-        let report_data = QuoteContentType::RaTlsCert.to_report_data(&pubkey);
-        attestation.set_report_data(report_data);
-
-        let csr = CertSigningRequestV2 {
-            confirm: "please sign cert:".to_string(),
-            pubkey,
-            config,
-            attestation,
-        };
-        let signature = csr.signed_by(key).context("Failed to sign the CSR")?;
-        self.sign_csr(&csr, &signature)
-            .await
-            .context("Failed to sign the CSR")
-    }
 }

dstack-util/src/system_setup.rs

@@ -27,8 +27,11 @@ use luks2::{
     LuksAf, LuksConfig, LuksDigest, LuksHeader, LuksJson, LuksKdf, LuksKeyslot, LuksSegment,
     LuksSegmentSize,
 };
-use ra_rpc::client::{CertInfo, RaClient, RaClientConfig};
-use ra_tls::cert::{generate_ra_cert, CertConfigV2};
+use ra_rpc::{client::{CertInfo, RaClient, RaClientConfig}, Attestation};
+use ra_tls::{
+    attestation::QuoteContentType,
+    cert::{generate_ra_cert, CertConfigV2, CertSigningRequestV2, Csr},
+};
 use rand::Rng as _;
 use safe_write::safe_write;
 use scopeguard::defer;
@@ -53,6 +56,30 @@ use ra_tls::rcgen::{KeyPair, PKCS_ECDSA_P256_SHA256};
 use serde_human_bytes as hex_bytes;
 use serde_json::Value;
 
+
+async fn sign_cert_request(
+    cert_client: &CertRequestClient,
+    key: &KeyPair,
+    config: CertConfigV2,
+) -> Result<Vec<String>> {
+    let pubkey = key.public_key_der();
+    let report_data = QuoteContentType::RaTlsCert.to_report_data(&pubkey);
+    let attestation = Attestation::quote(&report_data)
+        .context("Failed to get quote for cert pubkey")?
+        .into_versioned();
+    let csr = CertSigningRequestV2 {
+        confirm: "please sign cert:".to_string(),
+        pubkey,
+        config,
+        attestation,
+    };
+    let signature = csr.signed_by(key).context("Failed to sign the CSR")?;
+    cert_client
+        .sign_csr(&csr, &signature)
+        .await
+        .context("Failed to sign the CSR")
+}
+
 mod config_id_verifier;
 
 #[derive(clap::Parser)]
@@ -500,8 +527,7 @@ impl<'a> GatewayContext<'a> {
             not_before: None,
             not_after: Some(cert_not_after),
         };
-        let certs = cert_client
-            .request_cert(&key, config, None)
+        let certs = sign_cert_request(&cert_client, &key, config)
             .await
             .context("Failed to request cert")?;
         let client_cert = certs.join("\n");
@@ -520,8 +546,7 @@ impl<'a> GatewayContext<'a> {
             not_before: None,
             not_after: Some(cert_not_after),
         };
-        let certs_with_quote = cert_client
-            .request_cert(&key, config_with_quote, None)
+        let certs_with_quote = sign_cert_request(&cert_client, &key, config_with_quote)
             .await
             .context("Failed to request cert with quote")?;
         let client_cert_with_quote = certs_with_quote.join("\n");

guest-agent/src/rpc_service.rs

@@ -27,7 +27,7 @@ use or_panic::ResultOrPanic;
 use ra_rpc::{CallContext, RpcCall};
 use ra_tls::{
     attestation::{QuoteContentType, DEFAULT_HASH_ALGORITHM},
-    cert::CertConfigV2,
+    cert::{CertConfigV2, CertSigningRequestV2, Csr},
     kdf::{derive_key, derive_p256_key_pair_from_bytes},
 };
 use rcgen::KeyPair;
@@ -67,12 +67,20 @@ impl AppStateInner {
     }
 
     async fn issue_cert(&self, key: &KeyPair, config: CertConfigV2) -> Result<Vec<String>> {
+        let pubkey = key.public_key_der();
         let attestation = self
             .platform
-            .certificate_attestation(&key.public_key_der())
+            .certificate_attestation(&pubkey)
             .context("Failed to get certificate attestation")?;
+        let csr = CertSigningRequestV2 {
+            confirm: "please sign cert:".to_string(),
+            pubkey,
+            config,
+            attestation,
+        };
+        let signature = csr.signed_by(key).context("Failed to sign the CSR")?;
         self.cert_client
-            .sign_cert_with_attestation(key, config, attestation)
+            .sign_csr(&csr, &signature)
             .await
             .context("Failed to sign the CSR")
     }