From 9e9c30b751f27e5af97496ade23f9602764cbb90 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Sun, 20 Sep 2026 08:17:51 -0700 Subject: [PATCH 1/2] fix(kms): give the auth API call a timeout and one client `KMS.GetMeta` takes no client certificate. `ra_rpc::ratls_client_verifier` makes client auth optional, the handler checks nothing, and the RPC calls the auth API on every request with no decision cache. So an anonymous caller on the public RPC listener decides how many of these are in flight -- `GetAppEnvEncryptPubKey` and `GetTempCaCert` reach the same backend the same way. Each call built a fresh `reqwest::Client`, which sets no timeout of any kind. A backend that accepts the connection and never answers therefore held its KMS request open for the life of the process; the test below sat on one for ten seconds and gave up before the client did. `ra_rpc::client` already sets connect and request timeouts for the other direction, so this was the one outbound path without them. The request timeout is config, because it is what an operator has to match to their backend, and it is defaulted so an existing `[core.auth_api.webhook]` naming only `url` keeps parsing. The connect timeout is not: a TCP connect that has not completed in five seconds is a backend that is down, not one that is slow. Building a client per call also threw away the connection pool and the TLS configuration every time, so every authorization decision paid a fresh connect -- and over HTTPS a fresh handshake. One client now serves them all. Signed-off-by: Kevin Wang --- dstack/kms/kms.toml | 5 ++ dstack/kms/src/config.rs | 14 ++++ .../kms/src/main_service/upgrade_authority.rs | 83 +++++++++++++++++-- 3 files changed, 95 insertions(+), 7 deletions(-) diff --git a/dstack/kms/kms.toml b/dstack/kms/kms.toml index 5beb28d80..e5e8775a5 100644 --- a/dstack/kms/kms.toml +++ b/dstack/kms/kms.toml @@ -101,6 +101,11 @@ type = "webhook" [core.auth_api.webhook] url = "http://auth-api:8000" +# How long one authorization call may take. GetMeta and GetAppEnvEncryptPubKey +# reach this backend without a client certificate, so an anonymous caller picks +# how many of these are in flight; a backend that accepts the connection and +# never answers would otherwise hold each one for the life of the process. +timeout = "60s" [core.auth_api.dev] gateway_app_id = "any" diff --git a/dstack/kms/src/config.rs b/dstack/kms/src/config.rs index ba3f9f93e..4ad09a46c 100644 --- a/dstack/kms/src/config.rs +++ b/dstack/kms/src/config.rs @@ -177,6 +177,20 @@ impl AuthApi { #[derive(Debug, Clone, Deserialize)] pub(crate) struct Webhook { pub url: String, + /// How long one auth API call may take. + /// + /// Defaulted rather than required: an existing `[core.auth_api.webhook]` + /// section names only `url`, and a config that stops parsing on upgrade is + /// worse than one that keeps the old default. `serde` fills this in for + /// them. + #[serde(default = "default_webhook_timeout", with = "serde_duration")] + pub timeout: Duration, +} + +/// Matches the RPC client's request timeout (`ra_rpc::client`), so a KMS +/// waiting on its auth API gives up no later than the guest waiting on the KMS. +fn default_webhook_timeout() -> Duration { + Duration::from_secs(60) } #[derive(Debug, Clone, Deserialize)] diff --git a/dstack/kms/src/main_service/upgrade_authority.rs b/dstack/kms/src/main_service/upgrade_authority.rs index 8b1c20742..abd5eaef1 100644 --- a/dstack/kms/src/main_service/upgrade_authority.rs +++ b/dstack/kms/src/main_service/upgrade_authority.rs @@ -12,6 +12,8 @@ use http_client::prpc::PrpcClient; use ra_tls::attestation::{AttestationVerifier, VerifiedAttestation, VersionedAttestation}; use serde::de::DeserializeOwned; use serde::{Deserialize, Serialize}; +use std::sync::OnceLock; +use std::time::Duration; /// The KMS `bootAuth` payload. This is the verifier's `PolicyBootInfo` — the one /// canonical struct shared by the producer (KMS) and the policy input the @@ -94,12 +96,43 @@ pub(crate) struct GetInfoResponse { pub app_implementation: Option, } -async fn http_get(url: &str) -> Result { - send_request(reqwest::Client::new().get(url), url).await +/// How long the auth API has to accept a connection. +/// +/// Separate from the configured request timeout and not configurable: a TCP +/// connect that has not completed in this long is a backend that is down, not +/// one that is slow, and no deployment needs to tune that apart from the rest. +const AUTH_API_CONNECT_TIMEOUT: Duration = Duration::from_secs(5); + +/// The one client every auth API call goes through. +/// +/// A `reqwest::Client` owns the connection pool and the TLS configuration, so +/// building one per call threw both away and paid a fresh connect -- and, over +/// HTTPS, a fresh handshake -- for every authorization decision. There is +/// deliberately no decision cache above this (pinned by a test), which is what +/// makes the pool worth keeping. +fn auth_api_client() -> &'static reqwest::Client { + static CLIENT: OnceLock = OnceLock::new(); + CLIENT.get_or_init(|| { + reqwest::Client::builder() + .connect_timeout(AUTH_API_CONNECT_TIMEOUT) + .build() + // A builder that cannot produce a client means no TLS backend; the + // default client is the same thing without our connect timeout, and + // failing every call from here would be worse than losing it. + .unwrap_or_default() + }) +} + +async fn http_get(url: &str, timeout: Duration) -> Result { + send_request(auth_api_client().get(url).timeout(timeout), url).await } -async fn http_post(url: &str, body: &impl Serialize) -> Result { - send_request(reqwest::Client::new().post(url).json(body), url).await +async fn http_post( + url: &str, + body: &impl Serialize, + timeout: Duration, +) -> Result { + send_request(auth_api_client().post(url).json(body).timeout(timeout), url).await } async fn send_request(req: reqwest::RequestBuilder, url: &str) -> Result { @@ -131,7 +164,7 @@ impl AuthApi { "bootAuth/app" }; let url = url_join(&webhook.url, path); - http_post(&url, &boot_info).await + http_post(&url, &boot_info, webhook.timeout).await } } } @@ -147,7 +180,7 @@ impl AuthApi { app_implementation: None, }), AuthApi::Webhook { webhook } => { - let info: AuthApiInfoResponse = http_get(&webhook.url).await?; + let info: AuthApiInfoResponse = http_get(&webhook.url, webhook.timeout).await?; let eth_rpc_url = if info.eth_rpc_url.is_empty() { None } else { @@ -328,10 +361,46 @@ mod tests { fn webhook(url: String) -> AuthApi { AuthApi::Webhook { - webhook: Webhook { url }, + webhook: Webhook { + url, + timeout: std::time::Duration::from_millis(500), + }, } } + /// `KMS.GetMeta` needs no client certificate (`ra_rpc::ratls_client_verifier` + /// makes client auth optional and the handler checks nothing), and it calls + /// the auth API on every request with no decision cache. So an anonymous + /// caller on the public RPC listener decides how many of these are in + /// flight. A backend that accepts the connection and never answers must not + /// keep each one open for the life of the process. + #[rocket::async_test] + async fn a_backend_that_never_answers_does_not_hold_the_request_open() { + let listener = TcpListener::bind("127.0.0.1:0").unwrap(); + let address = listener.local_addr().unwrap(); + let _blackhole = thread::spawn(move || { + // Accept, then hold the socket open without ever writing a byte. + let held: Vec<_> = listener.incoming().take(1).collect(); + thread::sleep(std::time::Duration::from_secs(30)); + drop(held); + }); + + let auth = webhook(format!("http://{address}")); + let outcome = + tokio::time::timeout(std::time::Duration::from_secs(10), auth.get_info()).await; + + assert!( + outcome.is_ok(), + "the auth API call carried no timeout of its own and was still \ + running after 10s; an anonymous GetMeta flood accumulates these" + ); + let error = outcome.unwrap().expect_err("a silent backend must fail"); + assert!( + format!("{error:#}").contains("timed out") || format!("{error:#}").contains("timeout"), + "unexpected error: {error:#}" + ); + } + #[test] fn app_id_len_must_be_20_bytes() { assert!(ensure_app_id_len(&[0u8; 20]).is_ok()); From 8092ae5c2d6535c18bb58d0edc6882e02dbf24e8 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 24 Sep 2026 02:01:11 -0700 Subject: [PATCH 2/2] refactor(kms): trim auth API timeout comments --- dstack/kms/kms.toml | 5 +--- dstack/kms/src/config.rs | 10 ++----- .../kms/src/main_service/upgrade_authority.rs | 27 +++---------------- 3 files changed, 6 insertions(+), 36 deletions(-) diff --git a/dstack/kms/kms.toml b/dstack/kms/kms.toml index e5e8775a5..b06e95a54 100644 --- a/dstack/kms/kms.toml +++ b/dstack/kms/kms.toml @@ -101,10 +101,7 @@ type = "webhook" [core.auth_api.webhook] url = "http://auth-api:8000" -# How long one authorization call may take. GetMeta and GetAppEnvEncryptPubKey -# reach this backend without a client certificate, so an anonymous caller picks -# how many of these are in flight; a backend that accepts the connection and -# never answers would otherwise hold each one for the life of the process. +# Per-call timeout; GetMeta reaches the auth API without a client certificate. timeout = "60s" [core.auth_api.dev] diff --git a/dstack/kms/src/config.rs b/dstack/kms/src/config.rs index 4ad09a46c..4ee87e2f8 100644 --- a/dstack/kms/src/config.rs +++ b/dstack/kms/src/config.rs @@ -177,18 +177,12 @@ impl AuthApi { #[derive(Debug, Clone, Deserialize)] pub(crate) struct Webhook { pub url: String, - /// How long one auth API call may take. - /// - /// Defaulted rather than required: an existing `[core.auth_api.webhook]` - /// section names only `url`, and a config that stops parsing on upgrade is - /// worse than one that keeps the old default. `serde` fills this in for - /// them. + /// Per-call timeout. Defaulted so an existing section with only `url` still parses. #[serde(default = "default_webhook_timeout", with = "serde_duration")] pub timeout: Duration, } -/// Matches the RPC client's request timeout (`ra_rpc::client`), so a KMS -/// waiting on its auth API gives up no later than the guest waiting on the KMS. +/// Matches the `ra_rpc::client` request timeout. fn default_webhook_timeout() -> Duration { Duration::from_secs(60) } diff --git a/dstack/kms/src/main_service/upgrade_authority.rs b/dstack/kms/src/main_service/upgrade_authority.rs index abd5eaef1..ce604d2b3 100644 --- a/dstack/kms/src/main_service/upgrade_authority.rs +++ b/dstack/kms/src/main_service/upgrade_authority.rs @@ -96,29 +96,15 @@ pub(crate) struct GetInfoResponse { pub app_implementation: Option, } -/// How long the auth API has to accept a connection. -/// -/// Separate from the configured request timeout and not configurable: a TCP -/// connect that has not completed in this long is a backend that is down, not -/// one that is slow, and no deployment needs to tune that apart from the rest. const AUTH_API_CONNECT_TIMEOUT: Duration = Duration::from_secs(5); -/// The one client every auth API call goes through. -/// -/// A `reqwest::Client` owns the connection pool and the TLS configuration, so -/// building one per call threw both away and paid a fresh connect -- and, over -/// HTTPS, a fresh handshake -- for every authorization decision. There is -/// deliberately no decision cache above this (pinned by a test), which is what -/// makes the pool worth keeping. +/// Shared so auth API calls reuse pooled connections instead of reconnecting per decision. fn auth_api_client() -> &'static reqwest::Client { static CLIENT: OnceLock = OnceLock::new(); CLIENT.get_or_init(|| { reqwest::Client::builder() .connect_timeout(AUTH_API_CONNECT_TIMEOUT) .build() - // A builder that cannot produce a client means no TLS backend; the - // default client is the same thing without our connect timeout, and - // failing every call from here would be worse than losing it. .unwrap_or_default() }) } @@ -368,18 +354,12 @@ mod tests { } } - /// `KMS.GetMeta` needs no client certificate (`ra_rpc::ratls_client_verifier` - /// makes client auth optional and the handler checks nothing), and it calls - /// the auth API on every request with no decision cache. So an anonymous - /// caller on the public RPC listener decides how many of these are in - /// flight. A backend that accepts the connection and never answers must not - /// keep each one open for the life of the process. + /// `GetMeta` reaches the auth API unauthenticated, so a silent backend must not hold it open. #[rocket::async_test] async fn a_backend_that_never_answers_does_not_hold_the_request_open() { let listener = TcpListener::bind("127.0.0.1:0").unwrap(); let address = listener.local_addr().unwrap(); let _blackhole = thread::spawn(move || { - // Accept, then hold the socket open without ever writing a byte. let held: Vec<_> = listener.incoming().take(1).collect(); thread::sleep(std::time::Duration::from_secs(30)); drop(held); @@ -391,8 +371,7 @@ mod tests { assert!( outcome.is_ok(), - "the auth API call carried no timeout of its own and was still \ - running after 10s; an anonymous GetMeta flood accumulates these" + "the auth API call was still running after 10s" ); let error = outcome.unwrap().expect_err("a silent backend must fail"); assert!(