dstack_sdk/
dstack_client.rs

1// SPDX-FileCopyrightText: © 2025 Created-for-a-purpose <rachitchahar@gmail.com>
2// SPDX-FileCopyrightText: © 2025 Daniel Sharifi <daniel.sharifi@nearone.org>
3// SPDX-FileCopyrightText: © 2025 tuddman <tuddman@users.noreply.github.com>
4//
5// SPDX-License-Identifier: Apache-2.0
6
7use anyhow::Result;
8use hex::encode as hex_encode;
9use http_client_unix_domain_socket::{ClientUnix, Method};
10use reqwest::Client;
11use serde::{de::DeserializeOwned, Serialize};
12use serde_json::{json, Value};
13use std::env;
14
15pub use dstack_sdk_types::dstack::*;
16
17// Internal request structs for hex encoding
18#[derive(Debug, Serialize)]
19struct SignRequest<'a> {
20    algorithm: &'a str,
21    data: String,
22}
23
24#[derive(Debug, Serialize)]
25struct VerifyRequest<'a> {
26    algorithm: &'a str,
27    data: String,
28    signature: String,
29    public_key: String,
30}
31
32fn get_endpoint(endpoint: Option<&str>) -> String {
33    if let Some(e) = endpoint {
34        return e.to_string();
35    }
36    if let Ok(sim_endpoint) = env::var("DSTACK_SIMULATOR_ENDPOINT") {
37        return sim_endpoint;
38    }
39    // Try paths in order: legacy paths first, then namespaced paths
40    const SOCKET_PATHS: &[&str] = &[
41        "/var/run/dstack.sock",
42        "/run/dstack.sock",
43        "/var/run/dstack/dstack.sock",
44        "/run/dstack/dstack.sock",
45    ];
46    for path in SOCKET_PATHS {
47        if std::path::Path::new(path).exists() {
48            return path.to_string();
49        }
50    }
51    // Default to new path even if not exists (will fail with clear error)
52    SOCKET_PATHS[0].to_string()
53}
54
55#[derive(Debug)]
56pub enum ClientKind {
57    Http,
58    Unix,
59}
60
61pub trait BaseClient {}
62
63/// The main client for interacting with the dstack service
64pub struct DstackClient {
65    /// The base URL for HTTP requests
66    base_url: String,
67    /// The endpoint for Unix domain socket communication
68    endpoint: String,
69    /// The type of client (HTTP or Unix domain socket)
70    client: ClientKind,
71}
72
73impl BaseClient for DstackClient {}
74
75impl DstackClient {
76    pub fn new(endpoint: Option<&str>) -> Self {
77        let endpoint = get_endpoint(endpoint);
78        let (base_url, client) = match endpoint {
79            ref e if e.starts_with("http://") || e.starts_with("https://") => {
80                (e.to_string(), ClientKind::Http)
81            }
82            _ => ("http://localhost".to_string(), ClientKind::Unix),
83        };
84
85        DstackClient {
86            base_url,
87            endpoint,
88            client,
89        }
90    }
91
92    async fn send_rpc_request<S: Serialize, D: DeserializeOwned>(
93        &self,
94        path: &str,
95        payload: &S,
96    ) -> anyhow::Result<D> {
97        match &self.client {
98            ClientKind::Http => {
99                let client = Client::new();
100                let url = format!(
101                    "{}/{}",
102                    self.base_url.trim_end_matches('/'),
103                    path.trim_start_matches('/')
104                );
105                let res = client
106                    .post(&url)
107                    .json(payload)
108                    .header("Content-Type", "application/json")
109                    .send()
110                    .await?
111                    .error_for_status()?;
112                Ok(res.json().await?)
113            }
114            ClientKind::Unix => {
115                let mut unix_client = ClientUnix::try_new(&self.endpoint).await?;
116                let res = unix_client
117                    .send_request_json::<_, _, Value>(
118                        path,
119                        Method::POST,
120                        &[("Content-Type", "application/json"), ("Host", "dstack")],
121                        Some(&payload),
122                    )
123                    .await?;
124                Ok(res.1)
125            }
126        }
127    }
128
129    pub async fn get_key(
130        &self,
131        path: Option<String>,
132        purpose: Option<String>,
133    ) -> Result<GetKeyResponse> {
134        let data = json!({
135            "path": path.unwrap_or_default(),
136            "purpose": purpose.unwrap_or_default(),
137            "algorithm": "secp256k1", // Default or specify as needed
138        });
139        let response = self.send_rpc_request("/GetKey", &data).await?;
140        let response = serde_json::from_value::<GetKeyResponse>(response)?;
141
142        Ok(response)
143    }
144
145    pub async fn get_quote(&self, report_data: Vec<u8>) -> Result<GetQuoteResponse> {
146        if report_data.is_empty() || report_data.len() > 64 {
147            anyhow::bail!("Invalid report data length")
148        }
149        let hex_data = hex_encode(report_data);
150        let data = json!({ "report_data": hex_data });
151        let response = self.send_rpc_request("/GetQuote", &data).await?;
152        let response = serde_json::from_value::<GetQuoteResponse>(response)?;
153
154        Ok(response)
155    }
156
157    pub async fn attest(&self, report_data: Vec<u8>) -> Result<AttestResponse> {
158        if report_data.is_empty() || report_data.len() > 64 {
159            anyhow::bail!("Invalid report data length")
160        }
161        let hex_data = hex_encode(report_data);
162        let data = json!({ "report_data": hex_data });
163        let response = self.send_rpc_request("/Attest", &data).await?;
164        let response = serde_json::from_value::<AttestResponse>(response)?;
165
166        Ok(response)
167    }
168
169    pub async fn info(&self) -> Result<InfoResponse> {
170        let response = self.send_rpc_request("/Info", &json!({})).await?;
171        Ok(InfoResponse::validated_from_value(response)?)
172    }
173
174    /// Query the guest-agent version.
175    ///
176    /// Returns `Ok(VersionResponse)` on OS >= 0.5.7.
177    /// Returns an error on older OS versions that lack the Version RPC.
178    pub async fn version(&self) -> Result<VersionResponse> {
179        let response = self.send_rpc_request("/Version", &json!({})).await?;
180        let response = serde_json::from_value::<VersionResponse>(response)?;
181        Ok(response)
182    }
183
184    pub async fn emit_event(&self, event: String, payload: Vec<u8>) -> Result<()> {
185        if event.is_empty() {
186            anyhow::bail!("Event name cannot be empty")
187        }
188        let hex_payload = hex_encode(payload);
189        let data = json!({ "event": event, "payload": hex_payload });
190        self.send_rpc_request::<_, ()>("/EmitEvent", &data).await?;
191        Ok(())
192    }
193
194    pub async fn get_tls_key(&self, tls_key_config: TlsKeyConfig) -> Result<GetTlsKeyResponse> {
195        let response = self.send_rpc_request("/GetTlsKey", &tls_key_config).await?;
196        let response = serde_json::from_value::<GetTlsKeyResponse>(response)?;
197
198        Ok(response)
199    }
200
201    /// Signs a payload using a derived key.
202    pub async fn sign(&self, algorithm: &str, data: Vec<u8>) -> Result<SignResponse> {
203        let payload = SignRequest {
204            algorithm,
205            data: hex_encode(data),
206        };
207        let response = self.send_rpc_request("/Sign", &payload).await?;
208        let response = serde_json::from_value::<SignResponse>(response)?;
209        Ok(response)
210    }
211
212    /// Verifies a payload signature.
213    pub async fn verify(
214        &self,
215        algorithm: &str,
216        data: Vec<u8>,
217        signature: Vec<u8>,
218        public_key: Vec<u8>,
219    ) -> Result<VerifyResponse> {
220        let payload = VerifyRequest {
221            algorithm,
222            data: hex_encode(data),
223            signature: hex_encode(signature),
224            public_key: hex_encode(public_key),
225        };
226        let response = self.send_rpc_request("/Verify", &payload).await?;
227        let response = serde_json::from_value::<VerifyResponse>(response)?;
228        Ok(response)
229    }
230}
dstack_sdk/dstack_client.rs

dstack_sdk/
dstack_client.rs
