From 79b8b8d24076541730e4ee4fccd7cb38a3f691f1 Mon Sep 17 00:00:00 2001 From: Leechael Yim Date: Tue, 3 Mar 2026 23:20:03 +0800 Subject: [PATCH 1/2] fix(kms): set user-agent for auth api requests and improve error logging --- kms/src/main_service/upgrade_authority.rs | 24 ++++++++++++++++++----- 1 file changed, 19 insertions(+), 5 deletions(-) diff --git a/kms/src/main_service/upgrade_authority.rs b/kms/src/main_service/upgrade_authority.rs index 169fd495..aefb2195 100644 --- a/kms/src/main_service/upgrade_authority.rs +++ b/kms/src/main_service/upgrade_authority.rs @@ -3,7 +3,7 @@ // SPDX-License-Identifier: Apache-2.0 use crate::config::AuthApi; -use anyhow::{bail, Result}; +use anyhow::{bail, Context, Result}; use ra_tls::attestation::AttestationMode; use serde::{Deserialize, Serialize}; use serde_human_bytes as hex_bytes; @@ -60,6 +60,14 @@ pub(crate) struct GetInfoResponse { pub app_implementation: Option, } +fn http_client() -> Result { + static USER_AGENT: &str = concat!("dstack-kms/", env!("CARGO_PKG_VERSION")); + reqwest::Client::builder() + .user_agent(USER_AGENT) + .build() + .context("failed to build http client") +} + impl AuthApi { pub async fn is_app_allowed(&self, boot_info: &BootInfo, is_kms: bool) -> Result { match self { @@ -69,7 +77,7 @@ impl AuthApi { gateway_app_id: dev.gateway_app_id.clone(), }), AuthApi::Webhook { webhook } => { - let client = reqwest::Client::new(); + let client = http_client()?; let path = if is_kms { "bootAuth/kms" } else { @@ -95,10 +103,16 @@ impl AuthApi { app_implementation: None, }), AuthApi::Webhook { webhook } => { - let client = reqwest::Client::new(); + let client = http_client()?; let response = client.get(&webhook.url).send().await?; - println!("url: {}", webhook.url); - let info: AuthApiInfoResponse = response.json().await?; + let status = response.status(); + let body = response.text().await?; + let info: AuthApiInfoResponse = serde_json::from_str(&body).with_context(|| { + format!( + "failed to decode auth api response from {}, status={status}, body={body}", + webhook.url + ) + })?; Ok(GetInfoResponse { is_dev: false, kms_contract_address: Some(info.kms_contract_addr.clone()), From 802e80ab8b659f61e8e9d1fbada1fe6a77da8315 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Wed, 4 Mar 2026 00:51:40 +0000 Subject: [PATCH 2/2] refactor(kms): extract generic http_get/http_post helpers - Unify User-Agent, status check, body truncation, and decode error context into a single send_request function - Both is_app_allowed and get_info now share consistent error handling - Truncate response body to 512 bytes in error messages to avoid noisy HTML pages in logs - Remove separate http_client() function --- kms/src/main_service/upgrade_authority.rs | 43 ++++++++++++----------- 1 file changed, 22 insertions(+), 21 deletions(-) diff --git a/kms/src/main_service/upgrade_authority.rs b/kms/src/main_service/upgrade_authority.rs index aefb2195..b461e8ad 100644 --- a/kms/src/main_service/upgrade_authority.rs +++ b/kms/src/main_service/upgrade_authority.rs @@ -5,6 +5,7 @@ use crate::config::AuthApi; use anyhow::{bail, Context, Result}; use ra_tls::attestation::AttestationMode; +use serde::de::DeserializeOwned; use serde::{Deserialize, Serialize}; use serde_human_bytes as hex_bytes; @@ -60,12 +61,26 @@ pub(crate) struct GetInfoResponse { pub app_implementation: Option, } -fn http_client() -> Result { +async fn http_get(url: &str) -> Result { + send_request(reqwest::Client::new().get(url), 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 send_request(req: reqwest::RequestBuilder, url: &str) -> Result { static USER_AGENT: &str = concat!("dstack-kms/", env!("CARGO_PKG_VERSION")); - reqwest::Client::builder() - .user_agent(USER_AGENT) - .build() - .context("failed to build http client") + let response = req.header("User-Agent", USER_AGENT).send().await?; + let status = response.status(); + let body = response.text().await?; + let short_body = &body[..body.len().min(512)]; + if !status.is_success() { + bail!("auth api {url} returned {status}: {short_body}"); + } + serde_json::from_str(&body).with_context(|| { + format!("failed to decode response from {url}, status={status}, body={short_body}") + }) } impl AuthApi { @@ -77,18 +92,13 @@ impl AuthApi { gateway_app_id: dev.gateway_app_id.clone(), }), AuthApi::Webhook { webhook } => { - let client = http_client()?; let path = if is_kms { "bootAuth/kms" } else { "bootAuth/app" }; let url = url_join(&webhook.url, path); - let response = client.post(&url).json(&boot_info).send().await?; - if !response.status().is_success() { - bail!("Failed to check boot auth: {}", response.text().await?); - } - Ok(response.json().await?) + http_post(&url, &boot_info).await } } } @@ -103,16 +113,7 @@ impl AuthApi { app_implementation: None, }), AuthApi::Webhook { webhook } => { - let client = http_client()?; - let response = client.get(&webhook.url).send().await?; - let status = response.status(); - let body = response.text().await?; - let info: AuthApiInfoResponse = serde_json::from_str(&body).with_context(|| { - format!( - "failed to decode auth api response from {}, status={status}, body={body}", - webhook.url - ) - })?; + let info: AuthApiInfoResponse = http_get(&webhook.url).await?; Ok(GetInfoResponse { is_dev: false, kms_contract_address: Some(info.kms_contract_addr.clone()),