Skip to content

Commit 9b590a3

Browse files
kvinwangkvinwang
committed
rust-sdk: Refactor TappdClient API
1 parent afb66e0 commit 9b590a3

5 files changed

Lines changed: 94 additions & 177 deletions

File tree

‎sdk/run-tests.sh‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,8 @@ popd
1414

1515
pushd rust/
1616
cargo test -- --show-output
17+
cargo run --example tappd_client_usage
18+
cargo run --example dstack_client_usage
1719
popd
1820

1921
pushd go/

‎sdk/rust/examples/dstack_client_usage.rs‎

Lines changed: 62 additions & 128 deletions
Original file line numberDiff line numberDiff line change
@@ -16,99 +16,55 @@ async fn main() -> anyhow::Result<()> {
1616
// Example usage (these will fail without a running dstack service):
1717

1818
// 1. Get system info
19-
match client.info().await {
20-
Ok(info) => {
21-
println!("System info retrieved successfully!");
22-
println!(" App ID: {}", info.app_id);
23-
println!(" Instance ID: {}", info.instance_id);
24-
println!(" App Name: {}", info.app_name);
25-
println!(" Device ID: {}", info.device_id);
26-
println!(" Compose Hash: {}", info.compose_hash);
27-
println!(" TCB Info - MRTD: {}", info.tcb_info.mrtd);
28-
println!(" TCB Info - RTMR0: {}", info.tcb_info.rtmr0);
29-
}
30-
Err(e) => {
31-
println!("Failed to get system info: {}", e);
32-
}
33-
}
19+
let info = client.info().await?;
20+
println!("System info retrieved successfully!");
21+
println!(" App ID: {}", info.app_id);
22+
println!(" Instance ID: {}", info.instance_id);
23+
println!(" App Name: {}", info.app_name);
24+
println!(" Device ID: {}", info.device_id);
25+
println!(" Compose Hash: {}", info.compose_hash);
26+
println!(" TCB Info - MRTD: {}", info.tcb_info.mrtd);
27+
println!(" TCB Info - RTMR0: {}", info.tcb_info.rtmr0);
3428

3529
// 2. Derive a key
36-
match client
30+
let response = client
3731
.get_key(Some("my-app".to_string()), Some("encryption".to_string()))
38-
.await
39-
{
40-
Ok(response) => {
41-
println!("Key derived successfully!");
42-
println!(" Key length: {}", response.key.len());
43-
println!(
44-
" Signature chain length: {}",
45-
response.signature_chain.len()
46-
);
47-
48-
// Decode the key
49-
match response.decode_key() {
50-
Ok(key_bytes) => {
51-
println!(" Decoded key bytes length: {}", key_bytes.len());
52-
}
53-
Err(e) => {
54-
println!(" Key decode error: {}", e);
55-
}
56-
}
57-
}
58-
Err(e) => {
59-
println!("Failed to derive key: {}", e);
60-
}
61-
}
32+
.await?;
33+
println!("Key derived successfully!");
34+
println!(" Key length: {}", response.key.len());
35+
println!(
36+
" Signature chain length: {}",
37+
response.signature_chain.len()
38+
);
39+
40+
// Decode the key
41+
let key_bytes = response.decode_key()?;
42+
println!(" Decoded key bytes length: {}", key_bytes.len());
6243

6344
// 3. Generate TDX quote
6445
let report_data = b"Hello, dstack world!".to_vec();
65-
match client.get_quote(report_data).await {
66-
Ok(response) => {
67-
println!("TDX quote generated successfully!");
68-
println!(" Quote length: {}", response.quote.len());
69-
println!(" Event log length: {}", response.event_log.len());
70-
71-
// Decode the quote
72-
match response.decode_quote() {
73-
Ok(quote_bytes) => {
74-
println!(" Decoded quote bytes length: {}", quote_bytes.len());
75-
}
76-
Err(e) => {
77-
println!(" Quote decode error: {}", e);
78-
}
79-
}
80-
81-
// Replay RTMRs from event log
82-
match response.replay_rtmrs() {
83-
Ok(rtmrs) => {
84-
println!(" Replayed RTMRs: {} entries", rtmrs.len());
85-
for (idx, rtmr) in rtmrs.iter() {
86-
println!(" RTMR{}: {}", idx, rtmr);
87-
}
88-
}
89-
Err(e) => {
90-
println!(" RTMR replay error: {}", e);
91-
}
92-
}
93-
}
94-
Err(e) => {
95-
println!("Failed to get TDX quote: {}", e);
96-
}
46+
let response = client.get_quote(report_data).await?;
47+
println!("TDX quote generated successfully!");
48+
println!(" Quote length: {}", response.quote.len());
49+
println!(" Event log length: {}", response.event_log.len());
50+
51+
// Decode the quote
52+
let quote_bytes = response.decode_quote()?;
53+
println!(" Decoded quote bytes length: {}", quote_bytes.len());
54+
55+
// Replay RTMRs from event log
56+
let rtmrs = response.replay_rtmrs()?;
57+
println!(" Replayed RTMRs: {} entries", rtmrs.len());
58+
for (idx, rtmr) in rtmrs.iter() {
59+
println!(" RTMR{}: {}", idx, rtmr);
9760
}
9861

9962
// 4. Emit an event
10063
let event_payload = b"Application started successfully".to_vec();
101-
match client
64+
client
10265
.emit_event("AppStart".to_string(), event_payload)
103-
.await
104-
{
105-
Ok(()) => {
106-
println!("Event emitted successfully!");
107-
}
108-
Err(e) => {
109-
println!("Failed to emit event: {}", e);
110-
}
111-
}
66+
.await?;
67+
println!("Event emitted successfully!");
11268

11369
// 5. Get TLS key for server authentication
11470
let tls_config = TlsKeyConfig::builder()
@@ -122,57 +78,35 @@ async fn main() -> anyhow::Result<()> {
12278
.usage_ra_tls(true)
12379
.build();
12480

125-
match client.get_tls_key(tls_config).await {
126-
Ok(response) => {
127-
println!("TLS key generated successfully!");
128-
println!(" Key length: {}", response.key.len());
129-
println!(
130-
" Certificate chain length: {}",
131-
response.certificate_chain.len()
132-
);
133-
}
134-
Err(e) => {
135-
println!("Failed to get TLS key: {}", e);
136-
}
137-
}
81+
let response = client.get_tls_key(tls_config).await?;
82+
println!("TLS key generated successfully!");
83+
println!(" Key length: {}", response.key.len());
84+
println!(
85+
" Certificate chain length: {}",
86+
response.certificate_chain.len()
87+
);
13888

13989
// 6. Get a simple key without purpose
140-
match client.get_key(Some("simple-key".to_string()), None).await {
141-
Ok(response) => {
142-
println!("Simple key derived successfully!");
143-
println!(" Key: {}", response.key);
144-
}
145-
Err(e) => {
146-
println!("Failed to derive simple key: {}", e);
147-
}
148-
}
90+
let response = client.get_key(Some("simple-key".to_string()), None).await?;
91+
println!("Simple key derived successfully!");
92+
println!(" Key: {}", response.key);
14993

15094
// 7. Generate quote with minimal report data
15195
let minimal_data = vec![0x01, 0x02, 0x03, 0x04];
152-
match client.get_quote(minimal_data).await {
153-
Ok(response) => {
154-
println!("Minimal quote generated successfully!");
155-
156-
// Parse and display event log
157-
match response.decode_event_log() {
158-
Ok(events) => {
159-
println!(" Event log contains {} events", events.len());
160-
for (i, event) in events.iter().enumerate().take(3) {
161-
// Show first 3 events
162-
println!(
163-
" Event {}: IMR={}, Type={}, Event='{}'",
164-
i, event.imr, event.event_type, event.event
165-
);
166-
}
167-
}
168-
Err(e) => {
169-
println!(" Failed to parse event log: {}", e);
170-
}
171-
}
172-
}
173-
Err(e) => {
174-
println!("Failed to get minimal quote: {}", e);
175-
}
96+
let response = client.get_quote(minimal_data).await?;
97+
println!("Minimal quote generated successfully!");
98+
println!(" Quote length: {}", response.quote.len());
99+
println!(" Event log length: {}", response.event_log.len());
100+
101+
// Parse and display event log
102+
let events = response.decode_event_log()?;
103+
println!(" Event log contains {} events", events.len());
104+
for (i, event) in events.iter().enumerate().take(3) {
105+
// Show first 3 events
106+
println!(
107+
" Event {}: IMR={}, Type={}, Event='{}'",
108+
i, event.imr, event.event_type, event.event
109+
);
176110
}
177111

178112
Ok(())

‎sdk/rust/examples/tappd_client_usage.rs‎

Lines changed: 24 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -16,45 +16,37 @@ async fn main() -> anyhow::Result<()> {
1616
// Example usage (these will fail without a running tappd service):
1717

1818
// 1. Derive a key
19-
match client.derive_key("my/key/path").await {
20-
Ok(response) => {
21-
println!("Key derived successfully!");
22-
println!(
23-
"Certificate chain length: {}",
24-
response.certificate_chain.len()
25-
);
26-
}
27-
Err(e) => {
28-
println!("Failed to derive key: {}", e);
29-
}
30-
}
19+
let response = client.derive_key("my/key/path").await?;
20+
println!("Key derived successfully!");
21+
println!(
22+
"Certificate chain length: {}",
23+
response.certificate_chain.len()
24+
);
25+
let ecdsa_p256_key = response.decode_key().unwrap();
26+
println!("ECDSA P-256 key length: {}", ecdsa_p256_key.len());
3127

3228
// 2. Get a quote with 64 bytes of report data
3329
let mut report_data = b"Hello, world!".to_vec();
3430
// Pad to exactly 64 bytes for get_quote
3531
report_data.resize(64, 0);
36-
match client.get_quote(report_data).await {
37-
Ok(response) => {
38-
println!("Quote generated successfully!");
39-
println!("Quote length: {}", response.quote.len());
40-
}
41-
Err(e) => {
42-
println!("Failed to get quote: {}", e);
43-
}
44-
}
32+
let response = client.get_quote(report_data).await?;
33+
println!("Quote generated successfully!");
34+
println!("Quote length: {}", response.quote.len());
35+
36+
// 2. Get a quote with 64 bytes of report data
37+
let mut report_data = b"Hello, world!".to_vec();
38+
// Pad to exactly 64 bytes for get_quote
39+
report_data.resize(64, 0);
40+
let response = client.get_quote(report_data).await?;
41+
println!("Quote generated successfully!");
42+
println!("Quote length: {}", response.quote.len());
4543

4644
// 3. Get instance info
47-
match client.info().await {
48-
Ok(info) => {
49-
println!("Instance info retrieved successfully!");
50-
println!("App ID: {}", info.app_id);
51-
println!("Instance ID: {}", info.instance_id);
52-
println!("App Name: {}", info.app_name);
53-
}
54-
Err(e) => {
55-
println!("Failed to get instance info: {}", e);
56-
}
57-
}
45+
let response = client.info().await?;
46+
println!("Instance info retrieved successfully!");
47+
println!("App ID: {}", response.app_id);
48+
println!("Instance ID: {}", response.instance_id);
49+
println!("App Name: {}", response.app_name);
5850

5951
Ok(())
6052
}

‎sdk/rust/src/tappd_client.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -87,7 +87,7 @@ pub struct DeriveKeyResponse {
8787

8888
impl DeriveKeyResponse {
8989
/// Decodes the key from PEM format and extracts the raw ECDSA P-256 private key bytes
90-
pub fn to_bytes(&self) -> Result<Vec<u8>, anyhow::Error> {
90+
pub fn decode_key(&self) -> Result<Vec<u8>, anyhow::Error> {
9191
use x509_parser::der_parser::der::parse_der;
9292
use x509_parser::pem::parse_x509_pem;
9393

‎sdk/rust/tests/test_tappd_client.rs‎

Lines changed: 5 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -70,15 +70,9 @@ async fn test_tappd_client_derive_key_integration() {
7070
assert!(!response.key.is_empty());
7171

7272
// Test key decoding
73-
match response.to_bytes() {
74-
Ok(key_bytes) => {
75-
println!(" Decoded key bytes length: {}", key_bytes.len());
76-
assert!(!key_bytes.is_empty());
77-
}
78-
Err(e) => {
79-
println!(" Key decode error: {}", e);
80-
}
81-
}
73+
let key_bytes = response.decode_key().unwrap();
74+
println!("✓ Decoded key bytes length: {}", key_bytes.len());
75+
assert_eq!(key_bytes.len(), 32);
8276
}
8377

8478
#[tokio::test]
@@ -167,7 +161,6 @@ async fn test_tappd_client_get_quote_integration() {
167161
}
168162
}
169163

170-
171164
// Helper function to get a test client
172165
fn get_test_client() -> TappdClient {
173166
// Check for simulator endpoint first
@@ -187,7 +180,6 @@ fn get_test_client() -> TappdClient {
187180
TappdClient::new(None)
188181
}
189182

190-
191183
#[test]
192184
fn test_derive_key_response_decode() {
193185
use dstack_sdk::tappd_client::DeriveKeyResponse;
@@ -200,12 +192,9 @@ fn test_derive_key_response_decode() {
200192
};
201193

202194
// The implementation should return the decoded ECDSA P-256 private key bytes
203-
let bytes = response.to_bytes().unwrap();
195+
let bytes = response.decode_key().unwrap();
204196
assert!(!bytes.is_empty());
205197
// For a valid ECDSA P-256 key, we should get either 32 bytes (the private key)
206198
// or fall back to the full DER contents if parsing fails
207-
assert!(bytes.len() == 32 || bytes.len() > 32);
199+
assert_eq!(bytes.len(), 32);
208200
}
209-
210-
211-

0 commit comments

Comments
 (0)