Skip to content

Commit c65b899

Browse files
committed
Add dstack simulator config
1 parent bf0494b commit c65b899

File tree

15 files changed

+179
-30
lines changed

15 files changed

+179
-30
lines changed

cert-client/src/lib.rs

Lines changed: 15 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -79,12 +79,23 @@ impl CertRequestClient {
7979
}
8080
}
8181

82-
pub async fn request_cert(&self, key: &KeyPair, config: CertConfig) -> Result<Vec<String>> {
82+
pub async fn request_cert(
83+
&self,
84+
key: &KeyPair,
85+
config: CertConfig,
86+
no_ra: bool,
87+
) -> Result<Vec<String>> {
8388
let pubkey = key.public_key_der();
8489
let report_data = QuoteContentType::RaTlsCert.to_report_data(&pubkey);
85-
let (_, quote) = get_quote(&report_data, None).context("Failed to get quote")?;
86-
let event_log = read_event_logs().context("Failed to decode event log")?;
87-
let event_log = serde_json::to_vec(&event_log).context("Failed to serialize event log")?;
90+
let (quote, event_log) = if !no_ra {
91+
let (_, quote) = get_quote(&report_data, None).context("Failed to get quote")?;
92+
let event_log = read_event_logs().context("Failed to decode event log")?;
93+
let event_log =
94+
serde_json::to_vec(&event_log).context("Failed to serialize event log")?;
95+
(quote, event_log)
96+
} else {
97+
(vec![], vec![])
98+
};
8899

89100
let csr = CertSigningRequest {
90101
confirm: "please sign cert:".to_string(),

guest-agent/dstack.toml

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,11 @@ compose_file = "/tapp/.host-shared/app-compose.json"
1313
public_logs = true
1414
public_sysinfo = true
1515

16+
[default.core.simulator]
17+
enabled = false
18+
quote_file = "quote.hex"
19+
event_log_file = "eventlog.json"
20+
1621
[internal-v0]
1722
address = "unix:/var/run/tappd.sock"
1823
reuse = true

guest-agent/rpc/proto/agent_rpc.proto

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -193,7 +193,7 @@ message WorkerInfo {
193193

194194
// The response to a WorkerInfo request
195195
message WorkerVersion {
196-
// Tappd version
196+
// Dstack version
197197
string version = 1;
198198
// Git revision
199199
string rev = 2;

guest-agent/src/config.rs

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,4 +22,12 @@ pub struct Config {
2222
pub compose_file: String,
2323
#[serde(default)]
2424
pub pccs_url: Option<String>,
25+
pub simulator: Simulator,
26+
}
27+
28+
#[derive(Debug, Clone, Deserialize)]
29+
pub struct Simulator {
30+
pub enabled: bool,
31+
pub quote_file: String,
32+
pub event_log_file: String,
2533
}

guest-agent/src/main.rs

Lines changed: 16 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -124,12 +124,22 @@ async fn run_guest_api(state: AppState, figment: Figment) -> Result<()> {
124124
.ignite()
125125
.await
126126
.map_err(|err| anyhow!("Failed to ignite rocket: {err}"))?;
127-
let listener = VsockListener::bind_rocket(&ignite)
128-
.map_err(|err| anyhow!("Failed to bind guest API : {err}"))?;
129-
ignite
130-
.launch_on(listener)
131-
.await
132-
.map_err(|err| anyhow!(err.to_string()))?;
127+
if DefaultListener::bind_endpoint(&ignite).is_ok() {
128+
let listener = DefaultListener::bind(&ignite)
129+
.await
130+
.map_err(|err| anyhow!("Failed to bind guest API : {err}"))?;
131+
ignite
132+
.launch_on(listener)
133+
.await
134+
.map_err(|err| anyhow!(err.to_string()))?;
135+
} else {
136+
let listener = VsockListener::bind_rocket(&ignite)
137+
.map_err(|err| anyhow!("Failed to bind guest API : {err}"))?;
138+
ignite
139+
.launch_on(listener)
140+
.await
141+
.map_err(|err| anyhow!(err.to_string()))?;
142+
}
133143
Ok(())
134144
}
135145

guest-agent/src/rpc_service.rs

Lines changed: 42 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -58,6 +58,7 @@ impl AppState {
5858
usage_client_auth: true,
5959
ext_quote: true,
6060
},
61+
config.simulator.enabled,
6162
)
6263
.await
6364
.context("Failed to get app cert")?
@@ -101,7 +102,7 @@ impl DstackGuestRpc for InternalRpcHandler {
101102
.state
102103
.inner
103104
.cert_client
104-
.request_cert(&derived_key, config)
105+
.request_cert(&derived_key, config, self.state.config().simulator.enabled)
105106
.await
106107
.context("Failed to sign the CSR")?;
107108
Ok(GetTlsKeyResponse {
@@ -145,6 +146,9 @@ impl DstackGuestRpc for InternalRpcHandler {
145146
Some(padded)
146147
}
147148
let report_data = pad64(&request.report_data).context("Report data is too long")?;
149+
if self.state.config().simulator.enabled {
150+
return simulate_quote(self.state.config(), report_data);
151+
}
148152
let (_, quote) =
149153
tdx_attest::get_quote(&report_data, None).context("Failed to get quote")?;
150154
let event_log = read_event_logs().context("Failed to decode event log")?;
@@ -162,6 +166,23 @@ impl DstackGuestRpc for InternalRpcHandler {
162166
}
163167
}
164168

169+
fn simulate_quote(config: &Config, report_data: [u8; 64]) -> Result<GetQuoteResponse> {
170+
let quote_file =
171+
fs::read_to_string(&config.simulator.quote_file).context("Failed to read quote file")?;
172+
let mut quote = hex::decode(quote_file.trim()).context("Failed to decode quote")?;
173+
let event_log = fs::read_to_string(&config.simulator.event_log_file)
174+
.context("Failed to read event log file")?;
175+
if quote.len() < 632 {
176+
return Err(anyhow::anyhow!("Quote is too short"));
177+
}
178+
quote[568..632].copy_from_slice(&report_data);
179+
Ok(GetQuoteResponse {
180+
quote,
181+
event_log,
182+
report_data: report_data.to_vec(),
183+
})
184+
}
185+
165186
impl RpcCall<AppState> for InternalRpcHandler {
166187
type PrpcService = DstackGuestServer<Self>;
167188

@@ -201,7 +222,7 @@ impl TappdRpc for InternalRpcHandlerV0 {
201222
.state
202223
.inner
203224
.cert_client
204-
.request_cert(&derived_key, config)
225+
.request_cert(&derived_key, config, self.state.config().simulator.enabled)
205226
.await
206227
.context("Failed to sign the CSR")?;
207228
Ok(GetTlsKeyResponse {
@@ -221,28 +242,37 @@ impl TappdRpc for InternalRpcHandlerV0 {
221242
}
222243

223244
async fn tdx_quote(self, request: TdxQuoteArgs) -> Result<TdxQuoteResponse> {
245+
let hash_algorithm = if request.hash_algorithm.is_empty() {
246+
DEFAULT_HASH_ALGORITHM
247+
} else {
248+
&request.hash_algorithm
249+
};
250+
let prefix = if hash_algorithm == "raw" {
251+
"".into()
252+
} else {
253+
QuoteContentType::AppData.tag().to_string()
254+
};
224255
let content_type = if request.prefix.is_empty() {
225256
QuoteContentType::AppData
226257
} else {
227258
QuoteContentType::Custom(&request.prefix)
228259
};
229260
let report_data =
230261
content_type.to_report_data_with_hash(&request.report_data, &request.hash_algorithm)?;
262+
if self.state.config().simulator.enabled {
263+
let response = simulate_quote(self.state.config(), report_data)?;
264+
return Ok(TdxQuoteResponse {
265+
quote: response.quote,
266+
event_log: response.event_log,
267+
hash_algorithm: hash_algorithm.to_string(),
268+
prefix,
269+
});
270+
}
231271
let event_log = read_event_logs().context("Failed to decode event log")?;
232272
let event_log =
233273
serde_json::to_string(&event_log).context("Failed to serialize event log")?;
234274
let (_, quote) =
235275
tdx_attest::get_quote(&report_data, None).context("Failed to get quote")?;
236-
let hash_algorithm = if request.hash_algorithm.is_empty() {
237-
DEFAULT_HASH_ALGORITHM
238-
} else {
239-
&request.hash_algorithm
240-
};
241-
let prefix = if hash_algorithm == "raw" {
242-
"".into()
243-
} else {
244-
QuoteContentType::AppData.tag().to_string()
245-
};
246276
Ok(TdxQuoteResponse {
247277
quote,
248278
event_log,

sdk/simulator/.gitignore

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
dstack-simulator
2+
dstack-guest-agent
3+
*.lock

sdk/simulator/app-compose.json

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,14 @@
1+
{
2+
"manifest_version": 2,
3+
"name": "kvin-nb",
4+
"runner": "docker-compose",
5+
"docker_compose_file": "services:\n jupyter:\n image: quay.io/jupyter/base-notebook\n user: root\n environment:\n - GRANT_SUDO=yes\n ports:\n - \"8888:8888\"\n volumes:\n - /:/host/\n - /var/run/tappd.sock:/var/run/tappd.sock\n - /var/run/dstack.sock:/var/run/dstack.sock\n logging:\n driver: journald\n options:\n tag: jupyter-notebook\n",
6+
"docker_config": {},
7+
"kms_enabled": true,
8+
"tproxy_enabled": true,
9+
"public_logs": true,
10+
"public_sysinfo": true,
11+
"local_key_provider_enabled": false,
12+
"allowed_envs": [],
13+
"no_instance_id": false
14+
}

sdk/simulator/appkeys.json

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
{
2+
"disk_crypt_key": "1122e1f340c19407adc5ec531ac98d72bcf702bf7858f6fa49b5be79b61e4d5b",
3+
"env_crypt_key": "ca1a3895d9d613287fc14034d0ec60abb5089896e7c8fd7c2f02bd91fa0076aa",
4+
"k256_key": "e0e5d254fb944dcc370a2e5288b336a1e809871545a73ee645368957fefa31f9",
5+
"k256_signature": "2f431c7956869a4fe3e028c5f9518a935e2d01e81a3628f8b1d178fc2fac7b6d2405ace433624e5568e23c4ed291dbaf60dac79b756837c0fe745154ebfdc0a601",
6+
"gateway_app_id": "any",
7+
"ca_cert": "-----BEGIN CERTIFICATE-----\nMIIBmTCCAUCgAwIBAgIUU7801+krCs2OpIdne3t6OWrJ2fMwCgYIKoZIzj0EAwIw\nKTEPMA0GA1UECgwGRHN0YWNrMRYwFAYDVQQDDA1Ec3RhY2sgS01TIENBMB4XDTc1\nMDEwMTAwMDAwMFoXDTM1MDMxNzA5NDQ0MlowKTEPMA0GA1UECgwGRHN0YWNrMRYw\nFAYDVQQDDA1Ec3RhY2sgS01TIENBMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAE\nGbJFfdm4qmRG2YDxNv/3gS7NbHd0DusOKLENVsDAACiltuWdzqMH1YO9H3B2npwR\nbfK8+xdYqV2GE+feHISCwKNGMEQwDwYDVR0PAQH/BAUDAweAADAdBgNVHQ4EFgQU\nevjJ+VZPvDxHJ2ejjeIaUYMMcEcwEgYDVR0TAQH/BAgwBgEB/wIBATAKBggqhkjO\nPQQDAgNHADBEAiAhQHQNbmyvx9BDBXRjW1eCkPCpFs/2Vt/nvbi+M69FPAIgQ13F\n3pmxicxyFeVW2iOjrbG1cxLdT9Kh+9ICF9zn8kA=\n-----END CERTIFICATE-----\n",
8+
"key_provider": {
9+
"Local": {
10+
"key": "-----BEGIN PRIVATE KEY-----\nMIGHAgEAMBMGByqGSM49AgEGCCqGSM49AwEHBG0wawIBAQQg1PYCFKYfDmUfv5fk\nstppasf4mPGqnz0fEoLEnGx8CnKhRANCAAQZskV92biqZEbZgPE2//eBLs1sd3QO\n6w4osQ1WwMAAKKW25Z3OowfVg70fcHaenBFt8rz7F1ipXYYT594chILA\n-----END PRIVATE KEY-----\n"
11+
}
12+
}
13+
}

sdk/simulator/build.sh

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
#!/bin/bash
2+
cd $(dirname $0)
3+
cargo build --release -p dstack-guest-agent
4+
cp ../../target/release/dstack-guest-agent .
5+
ln -sf dstack-guest-agent dstack-simulator
6+

0 commit comments

Comments
 (0)