agent: Always generate random tls key

Author: kvinwang

gateway/src/main.rs

@@ -1,7 +1,7 @@
 use anyhow::{anyhow, bail, Context, Result};
 use clap::Parser;
 use config::{Config, TlsConfig};
-use dstack_guest_agent_rpc::dstack_guest_client::DstackGuestClient;
+use dstack_guest_agent_rpc::{dstack_guest_client::DstackGuestClient, GetTlsKeyArgs};
 use http_client::prpc::PrpcClient;
 use ra_rpc::{client::RaClient, rocket_helper::QuoteVerifier};
 use rocket::{
@@ -67,14 +67,12 @@ async fn maybe_gen_certs(config: &Config, tls_config: &TlsConfig) -> Result<()>
         info!("Using dstack guest agent for certificate generation");
         let agent_client = dstack_agent().context("Failed to create dstack client")?;
         let response = agent_client
-            .get_tls_key(dstack_guest_agent_rpc::GetTlsKeyArgs {
-                path: "".to_string(),
+            .get_tls_key(GetTlsKeyArgs {
                 subject: "dstack-gateway".to_string(),
                 alt_names: vec![config.rpc_domain.clone()],
                 usage_ra_tls: true,
                 usage_server_auth: true,
                 usage_client_auth: false,
-                random_seed: true,
             })
             .await?;
 

gateway/src/main_service/sync_client.rs

@@ -92,13 +92,11 @@ pub(crate) async fn sync_task(
         let agent = dstack_agent().context("Failed to create dstack agent client")?;
         let keys = agent
             .get_tls_key(GetTlsKeyArgs {
-                path: "/sync-state-client".into(),
-                subject: "".into(),
+                subject: "dstack-gateway-sync-client".into(),
                 alt_names: vec![],
                 usage_ra_tls: false,
                 usage_server_auth: false,
                 usage_client_auth: true,
-                random_seed: true,
             })
             .await
             .context("Failed to get sync-client keys")?;

guest-agent/rpc/proto/agent_rpc.proto

@@ -8,7 +8,7 @@ package dstack_guest;
 service Tappd {
   // Derives a cryptographic key from the specified key path.
   // Returns the derived key along with its certificate chain.
-  rpc DeriveKey(GetTlsKeyArgs) returns (GetTlsKeyResponse) {}
+  rpc DeriveKey(DeriveKeyArgs) returns (GetTlsKeyResponse) {}
 
   // Derives a new ECDSA key with k256 EC curve.
   rpc DeriveK256Key(GetKeyArgs) returns (DeriveK256KeyResponse) {}
@@ -43,9 +43,22 @@ service DstackGuest {
   rpc Info(google.protobuf.Empty) returns (WorkerInfo) {}
 }
 
-
 // The request to derive a key
 message GetTlsKeyArgs {
+  // Subject of the certificate to request
+  string subject = 1;
+  // DNS alternative names for the certificate
+  repeated string alt_names = 2;
+  // Includes quote in the certificate
+  bool usage_ra_tls = 3;
+  // Key usage server auth
+  bool usage_server_auth = 4;
+  // Key usage client auth
+  bool usage_client_auth = 5;
+}
+
+// The request to derive a key
+message DeriveKeyArgs {
   // Path used to derive the private key
   string path = 1;
   // Bellow fields are used to generate the certificate

guest-agent/src/rpc_service.rs

@@ -6,8 +6,9 @@ use dstack_guest_agent_rpc::{
     dstack_guest_server::{DstackGuestRpc, DstackGuestServer},
     tappd_server::{TappdRpc, TappdServer},
     worker_server::{WorkerRpc, WorkerServer},
-    DeriveK256KeyResponse, GetKeyArgs, GetKeyResponse, GetQuoteResponse, GetTlsKeyArgs,
-    GetTlsKeyResponse, RawQuoteArgs, TdxQuoteArgs, TdxQuoteResponse, WorkerInfo, WorkerVersion,
+    DeriveK256KeyResponse, DeriveKeyArgs, GetKeyArgs, GetKeyResponse, GetQuoteResponse,
+    GetTlsKeyArgs, GetTlsKeyResponse, RawQuoteArgs, TdxQuoteArgs, TdxQuoteResponse, WorkerInfo,
+    WorkerVersion,
 };
 use dstack_types::AppKeys;
 use fs_err as fs;
@@ -82,17 +83,12 @@ pub struct InternalRpcHandler {
 
 impl DstackGuestRpc for InternalRpcHandler {
     async fn get_tls_key(self, request: GetTlsKeyArgs) -> anyhow::Result<GetTlsKeyResponse> {
-        let mut mbuf = [0u8; 32];
-        let seed = if request.random_seed {
-            SystemRandom::new()
-                .fill(&mut mbuf)
-                .context("Failed to generate secure seed")?;
-            &mbuf[..]
-        } else {
-            &self.state.inner.keys.k256_key
-        };
-        let derived_key = derive_ecdsa_key_pair_from_bytes(seed, &[request.path.as_bytes()])
-            .context("Failed to derive key")?;
+        let mut seed = [0u8; 32];
+        SystemRandom::new()
+            .fill(&mut seed)
+            .context("Failed to generate secure seed")?;
+        let derived_key =
+            derive_ecdsa_key_pair_from_bytes(&seed, &[]).context("Failed to derive key")?;
         let config = CertConfig {
             org_name: None,
             subject: request.subject,
@@ -179,10 +175,37 @@ pub struct InternalRpcHandlerV0 {
 }
 
 impl TappdRpc for InternalRpcHandlerV0 {
-    async fn derive_key(self, request: GetTlsKeyArgs) -> anyhow::Result<GetTlsKeyResponse> {
-        InternalRpcHandler { state: self.state }
-            .get_tls_key(request)
+    async fn derive_key(self, request: DeriveKeyArgs) -> anyhow::Result<GetTlsKeyResponse> {
+        let mut mbuf = [0u8; 32];
+        let seed = if request.random_seed {
+            SystemRandom::new()
+                .fill(&mut mbuf)
+                .context("Failed to generate secure seed")?;
+            &mbuf[..]
+        } else {
+            &self.state.inner.keys.k256_key
+        };
+        let derived_key = derive_ecdsa_key_pair_from_bytes(seed, &[request.path.as_bytes()])
+            .context("Failed to derive key")?;
+        let config = CertConfig {
+            org_name: None,
+            subject: request.subject,
+            subject_alt_names: request.alt_names,
+            usage_server_auth: request.usage_server_auth,
+            usage_client_auth: request.usage_client_auth,
+            ext_quote: request.usage_ra_tls,
+        };
+        let certificate_chain = self
+            .state
+            .inner
+            .cert_client
+            .request_cert(&derived_key, config)
             .await
+            .context("Failed to sign the CSR")?;
+        Ok(GetTlsKeyResponse {
+            key: derived_key.serialize_pem(),
+            certificate_chain,
+        })
     }
 
     async fn derive_k256_key(self, request: GetKeyArgs) -> Result<DeriveK256KeyResponse> {