Skip to content

Commit 476cd6b

Browse files
kvinwangclaude
andcommitted
guest-agent: refuse a malformed token before paying for identity
The scope fairing resolved the app id before looking at the token, so any string in an Authorization header could drive a boot attestation and queue behind attest_lock. Split `jwt::decode` out of `Verifier::verify` so the fairing reads the header and claims first, and only then asks for the app id. `verify_decoded` checks the signature before the claims: a claim error from an unsigned token describes a guess at the caller's intent, not the token. Also: - the dashboard skips `Worker.Info` unless a block would render it, so an anonymous page view no longer costs an attestation; - the watchdog builds one HTTP client instead of one per heartbeat; - `EffectiveScopes` answers 401, not 403, when the fairing published nothing; - the gateway drops an `agent_token` longer than 4 KiB with a warning: it is self-declared, persisted and synced, and must not be sized by the CVM. Co-Authored-By: Claude Code <noreply@anthropic.com> Signed-off-by: Kevin Wang <wy721@qq.com>
1 parent 23a6ed6 commit 476cd6b

7 files changed

Lines changed: 253 additions & 48 deletions

File tree

‎dstack/crates/api-auth/src/jwt.rs‎

Lines changed: 95 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -145,6 +145,48 @@ pub fn internal_issuer(app_id_hex: &str) -> String {
145145
format!("dstack:guest-api:{app_id_hex}")
146146
}
147147

