From 25bc06fba53c33a4022ea2fa24126396928c5070 Mon Sep 17 00:00:00 2001
From: tuddman <tuddman@users.noreply.github.com>
Date: Wed, 14 May 2025 16:51:37 +0200
Subject: [PATCH] fix(sdk, rust): adds build-able TlsKeyConfig

---
 Cargo.lock                    |  1 +
 sdk/rust/Cargo.toml           |  1 +
 sdk/rust/src/dstack_client.rs | 31 +++++++++++++++----------------
 sdk/rust/tests/test_client.rs | 18 ++++++------------
 sdk/rust/tests/test_eth.rs    |  5 ++---
 5 files changed, 25 insertions(+), 31 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index beda34c1..c67c171f 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -2325,6 +2325,7 @@ version = "0.1.0"
 dependencies = [
  "alloy",
  "anyhow",
+ "bon",
  "dcap-qvl",
  "hex",
  "http",
diff --git a/sdk/rust/Cargo.toml b/sdk/rust/Cargo.toml
index b9499b03..889c2f36 100644
--- a/sdk/rust/Cargo.toml
+++ b/sdk/rust/Cargo.toml
@@ -9,6 +9,7 @@ authors = ["Encifher <encifher@rizelabs.io>"]
 [dependencies]
 alloy = { version = "0.15", features = ["signers", "signer-local"] }
 anyhow.workspace = true
+bon.workspace = true
 hex.workspace = true 
 http.workspace = true 
 http-client-unix-domain-socket = "0.1.1"
diff --git a/sdk/rust/src/dstack_client.rs b/sdk/rust/src/dstack_client.rs
index 0a93d052..48af3bc9 100644
--- a/sdk/rust/src/dstack_client.rs
+++ b/sdk/rust/src/dstack_client.rs
@@ -51,6 +51,19 @@ pub struct EventLog {
     event_payload: String,
 }
 
+#[derive(bon::Builder, Serialize)]
+pub struct TlsKeyConfig {
+    #[builder(into)]
+    pub subject: Option<String>,
+    pub alt_names: Option<Vec<String>>,
+    #[builder(default = false)]
+    pub usage_ra_tls: bool,
+    #[builder(default = true)]
+    pub usage_server_auth: bool,
+    #[builder(default = false)]
+    pub usage_client_auth: bool,
+}
+
 #[derive(Serialize, Deserialize)]
 pub struct GetKeyResponse {
     pub key: String,
@@ -248,22 +261,8 @@ impl DstackClient {
         Ok(())
     }
 
-    pub async fn get_tls_key(
-        &self,
-        subject: Option<String>,
-        alt_names: Option<Vec<String>>,
-        usage_ra_tls: bool,
-        usage_server_auth: bool,
-        usage_client_auth: bool,
-    ) -> Result<GetTlsKeyResponse> {
-        let data = json!({
-            "subject": subject.unwrap_or_default(),
-            "alt_names": alt_names.unwrap_or_default(),
-            "usage_ra_tls": usage_ra_tls,
-            "usage_server_auth": usage_server_auth,
-            "usage_client_auth": usage_client_auth
-        });
-        let response = self.send_rpc_request("/GetTlsKey", &data).await?;
+    pub async fn get_tls_key(&self, tls_key_config: TlsKeyConfig) -> Result<GetTlsKeyResponse> {
+        let response = self.send_rpc_request("/GetTlsKey", &tls_key_config).await?;
         let response = serde_json::from_value::<GetTlsKeyResponse>(response)?;
 
         Ok(response)
diff --git a/sdk/rust/tests/test_client.rs b/sdk/rust/tests/test_client.rs
index ce557927..02d2a061 100644
--- a/sdk/rust/tests/test_client.rs
+++ b/sdk/rust/tests/test_client.rs
@@ -19,10 +19,8 @@ async fn test_async_client_get_quote() {
 #[tokio::test]
 async fn test_async_client_get_tls_key() {
     let client = AsyncDstackClient::new(None);
-    let result = client
-        .get_tls_key(None, None, false, false, false)
-        .await
-        .unwrap();
+    let key_config = dstack_sdk::dstack_client::TlsKeyConfig::builder().build();
+    let result = client.get_tls_key(key_config).await.unwrap();
     assert!(result.key.starts_with("-----BEGIN PRIVATE KEY-----"));
     assert!(!result.certificate_chain.is_empty());
 }
@@ -30,14 +28,10 @@ async fn test_async_client_get_tls_key() {
 #[tokio::test]
 async fn test_tls_key_uniqueness() {
     let client = AsyncDstackClient::new(None);
-    let result1 = client
-        .get_tls_key(None, None, false, false, false)
-        .await
-        .unwrap();
-    let result2 = client
-        .get_tls_key(None, None, false, false, false)
-        .await
-        .unwrap();
+    let key_config_1 = dstack_sdk::dstack_client::TlsKeyConfig::builder().build();
+    let key_config_2 = dstack_sdk::dstack_client::TlsKeyConfig::builder().build();
+    let result1 = client.get_tls_key(key_config_1).await.unwrap();
+    let result2 = client.get_tls_key(key_config_2).await.unwrap();
     assert_ne!(result1.key, result2.key);
 }
 
diff --git a/sdk/rust/tests/test_eth.rs b/sdk/rust/tests/test_eth.rs
index b73b5d48..608e633a 100644
--- a/sdk/rust/tests/test_eth.rs
+++ b/sdk/rust/tests/test_eth.rs
@@ -1,6 +1,5 @@
-use dstack_rust::dstack_client::{DstackClient, GetKeyResponse};
-use dstack_rust::ethereum::to_account;
-use tokio;
+use dstack_sdk::dstack_client::{DstackClient, GetKeyResponse};
+use dstack_sdk::ethereum::to_account;
 
 #[tokio::test]
 async fn test_async_to_keypair() {