148+
/// A token whose shape and claims have been read, with the signature still
149+
/// unchecked.
150+
///
151+
/// Split out from verification so a caller can refuse a malformed token before
152+
/// it does anything expensive. The guest agent needs exactly that: checking a
153+
/// token it was handed must not cost a boot attestation.
154+
#[derive(Clone, Debug)]
155+
pub struct Decoded {
156+
signing_input: String,
157+
signature: Vec<u8>,
158+
pub claims: Claims,
159+
}
160+
161+
/// Read a token's header and claims without verifying it.
162+
///
163+
/// Everything here is cheap and local. The result says nothing about whether
164+
/// the token was signed by anyone; [`Verifier::verify_decoded`] answers that.
165+
pub fn decode(token: &str) -> Result<Decoded> {
166+
let mut parts = token.split('.');
167+
let (Some(header), Some(encoded_payload), Some(encoded_signature), None) =
168+
(parts.next(), parts.next(), parts.next(), parts.next())
169+
else {
170+
bail!("token is not a JWT");
171+
};
172+
if header != HEADER {
173+
// Covers `alg: none`, HMAC, and any header carrying key material.
174+
bail!("unsupported JWT header");
175+
}
176+
let payload = URL_SAFE_NO_PAD
177+
.decode(encoded_payload)
178+
.context("payload is not base64url")?;
179+
let claims = serde_json::from_slice(&payload).context("malformed claims")?;
180+
let signature = URL_SAFE_NO_PAD
181+
.decode(encoded_signature)
182+
.context("signature is not base64url")?;
183+
Ok(Decoded {
184+
signing_input: format!("{header}.{encoded_payload}"),
185+
signature,
186+
claims,
187+
})
188+
}
189+
148190
/// Mint a token. `signing_key` is the raw 32-byte Ed25519 seed.
149191
pub fn sign(signing_key: &[u8; 32], claims: &Claims) -> Result<String> {
150192
let payload = serde_json::to_vec(claims).context("failed to serialize claims")?;
@@ -184,20 +226,37 @@ impl Verifier {
184226
/// `now` is seconds since the UNIX epoch, passed in rather than read here
185227
/// so the expiry rules can be tested.
186228
pub fn verify(&self, token: &str, app_id_hex: &str, now: u64) -> Result<ScopeSet> {
187-
let mut parts = token.split('.');
188-
let (Some(header), Some(encoded_payload), Some(signature), None) =
189-
(parts.next(), parts.next(), parts.next(), parts.next())
190-
else {
191-
bail!("token is not a JWT");
192-
};
193-
if header != HEADER {
194-
// Covers `alg: none`, HMAC, and any header carrying key material.
195-
bail!("unsupported JWT header");
229+
self.verify_decoded(&decode(token)?, app_id_hex, now)
230+
}
231+
232+
/// Check the signature and claims of a token [`decode`] already read.
233+
///
234+
/// Separate from [`Verifier::verify`] so a caller can pay for [`decode`]
235+
/// alone -- it needs no key material and no app id -- and refuse the token
236+
/// before doing anything that costs.
237+
///
238+
/// The signature is checked before the claims: a claim error from an
239+
/// unsigned token is a guess about a caller's intent, not a fact.
240+
pub fn verify_decoded(&self, token: &Decoded, app_id_hex: &str, now: u64) -> Result<ScopeSet> {
241+
let Decoded {
242+
signing_input,
243+
signature,
244+
claims,
245+
} = token;
246+
let signature =
247+
Signature::from_slice(signature).context("signature has the wrong length")?;
248+
249+
let keys = self.keys_for(&claims.iss, app_id_hex);
250+
if keys.is_empty() {
251+
bail!("unknown token issuer");
252+
}
253+
if !keys.iter().any(|key| {
254+
key.verify_strict(signing_input.as_bytes(), &signature)
255+
.is_ok()
256+
}) {
257+
bail!("token signature does not verify");
196258
}
197-
let payload = URL_SAFE_NO_PAD
198-
.decode(encoded_payload)
199-
.context("payload is not base64url")?;
200-
let claims: Claims = serde_json::from_slice(&payload).context("malformed claims")?;
259+
201260
if claims.aud != app_id_hex {
202261
bail!("token was issued for a different application");
203262
}
@@ -210,23 +269,6 @@ impl Verifier {
210269
if claims.iat > now + CLOCK_LEEWAY_SECS {
211270
bail!("token was issued in the future");
212271
}
213-
let signature = URL_SAFE_NO_PAD
214-
.decode(signature)
215-
.context("signature is not base64url")?;
216-
let signature =
217-
Signature::from_slice(&signature).context("signature has the wrong length")?;
218-
let signing_input = format!("{header}.{encoded_payload}");
219-
220-
let keys = self.keys_for(&claims.iss, app_id_hex);
221-
if keys.is_empty() {
222-
bail!("unknown token issuer");
223-
}
224-
if !keys.iter().any(|key| {
225-
key.verify_strict(signing_input.as_bytes(), &signature)
226-
.is_ok()
227-
}) {
228-
bail!("token signature does not verify");
229-
}
230272
Ok(claims.scopes())
231273
}
232274

@@ -353,6 +395,29 @@ mod tests {
353395
assert!(verifier(1).verify(&future, APP_ID, 1000).is_err());
354396
}
355397

398+
#[test]
399+
fn decoding_refuses_a_shape_no_token_could_have() {
400+
for token in [
401+
"not-a-token",
402+
"a.b.c.d",
403+
"eyJhbGciOiJub25lIiwidHlwIjoiSldUIn0.eyJpc3MiOiJ4In0.",
404+
"eyJhbGciOiJFZERTQSIsInR5cCI6IkpXVCJ9.!!!.c2ln",
405+
] {
406+
assert!(decode(token).is_err(), "{token}");
407+
}
408+
// A well-formed token decodes without any key: that is what lets the
409+
// caller refuse it before doing work.
410+
let token = mint(
411+
1,
412+
&internal_issuer(APP_ID),
413+
APP_ID,
414+
1000,
415+
1300,
416+
"guest:info",
417+
);
418+
assert_eq!(decode(&token).unwrap().claims.aud, APP_ID);
419+
}
420+
356421
#[test]
357422
fn accepts_an_external_issuer_by_name() {
358423
let verifier = Verifier::new(

‎dstack/gateway/src/main_service.rs‎

Lines changed: 29 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2757,6 +2757,33 @@ fn gateway_nodes(kv_store: &KvStore, exclude_down: bool) -> Vec<GatewayNodeInfo>
27572757
.collect()
27582758
}
27592759

2760+
/// Longest `agent_token` a registration may carry.
2761+
///
2762+
/// The token is a short Ed25519 JWT -- a few hundred bytes. The field is
2763+
/// self-declared, persisted, synced to every peer and sent back out as a
2764+
/// request header, so it is bounded here rather than trusted.
2765+
const MAX_AGENT_TOKEN_BYTES: usize = 4096;
2766+
2767+
/// The `agent_token` a registration may store, or `None`.
2768+
///
2769+
/// Absent and empty both mean "this CVM has no token", which is what a CVM
2770+
/// that predates the field or could not mint one reports. A value past
2771+
/// [`MAX_AGENT_TOKEN_BYTES`] is not a token at all: it is dropped with a
2772+
/// warning rather than stored, so one CVM cannot inflate every record.
2773+
fn validated_agent_token(instance_id: &str, token: Option<String>) -> Option<String> {
2774+
match token {
2775+
Some(token) if token.len() > MAX_AGENT_TOKEN_BYTES => {
2776+
warn!(
2777+
"[{instance_id}] ignoring an agent token of {} bytes (limit {MAX_AGENT_TOKEN_BYTES})",
2778+
token.len()
2779+
);
2780+
None
2781+
}
2782+
Some(token) if token.is_empty() => None,
2783+
token => token,
2784+
}
2785+
}
2786+
27602787
pub struct RpcHandler {
27612788
remote_app_id: Option<Vec<u8>>,
27622789
remote_app_info: Option<AppInfo>,
@@ -2817,6 +2844,7 @@ impl GatewayRpc for RpcHandler {
28172844
})
28182845
})
28192846
.transpose()?;
2847+
let agent_token = validated_agent_token(&instance_id, request.agent_token);
28202848
self.state
28212849
.register_cvm_async(
28222850
app_id,
@@ -2827,7 +2855,7 @@ impl GatewayRpc for RpcHandler {
28272855
port_policy,
28282856
// A real CVM always states its intent, either way.
28292857
health_check: Some(request.health_check),
2830-
agent_token: request.agent_token.filter(|token| !token.is_empty()),
2858+
agent_token,
28312859
},
28322860
)
28332861
.await

‎dstack/guest-agent/src/api_auth.rs‎

Lines changed: 44 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@
66
//!
77
//! Every request that reaches `[external]` names a scope, and the caller has
88
//! to be granted it: either by an Ed25519 JWT, or -- when the host configured
9-
//! it -- by being anonymous. Nothing is granted by default.
9+
//! it -- by being anonymous. Nothing is granted without one or the other.
1010
//!
1111
//! This is deliberately not a trust boundary. The measured `public_*` switches
1212
//! in app-compose still decide whether confidential data (logs, sysinfo, TCB
@@ -17,7 +17,7 @@
1717
use std::time::{SystemTime, UNIX_EPOCH};
1818

1919
use anyhow::{Context, Result};
20-
use dstack_api_auth::jwt::{self, IssuerKeys, ScopeSet, Verifier};
20+
use dstack_api_auth::jwt::{self, Decoded, IssuerKeys, ScopeSet, Verifier};
2121
use dstack_types::GuestAgentApiConfig;
2222
use ed25519_dalek::VerifyingKey;
2323
use rocket::{
@@ -94,7 +94,7 @@ impl GuestApiAuth {
9494
/// authenticate and got it wrong should see a refusal, not whatever the
9595
/// host happened to leave open. `app_id` is `None` while the app identity
9696
/// is unavailable, in which case no token can be checked.
97-
fn resolve(&self, token: Option<&str>, app_id: Option<&str>) -> Result<ScopeSet> {
97+
fn resolve(&self, token: Option<&Decoded>, app_id: Option<&str>) -> Result<ScopeSet> {
9898
let Some(token) = token else {
9999
return Ok(self.anonymous.clone());
100100
};
@@ -105,7 +105,7 @@ impl GuestApiAuth {
105105
.unwrap_or(0);
106106
let granted = self
107107
.verifier
108-
.verify(token, app_id, now)
108+
.verify_decoded(token, app_id, now)
109109
.context("rejected bearer token")?;
110110
Ok(self.anonymous.union(&granted))
111111
}
@@ -139,7 +139,11 @@ impl<'r> FromRequest<'r> for EffectiveScopes {
139139
async fn from_request(request: &'r Request<'_>) -> Outcome<Self, ()> {
140140
match request.local_cache(|| None::<ScopeSet>) {
141141
Some(scopes) => Outcome::Success(EffectiveScopes(scopes.clone())),
142-
None => Outcome::Forward(rocket::http::Status::Forbidden),
142+
// Unreachable: `ScopeFairing` either publishes a set or rewrites
143+
// the request to the 401 route. Kept as a 401 rather than a 403 so
144+
// a request that somehow got past the fairing still looks like
145+
// every other refusal.
146+
None => Outcome::Error((rocket::http::Status::Unauthorized, ())),
143147
}
144148
}
145149
}
@@ -206,7 +210,17 @@ impl Fairing for ScopeFairing {
206210
}
207211

208212
async fn on_request(&self, req: &mut Request<'_>, _: &mut Data<'_>) {
209-
let token = bearer_token(req);
213+
// Read the token's shape first. It is the only cheap step, and it is
214+
// the one an unauthenticated caller can drive; the app id lookup below
215+
// can cost a boot attestation, so nothing that arrives here gets that
216+
// far unless it is at least a well-formed Ed25519 token.
217+
let token = match bearer_token(req).map(jwt::decode).transpose() {
218+
Ok(token) => token,
219+
Err(err) => {
220+
tracing::debug!("{err:#}");
221+
return deny(req);
222+
}
223+
};
210224
// Only a token needs the app id. Going through `identity()` rather
211225
// than a value bound at startup is what lets a boot whose first decode
212226
// failed recover: the retry is driven by the callers that need it.
@@ -219,10 +233,16 @@ impl Fairing for ScopeFairing {
219233
.map(|identity| hex::encode(&identity.app_id)),
220234
None => None,
221235
};
222-
let scopes = match self.state.api_auth().resolve(token, app_id.as_deref()) {
236+
let scopes = match self
237+
.state
238+
.api_auth()
239+
.resolve(token.as_ref(), app_id.as_deref())
240+
{
223241
Ok(scopes) => scopes,
224242
Err(err) => {
225-
warn!("{err:#}");
243+
// Every caller on the internet reaches this. A warning per
244+
// attempt is a log the public gets to size.
245+
tracing::debug!("{err:#}");
226246
return deny(req);
227247
}
228248
};
@@ -301,6 +321,13 @@ mod tests {
301321
.unwrap()
302322
}
303323

324+
/// Stands in for a token that is well-formed but not signed by anything
325+
/// this agent trusts.
326+
fn decoded(claims: jwt::Claims) -> Decoded {
327+
let token = jwt::sign(&[1u8; 32], &claims).unwrap();
328+
jwt::decode(&token).unwrap()
329+
}
330+
304331
#[test]
305332
fn anonymous_grants_only_what_was_configured() {
306333
assert!(auth(&[]).resolve(None, None).unwrap().is_empty());
@@ -311,8 +338,16 @@ mod tests {
311338

312339
#[test]
313340
fn a_bad_token_is_refused_rather_than_downgraded_to_anonymous() {
341+
let claims = jwt::Claims {
342+
iss: jwt::internal_issuer("00"),
343+
aud: "00".into(),
344+
iat: 0,
345+
exp: u64::MAX / 2,
346+
scope: jwt::scope::DASHBOARD.into(),
347+
};
348+
// Signed by the wrong key, so it decodes and fails to verify.
314349
assert!(auth(&["guest:dashboard"])
315-
.resolve(Some("not-a-token"), Some("00"))
350+
.resolve(Some(&decoded(claims)), Some("00"))
316351
.is_err());
317352
}
318353

‎dstack/guest-agent/src/http_routes.rs‎

Lines changed: 15 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -51,8 +51,11 @@ async fn index(
5151
let show_logs = public_logs && scopes.allows(jwt::scope::LOGS);
5252
let show_tcb = public_tcbinfo && show_identity;
5353
let context = CallContext::builder().state(&**state).build();
54-
let handler = ExternalRpcHandler::construct(context.clone())
55-
.map_err(|e| format!("Failed to construct RPC handler: {}", e))?;
54+
// `Info` costs a boot attestation (and may start a demo-cert request), so
55+
// it is only asked for when the identity block -- and with it the TCB
56+
// block, which requires `guest:info` too -- will be rendered.
57+
// `guest:dashboard` is the one scope an anonymous visitor is most likely
58+
// to be given, and on its own it renders every block hidden.
5659
let AppInfo {
5760
app_name,
5861
app_id,
@@ -67,10 +70,16 @@ async fn index(
6770
vm_config: _,
6871
cloud_vendor,
6972
cloud_product,
70-
} = handler
71-
.info()
72-
.await
73-
.map_err(|e| format!("Failed to get worker info: {}", e))?;
73+
} = if show_identity || show_tcb {
74+
let handler = ExternalRpcHandler::construct(context.clone())
75+
.map_err(|e| format!("Failed to construct RPC handler: {}", e))?;
76+
handler
77+
.info()
78+
.await
79+
.map_err(|e| format!("Failed to get worker info: {}", e))?
80+
} else {
81+
Default::default()
82+
};
7483

7584
// Only collected when the page will actually render it. Collecting
7685
// otherwise would scan every process and disk, and spawn a GPU collector,

0 commit comments

Comments
 (0)