public/galaxy
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

first pass of merging in warp (doesn't build)

Browse Source
This commit is contained in:
Ryan Ward
2026-07-01 16:08:58 -05:00
parent 2f64909469
commit 4770ac06b5
3662 changed files with 414574 additions and 89772 deletions
Download Patch File Download Diff File Expand all files Collapse all files
app/src/tracing/cloud_agent_auth.rs
+546
View File
@@ -0,0 +1,546 @@
//! Provides authenticated OTLP trace transport and credential refresh for opted-in cloud agents.
//!
//! Dispatch bootstraps tracing with a bearer token and expiry in the process environment. The
//! exporter is built once around [`AuthenticatedHttpClient`], which reads a shared token snapshot
//! immediately before every request so refresh never requires rebuilding the exporter. Processes
//! without the endpoint switch or a currently valid dispatch credential never initialize this
//! module.
//!
//! Refresh begins only after the application has an authenticated managed-secrets client. A
//! successful mint replaces the dispatch credential only after the returned JWT's unverified
//! payload contains a string `run_id` exactly matching the immutable startup `OZ_RUN_ID`. This
//! payload inspection is only a rejection gate; the collector remains responsible for verifying
//! the token's signature, audience, expiry, and trusted trace resource attributes. Every refresh
//! failure preserves the last valid credential and enters bounded jittered backoff.
//!
//! Tokens must never appear in diagnostics or formatted values. Cached authorization headers are
//! marked sensitive, manual `Debug` implementations omit secrets, and token-store locks are always
//! released before network I/O.
use std::fmt;
use std::sync::{Arc, Mutex, RwLock};
use std::time::Duration;
use anyhow::{anyhow, Context as _};
use async_channel::{Receiver, Sender};
use async_compat::Compat;
use async_trait::async_trait;
use base64::Engine as _;
use chrono::{DateTime, Utc};
use futures_util::stream::AbortHandle;
use http::header::{HeaderValue, AUTHORIZATION};
use instant::Instant;
use opentelemetry_http::{Bytes, HttpClient, HttpError, Request, Response};
use warp_managed_secrets::client::{IdentityTokenOptions, ManagedSecretsClient, TaskIdentityToken};
use warpui::r#async::{FutureExt as _, Timer};
use warpui::{AppContext, Entity, ModelContext, SingletonEntity};
/// The environment variables form the immutable dispatch-time authentication bootstrap.
const CLOUD_AGENT_OTLP_TOKEN: &str = "WARP_CLOUD_AGENT_OTLP_TOKEN";
const CLOUD_AGENT_OTLP_TOKEN_EXPIRES_AT: &str = "WARP_CLOUD_AGENT_OTLP_TOKEN_EXPIRES_AT";
const OZ_RUN_ID: &str = "OZ_RUN_ID";
/// The collector audience and requested lifetime are fixed by the cloud-agent trace contract.
const COLLECTOR_AUDIENCE: &str = "warp-cloud-agent-otel";
const REFRESHED_TOKEN_DURATION: Duration = Duration::from_secs(60 * 60);
/// Proactive refresh starts roughly twenty minutes before expiry, with jitter to spread load.
const PROACTIVE_REFRESH_BUFFER: Duration = Duration::from_secs(20 * 60);
const PROACTIVE_REFRESH_JITTER: Duration = Duration::from_secs(2 * 60);
const MIN_PROACTIVE_REFRESH_DELAY: Duration = Duration::from_secs(1);
/// Failed refreshes use bounded full-jitter exponential backoff and rate-limited diagnostics.
const INITIAL_FAILURE_BACKOFF: Duration = Duration::from_secs(1);
const MAX_FAILURE_BACKOFF: Duration = Duration::from_secs(5 * 60);
const FAILURE_LOG_INTERVAL: Duration = Duration::from_secs(60);
/// A stalled identity-token request must release the single in-flight refresh slot.
const REFRESH_REQUEST_TIMEOUT: Duration = Duration::from_secs(30);
/// Shared dispatch authentication state between the exporter and the later refresh coordinator.
///
/// The optional expected run ID intentionally does not gate initial tracing: a valid dispatch
/// credential remains usable when `OZ_RUN_ID` is missing or empty, but every refreshed credential
/// is rejected until an immutable expected run ID is available to the replacement gate.
#[derive(Clone)]
pub(super) struct AuthContext {
token_store: TokenStore,
expected_run_id: Option<Arc<str>>,
refresh_hint_sender: Sender<()>,
refresh_hint_receiver: Arc<Mutex<Option<Receiver<()>>>>,
}
impl AuthContext {
/// Seeds authentication from a currently valid dispatch credential in the environment.
///
/// The caller treats failure as an opt-out so normal processes and partially rolled-out cloud
/// agents retain no-op tracing behavior.
pub(super) fn from_environment() -> anyhow::Result<Self> {
let token =
std::env::var(CLOUD_AGENT_OTLP_TOKEN).context("Cloud-agent OTLP token is missing")?;
// Remove the bootstrap secret as soon as it is owned so child processes cannot inherit it.
std::env::remove_var(CLOUD_AGENT_OTLP_TOKEN);
let token = token.trim().to_owned();
anyhow::ensure!(!token.is_empty(), "Cloud-agent OTLP token is empty");
let expires_at = std::env::var(CLOUD_AGENT_OTLP_TOKEN_EXPIRES_AT)
.context("Cloud-agent OTLP token expiry is missing")?;
let expires_at = DateTime::parse_from_rfc3339(expires_at.trim())
.context("Cloud-agent OTLP token expiry is not valid RFC3339")?;
anyhow::ensure!(
expires_at.offset().local_minus_utc() == 0,
"Cloud-agent OTLP token expiry is not UTC"
);
let expires_at = expires_at.with_timezone(&Utc);
anyhow::ensure!(
expires_at > Utc::now(),
"Cloud-agent OTLP token is already expired"
);
let expected_run_id = std::env::var(OZ_RUN_ID)
.ok()
.filter(|run_id| !run_id.trim().is_empty());
let token_store = TokenStore::new(token, expires_at)?;
let (refresh_hint_sender, refresh_hint_receiver) = async_channel::bounded(1);
Ok(Self {
token_store,
expected_run_id: expected_run_id.map(Into::into),
refresh_hint_sender,
refresh_hint_receiver: Arc::new(Mutex::new(Some(refresh_hint_receiver))),
})
}
/// Creates a transport sharing the latest credential while leaving the exporter itself stable.
pub(super) fn http_client(&self) -> AuthenticatedHttpClient {
AuthenticatedHttpClient {
inner: reqwest::Client::new(),
token_store: self.token_store.clone(),
refresh_hint_sender: self.refresh_hint_sender.clone(),
}
}
/// Transfers the bounded refresh-hint receiver to the one allowed coordinator.
fn take_refresh_hint_receiver(&self) -> Option<Receiver<()>> {
self.refresh_hint_receiver
.lock()
.unwrap_or_else(|err| err.into_inner())
.take()
}
}
impl fmt::Debug for AuthContext {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter
.debug_struct("AuthContext")
.field("token_store", &self.token_store)
.finish_non_exhaustive()
}
}
/// A snapshot of the latest credential, stored behind a short-lived reader/writer lock.
///
/// Readers clone only the sensitive authorization header, and no caller holds this lock during
/// network I/O. Replacement constructs and validates a complete snapshot before taking the write
/// lock so failures preserve the last valid credential.
#[derive(Clone)]
struct TokenStore {
inner: Arc<RwLock<TokenSnapshot>>,
}
impl TokenStore {
/// Creates the initial store from the validated dispatch credential.
fn new(token: String, expires_at: DateTime<Utc>) -> anyhow::Result<Self> {
Ok(Self {
inner: Arc::new(RwLock::new(TokenSnapshot::new(token, expires_at)?)),
})
}
/// Returns a cloned sensitive header only while the current credential remains unexpired.
fn valid_authorization_header(&self) -> Option<HeaderValue> {
let snapshot = self.inner.read().unwrap_or_else(|err| err.into_inner());
(snapshot.expires_at > Utc::now()).then(|| snapshot.authorization_header.clone())
}
/// Atomically replaces the current snapshot only with a usable unexpired credential.
fn replace(&self, token: String, expires_at: DateTime<Utc>) -> anyhow::Result<()> {
anyhow::ensure!(
expires_at > Utc::now(),
"Refreshed cloud-agent OTLP token is already expired"
);
let snapshot = TokenSnapshot::new(token, expires_at)?;
*self.inner.write().unwrap_or_else(|err| err.into_inner()) = snapshot;
Ok(())
}
/// Applies the exact-run rejection gate before allowing a refreshed credential to replace the
/// dispatch or previous refresh credential.
fn replace_refreshed(
&self,
token: String,
expires_at: DateTime<Utc>,
expected_run_id: Option<&str>,
) -> anyhow::Result<()> {
validate_refreshed_token_run_id(&token, expected_run_id)?;
self.replace(token, expires_at)
}
}
impl fmt::Debug for TokenStore {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
let snapshot = self.inner.read().unwrap_or_else(|err| err.into_inner());
formatter
.debug_struct("TokenStore")
.field("expires_at", &snapshot.expires_at)
.finish_non_exhaustive()
}
}
/// An already-parsed sensitive authorization header and its trusted server expiry.
struct TokenSnapshot {
authorization_header: HeaderValue,
expires_at: DateTime<Utc>,
}
impl TokenSnapshot {
/// Constructs a snapshot whose header redacts its value from standard debug formatting.
fn new(token: String, expires_at: DateTime<Utc>) -> anyhow::Result<Self> {
let mut authorization_header = HeaderValue::from_str(&format!("Bearer {token}"))
.map_err(|_| anyhow!("Cloud-agent OTLP token cannot be used as an HTTP header"))?;
authorization_header.set_sensitive(true);
Ok(Self {
authorization_header,
expires_at,
})
}
}
impl fmt::Debug for TokenSnapshot {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter
.debug_struct("TokenSnapshot")
.field("authorization_header", &"<redacted>")
.field("expires_at", &self.expires_at)
.finish()
}
}
/// Validates that the unverified refreshed-token payload names the immutable expected run exactly.
///
/// This local decode never establishes token authenticity. The collector remains responsible for
/// cryptographically verifying the token, while malformed or mismatched tokens fail closed here
/// before replacement and leave the existing credential untouched.
fn validate_refreshed_token_run_id(
token: &str,
expected_run_id: Option<&str>,
) -> anyhow::Result<()> {
let expected_run_id = expected_run_id
.filter(|run_id| !run_id.trim().is_empty())
.context("Expected cloud-agent run ID is missing or empty")?;
let mut segments = token.split('.');
let _header = segments
.next()
.filter(|segment| !segment.is_empty())
.context("Refreshed cloud-agent OTLP token is not a valid JWT")?;
let payload = segments
.next()
.filter(|segment| !segment.is_empty())
.context("Refreshed cloud-agent OTLP token is not a valid JWT")?;
let _signature = segments
.next()
.filter(|segment| !segment.is_empty())
.context("Refreshed cloud-agent OTLP token is not a valid JWT")?;
anyhow::ensure!(
segments.next().is_none(),
"Refreshed cloud-agent OTLP token is not a valid JWT"
);
let payload = base64::engine::general_purpose::URL_SAFE_NO_PAD
.decode(payload)
.map_err(|_| anyhow!("Refreshed cloud-agent OTLP token payload is not valid base64"))?;
let payload: serde_json::Value = serde_json::from_slice(&payload)
.map_err(|_| anyhow!("Refreshed cloud-agent OTLP token payload is not valid JSON"))?;
let run_id = payload
.get("run_id")
.and_then(serde_json::Value::as_str)
.context("Refreshed cloud-agent OTLP token has no string run ID")?;
anyhow::ensure!(
run_id == expected_run_id,
"Refreshed cloud-agent OTLP token run ID does not match"
);
Ok(())
}
/// The set of errors that can occur when making an HTTP request using [`AuthenticatedHttpClient`].
#[derive(thiserror::Error, Debug)]
enum AuthenticatedHttpError {
#[error("No unexpired cloud-agent OTLP token is available")]
NoValidToken,
#[error("Cloud-agent OTLP request failed with HTTP status {0}")]
HttpStatus(u16),
}
/// An HTTP client that injects the latest valid token immediately before each request.
///
/// The token-store lock is released before network I/O begins. A manual `Debug` implementation
/// prevents the client from formatting cached state, while sensitive [`HeaderValue`] instances
/// redact request headers. Expired credentials are removed and refused rather than sent.
pub(super) struct AuthenticatedHttpClient {
inner: reqwest::Client,
token_store: TokenStore,
refresh_hint_sender: Sender<()>,
}
impl fmt::Debug for AuthenticatedHttpClient {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter
.debug_struct("AuthenticatedHttpClient")
.field("token_store", &self.token_store)
.finish_non_exhaustive()
}
}
impl AuthenticatedHttpClient {
/// Overwrites any supplied authorization header with the latest unexpired credential.
///
/// Removing the supplied header first ensures an expired store fails closed rather than
/// accidentally sending a stale or caller-provided credential.
fn authorize_request(
&self,
request: &mut Request<Bytes>,
) -> Result<(), AuthenticatedHttpError> {
request.headers_mut().remove(AUTHORIZATION);
let authorization = self
.token_store
.valid_authorization_header()
.ok_or(AuthenticatedHttpError::NoValidToken)?;
request.headers_mut().insert(AUTHORIZATION, authorization);
Ok(())
}
}
#[async_trait]
impl HttpClient for AuthenticatedHttpClient {
async fn send_bytes(&self, mut request: Request<Bytes>) -> Result<Response<Bytes>, HttpError> {
self.authorize_request(&mut request)?;
let request: reqwest::Request = request.try_into()?;
// Reqwest requires a Tokio-compatible context, while the exporter may use another executor.
let (status, response) = Compat::new(async {
let mut response = self.inner.execute(request).await?;
let status = response.status();
let response = if status.is_success() {
let headers = std::mem::take(response.headers_mut());
Some((headers, response.bytes().await?))
} else {
None
};
Ok::<_, reqwest::Error>((status, response))
})
.await?;
if status == http::StatusCode::UNAUTHORIZED {
// The bounded nonblocking hint cannot recurse into or delay this export request.
let _ = self.refresh_hint_sender.try_send(());
}
let Some((headers, body)) = response else {
return Err(AuthenticatedHttpError::HttpStatus(status.as_u16()).into());
};
let mut response = Response::builder().status(status).body(body)?;
*response.headers_mut() = headers;
Ok(response)
}
}
/// Starts the one refresh coordinator after authenticated server connectivity is available.
///
/// Consuming the bounded hint receiver coalesces concurrent starts, and the coordinator immediately
/// mints once so the short-lived dispatch credential is replaced as soon as possible.
pub(super) fn start_refresh_coordinator(
auth_context: AuthContext,
client: Arc<dyn ManagedSecretsClient>,
ctx: &mut AppContext,
) {
let Some(refresh_hint_receiver) = auth_context.take_refresh_hint_receiver() else {
return;
};
ctx.add_singleton_model(move |ctx| {
AuthRefreshCoordinator::new(
auth_context.token_store,
auth_context.expected_run_id,
refresh_hint_receiver,
client,
ctx,
)
});
}
/// Owns serialized credential minting, proactive scheduling, failure backoff, and diagnostics.
///
/// At most one mint is in flight and one scheduled wakeup is retained. A bounded nonblocking 401
/// hint can accelerate refresh without recursing into or blocking the export request.
struct AuthRefreshCoordinator {
token_store: TokenStore,
expected_run_id: Option<Arc<str>>,
client: Arc<dyn ManagedSecretsClient>,
refresh_in_flight: bool,
consecutive_failures: u32,
scheduled_refresh: Option<AbortHandle>,
last_failure_diagnostic: Option<Instant>,
}
impl AuthRefreshCoordinator {
/// Installs the hint stream and immediately starts the first bounded refresh request.
fn new(
token_store: TokenStore,
expected_run_id: Option<Arc<str>>,
refresh_hint_receiver: Receiver<()>,
client: Arc<dyn ManagedSecretsClient>,
ctx: &mut ModelContext<Self>,
) -> Self {
let mut coordinator = Self {
token_store,
expected_run_id,
client,
refresh_in_flight: false,
consecutive_failures: 0,
scheduled_refresh: None,
last_failure_diagnostic: None,
};
let _ = ctx.spawn_stream_local(
refresh_hint_receiver,
|coordinator, (), ctx| coordinator.start_refresh(ctx),
|_, _| {},
);
coordinator.start_refresh(ctx);
coordinator
}
/// Starts one mint and coalesces all triggers while it remains in flight.
///
/// Each request asks for the fixed collector audience and principal-only subject, and the
/// timeout guarantees a stalled request eventually enters the ordinary failure path.
fn start_refresh(&mut self, ctx: &mut ModelContext<Self>) {
if self.refresh_in_flight {
return;
}
self.cancel_scheduled_refresh();
self.refresh_in_flight = true;
let client = self.client.clone();
ctx.spawn(
async move {
client
.issue_task_identity_token(IdentityTokenOptions {
audience: COLLECTOR_AUDIENCE.to_owned(),
requested_duration: REFRESHED_TOKEN_DURATION,
subject_template: vec1::vec1!["principal".to_owned()],
})
.with_timeout(REFRESH_REQUEST_TIMEOUT)
.await
.map_err(|_| anyhow!("Cloud-agent OTLP authorization refresh timed out"))?
},
|coordinator, result, ctx| coordinator.finish_refresh(result, ctx),
);
}
/// Accepts a refreshed credential only after all replacement gates succeed.
///
/// Any mint, timeout, expiry, header, or run-ID failure retains the last valid token and enters
/// the same bounded retry path without logging token contents.
fn finish_refresh(
&mut self,
result: anyhow::Result<TaskIdentityToken>,
ctx: &mut ModelContext<Self>,
) {
self.refresh_in_flight = false;
match result {
Ok(token) => {
let expires_at = token.expires_at;
if self
.token_store
.replace_refreshed(token.token, expires_at, self.expected_run_id.as_deref())
.is_ok()
{
self.consecutive_failures = 0;
log::info!("Cloud-agent OTLP authorization refreshed");
self.schedule_proactive_refresh(expires_at, ctx);
} else {
self.warn_refresh_failure();
self.schedule_failure_retry(ctx);
}
}
Err(_) => {
self.warn_refresh_failure();
self.schedule_failure_retry(ctx);
}
}
}
/// Schedules a refresh to occur before the current token expires.
///
/// This leaves some buffer for retries in case the refresh fails, but also guarantees
/// some minimum amount of time before the first refresh attempt.
fn schedule_proactive_refresh(
&mut self,
expires_at: DateTime<Utc>,
ctx: &mut ModelContext<Self>,
) {
let jitter = PROACTIVE_REFRESH_JITTER.mul_f64(rand::random::<f64>());
let refresh_buffer = PROACTIVE_REFRESH_BUFFER.saturating_add(jitter);
let remaining = (expires_at - Utc::now()).to_std().unwrap_or_default();
let delay = remaining
.saturating_sub(refresh_buffer)
.max(remaining.mul_f64(0.5))
.max(MIN_PROACTIVE_REFRESH_DELAY);
self.schedule_refresh(delay, ctx);
}
/// Schedules a full-jitter exponential retry capped at five minutes.
fn schedule_failure_retry(&mut self, ctx: &mut ModelContext<Self>) {
let exponent = self.consecutive_failures.min(31);
let upper_bound = INITIAL_FAILURE_BACKOFF
.saturating_mul(1u32 << exponent)
.min(MAX_FAILURE_BACKOFF);
self.consecutive_failures = self.consecutive_failures.saturating_add(1);
let delay = upper_bound.mul_f64(rand::random::<f64>());
self.schedule_refresh(delay, ctx);
}
/// Replaces the one scheduled wakeup so proactive, retry, and hint triggers stay coalesced.
fn schedule_refresh(&mut self, delay: Duration, ctx: &mut ModelContext<Self>) {
self.cancel_scheduled_refresh();
let task = ctx.spawn(
async move {
Timer::after(delay).await;
},
|coordinator, _, ctx| {
coordinator.scheduled_refresh = None;
coordinator.start_refresh(ctx);
},
);
self.scheduled_refresh = Some(task.abort_handle());
}
/// Cancels the prior wakeup without affecting a refresh already in flight.
fn cancel_scheduled_refresh(&mut self) {
if let Some(handle) = self.scheduled_refresh.take() {
handle.abort();
}
}
/// Emits a local token-free failure diagnostic at most once per configured interval.
fn warn_refresh_failure(&mut self) {
let now = Instant::now();
if self
.last_failure_diagnostic
.is_none_or(|last| now.duration_since(last) >= FAILURE_LOG_INTERVAL)
{
self.last_failure_diagnostic = Some(now);
log::warn!("Cloud-agent OTLP authorization refresh failed");
}
}
}
impl Entity for AuthRefreshCoordinator {
type Event = ();
}
impl SingletonEntity for AuthRefreshCoordinator {}
#[cfg(test)]
#[path = "cloud_agent_auth_tests.rs"]
mod tests;
app/src/tracing/cloud_agent_auth_tests.rs
+157
View File
@@ -0,0 +1,157 @@
use base64::Engine as _;
use chrono::TimeDelta;
use super::*;
fn jwt_with_payload(payload: serde_json::Value) -> String {
let encoder = base64::engine::general_purpose::URL_SAFE_NO_PAD;
let header = encoder.encode(br#"{"alg":"none"}"#);
let payload = encoder.encode(serde_json::to_vec(&payload).unwrap());
format!("{header}.{payload}.test-signature")
}
fn client_with_expiry(token: &str, expires_at: DateTime<Utc>) -> AuthenticatedHttpClient {
let (refresh_hint_sender, _) = async_channel::bounded(1);
AuthenticatedHttpClient {
inner: reqwest::Client::new(),
token_store: TokenStore::new(token.to_owned(), expires_at).unwrap(),
refresh_hint_sender,
}
}
#[test]
fn authorization_overwrites_supplied_header() {
let client = client_with_expiry(
"current-test-token",
Utc::now() + TimeDelta::try_minutes(5).unwrap(),
);
let mut request = Request::builder()
.header(AUTHORIZATION, "Bearer stale-test-token")
.body(Bytes::new())
.unwrap();
client.authorize_request(&mut request).unwrap();
assert_eq!(
request.headers().get(AUTHORIZATION).unwrap(),
"Bearer current-test-token"
);
}
#[test]
fn expired_token_is_refused_and_supplied_header_is_removed() {
let client = client_with_expiry(
"expired-test-token",
Utc::now() - TimeDelta::try_minutes(5).unwrap(),
);
let mut request = Request::builder()
.header(AUTHORIZATION, "Bearer stale-test-token")
.body(Bytes::new())
.unwrap();
assert!(matches!(
client.authorize_request(&mut request),
Err(AuthenticatedHttpError::NoValidToken)
));
assert!(!request.headers().contains_key(AUTHORIZATION));
}
#[test]
fn debug_output_redacts_token() {
let client = client_with_expiry(
"secret-test-token",
Utc::now() + TimeDelta::try_minutes(5).unwrap(),
);
let debug_output = format!("{client:?}");
assert!(!debug_output.contains("secret-test-token"));
assert!(debug_output.contains("expires_at"));
}
#[test]
fn authorized_request_debug_redacts_token() {
let client = client_with_expiry(
"secret-request-test-token",
Utc::now() + TimeDelta::try_minutes(5).unwrap(),
);
let mut request = Request::builder().body(Bytes::new()).unwrap();
client.authorize_request(&mut request).unwrap();
let request_debug = format!("{request:?}");
let headers_debug = format!("{:?}", request.headers());
assert!(!request_debug.contains("secret-request-test-token"));
assert!(!headers_debug.contains("secret-request-test-token"));
assert!(request_debug.contains("Sensitive"));
assert!(headers_debug.contains("Sensitive"));
}
#[test]
fn refreshed_token_run_id_exactly_matches() {
let token = jwt_with_payload(serde_json::json!({ "run_id": "expected-run-id" }));
validate_refreshed_token_run_id(&token, Some("expected-run-id")).unwrap();
}
#[test]
fn refreshed_token_run_id_is_required() {
let token = jwt_with_payload(serde_json::json!({}));
assert!(validate_refreshed_token_run_id(&token, Some("expected-run-id")).is_err());
}
#[test]
fn expected_run_id_is_required() {
let token = jwt_with_payload(serde_json::json!({ "run_id": "expected-run-id" }));
assert!(validate_refreshed_token_run_id(&token, None).is_err());
assert!(validate_refreshed_token_run_id(&token, Some("")).is_err());
}
#[test]
fn refreshed_token_run_id_must_match() {
let token = jwt_with_payload(serde_json::json!({ "run_id": "wrong-run-id" }));
assert!(validate_refreshed_token_run_id(&token, Some("expected-run-id")).is_err());
}
#[test]
fn refreshed_token_run_id_must_be_a_string() {
let token = jwt_with_payload(serde_json::json!({ "run_id": 123 }));
assert!(validate_refreshed_token_run_id(&token, Some("expected-run-id")).is_err());
}
#[test]
fn malformed_refreshed_tokens_are_rejected() {
let invalid_json = {
let encoder = base64::engine::general_purpose::URL_SAFE_NO_PAD;
let payload = encoder.encode(b"not-json");
format!("header.{payload}.signature")
};
for token in ["not-a-jwt", "header.!!!.signature", &invalid_json] {
assert!(validate_refreshed_token_run_id(token, Some("expected-run-id")).is_err());
}
}
#[test]
fn rejected_refreshed_token_preserves_previous_token() {
let token_store = TokenStore::new(
"current-test-token".to_owned(),
Utc::now() + TimeDelta::try_minutes(5).unwrap(),
)
.unwrap();
let wrong_run_token = jwt_with_payload(serde_json::json!({ "run_id": "wrong-run-id" }));
assert!(token_store
.replace_refreshed(
wrong_run_token,
Utc::now() + TimeDelta::try_minutes(5).unwrap(),
Some("expected-run-id"),
)
.is_err());
assert_eq!(
token_store.valid_authorization_header().unwrap(),
"Bearer current-test-token"
);
}
app/src/tracing/native.rs
+542
View File
@@ -0,0 +1,542 @@
//! Configures opt-in OpenTelemetry export for cloud-agent traces on native platforms.
//!
//! The global `tracing` subscriber observes the whole application, while
//! [`CloudAgentSpanExporter`] limits OTLP export to spans explicitly marked with
//! [`CLOUD_AGENT_MARKER`]. Keeping this selection at the exporter boundary lets callers use the
//! normal `tracing` macros and propagation machinery without installing a second subscriber or
//! coupling generic task executors to cloud-agent tracing.
//!
//! # Why spans must be ended during shutdown
//!
//! An OpenTelemetry span becomes exportable and passes owned span data to a processor's `on_end`
//! callback only after it ends. Shutting down an [`SdkTracerProvider`] flushes spans that have
//! reached `on_end`, but it does not end spans that are still active. Some `tracing::Span`
//! references are intentionally propagated into asynchronous task machinery and can therefore
//! remain alive when the application terminates. Shutting down the provider before those spans end
//! would silently discard them.
//!
//! [`ShutdownAwareTracer`] and [`ShutdownAwareSpan`] wrap the SDK tracer and spans used by
//! `tracing-opentelemetry`. This keeps existing `tracing` instrumentation unchanged while allowing
//! [`ActiveSpanRegistry`] to explicitly end still-reachable, registered SDK spans before shutting
//! down the provider. The standard application lifecycle retains [`Initialization`] in its
//! termination callback so this ordering happens before platforms that terminate the process
//! without running Rust destructors. [`Initialization`]'s `Drop` implementation remains a fallback
//! for ordinary returns. Explicit process exits bypass both forms of cleanup.
//!
//! # Span ownership and synchronization
//!
//! `tracing-opentelemetry` creates SDK spans lazily during several `tracing` span lifecycle
//! operations, including when it needs a span's context and when a span closes. Every SDK span that
//! reaches [`ShutdownAwareTracer::build_with_context`] is wrapped in an `Arc<Mutex<_>>`. Before
//! shutdown begins, it is weakly registered; after shutdown begins, it is immediately ended
//! instead. The wrapper remains the span's owner; the registry uses weak references so tracking
//! does not extend normal span lifetimes. An SDK span that has not yet been built when shutdown
//! begins cannot reach `on_end` before provider shutdown. If it materializes later, it is ended too
//! late for export.
//!
//! SDK-span creation and shutdown are serialized by the registry-state mutex. Shutdown keeps that
//! mutex locked while it ends every still-upgradeable registered span and shuts down the provider,
//! preventing an SDK span from being created in the otherwise-dangerous gap between those
//! operations. A final span owner can begin dropping after its weak reference becomes impossible
//! to upgrade, so shutdown cannot strictly guarantee that every previously registered span has
//! finished ending. The lock order is always registry state followed by an individual SDK span.
//! Normal span operations lock only the individual SDK span and never attempt to lock the registry.
//! Mutex acquisition recovers poisoned inner values because trace export and shutdown are
//! best-effort cleanup that should continue after an unrelated panic.
use std::borrow::Cow;
use std::sync::{Arc, Mutex, OnceLock, Weak};
use std::time::{Duration, SystemTime};
use anyhow::{anyhow, Context as _};
use instant::Instant;
use opentelemetry::trace::{
Span as _, SpanBuilder, SpanContext, Status, Tracer as _, TracerProvider as _,
};
use opentelemetry::{Context as OtelContext, KeyValue, Value};
use opentelemetry_otlp::{Protocol, WithExportConfig as _, WithHttpConfig as _};
use opentelemetry_sdk::error::OTelSdkResult;
use opentelemetry_sdk::resource::{EnvResourceDetector, TelemetryResourceDetector};
use opentelemetry_sdk::trace::{
SdkTracer, SdkTracerProvider, Span as SdkSpan, SpanData, SpanExporter,
};
use opentelemetry_sdk::Resource;
use tracing::subscriber;
use tracing_subscriber::layer::SubscriberExt as _;
use tracing_subscriber::EnvFilter;
use url::{Host, Url};
use warp_managed_secrets::client::ManagedSecretsClient;
use warpui::AppContext;
use super::cloud_agent_auth::{self, AuthContext};
use super::Initialization;
use crate::channel::ChannelState;
use crate::tracing::install_no_subscriber;
/// The tag used to mark spans related to cloud agents, which we use to filter out
/// spans we don't care about (e.g.: ones from dependencies).
const CLOUD_AGENT_MARKER: &str = "tags.cloud_agent";
/// The environment variable used to configure the cloud agent OTLP endpoint.
const CLOUD_AGENT_OTLP_ENDPOINT: &str = "WARP_CLOUD_AGENT_OTLP_ENDPOINT";
/// The environment variable used to configure the OTel service name.
const OTEL_SERVICE_NAME: &str = "OTEL_SERVICE_NAME";
/// The minimum interval between local export failure diagnostics.
const EXPORT_FAILURE_LOG_INTERVAL: Duration = Duration::from_secs(60);
/// Process-global authentication context for cloud-agent OTLP export.
///
/// The exporter is built once during [`init`], while the stored context later starts dynamic
/// credential refresh after authenticated application services become available, and this static
/// remains unset for processes that did not opt in.
static AUTH_CONTEXT: OnceLock<AuthContext> = OnceLock::new();
/// Installs the native tracing subscriber and optional cloud-agent OTLP exporter.
///
/// Export is deliberately opt-in through [`CLOUD_AGENT_OTLP_ENDPOINT`] plus a currently valid
/// dispatch token. When either is absent or the exporter cannot be constructed, a no-op subscriber
/// is installed so tracing instrumentation remains safe without producing output or partially
/// initializing export.
pub fn init() -> anyhow::Result<Initialization> {
// INFO is the default because this is a global subscriber and DEBUG-level application spans
// would otherwise create substantial work even though only marked cloud-agent spans are
// exported. RUST_LOG can still override this when deeper tracing is needed.
let env_filter = EnvFilter::builder()
.with_default_directive(tracing::Level::INFO.into())
.from_env_lossy();
let Some(base_endpoint) = std::env::var(CLOUD_AGENT_OTLP_ENDPOINT)
.ok()
.filter(|endpoint| !endpoint.trim().is_empty())
else {
install_no_subscriber()?;
return Ok(Initialization::default());
};
let Ok(auth_context) = AuthContext::from_environment() else {
install_no_subscriber()?;
return Ok(Initialization::default());
};
let shutdown_timeout = export_timeout();
let provider = match build_provider(base_endpoint.trim(), &auth_context) {
Ok(provider) => provider,
Err(err) => {
install_no_subscriber()?;
return Ok(Initialization {
initialization_warning: Some(err),
active_spans: None,
provider: None,
shutdown_timeout,
});
}
};
let _ = AUTH_CONTEXT.set(auth_context);
let active_spans = ActiveSpanRegistry::default();
let tracer =
ShutdownAwareTracer::new(provider.tracer("warp-cloud-agent"), active_spans.clone());
let subscriber = tracing_subscriber::registry()
.with(env_filter)
.with(tracing_opentelemetry::layer().with_tracer(tracer));
subscriber::set_global_default(subscriber)?;
Ok(Initialization {
initialization_warning: None,
active_spans: Some(active_spans),
provider: Some(provider),
shutdown_timeout,
})
}
/// Builds the SDK provider and its batch exporter.
///
/// A batch exporter keeps network export off instrumentation call sites. The provider is retained
/// by [`Initialization`] so application termination can explicitly shut it down after attempting
/// to end registered active spans. Exported resources include Warp's version and channel alongside
/// standard environment-detected OpenTelemetry attributes, with [`OTEL_SERVICE_NAME`] taking
/// precedence over the default service name. The exporter is built once with a dynamic HTTP client
/// so credential refresh can update requests without reconstructing provider state.
fn build_provider(
base_endpoint: &str,
auth_context: &AuthContext,
) -> anyhow::Result<SdkTracerProvider> {
let endpoint = traces_endpoint(base_endpoint)?;
let exporter = opentelemetry_otlp::SpanExporter::builder()
.with_http()
.with_http_client(auth_context.http_client())
.with_protocol(Protocol::HttpBinary)
.with_endpoint(endpoint)
.build()
.context("Failed to build the OTLP span exporter")?;
let resource = Resource::builder_empty()
.with_service_name("warp-cloud-agent")
.with_attribute(KeyValue::new(
"service.version",
ChannelState::app_version().unwrap_or("<no tag>"),
))
.with_attribute(KeyValue::new(
"warp.channel",
ChannelState::channel().to_string(),
))
.with_detector(Box::new(TelemetryResourceDetector))
.with_detector(Box::new(EnvResourceDetector::new()));
let resource = match std::env::var(OTEL_SERVICE_NAME) {
Ok(service_name) if !service_name.is_empty() => resource.with_service_name(service_name),
Ok(_) | Err(_) => resource,
}
.build();
Ok(SdkTracerProvider::builder()
.with_batch_exporter(CloudAgentSpanExporter {
inner: exporter,
diagnostics: RateLimitedDiagnostics::default(),
})
.with_resource(resource)
.build())
}
/// Starts the single refresh coordinator after the authenticated server client exists.
///
/// Processes that did not opt in with both an endpoint and valid dispatch credential have no
/// retained [`AUTH_CONTEXT`] and remain no-ops here.
pub(super) fn start_auth_refresh(client: Arc<dyn ManagedSecretsClient>, ctx: &mut AppContext) {
if let Some(auth_context) = AUTH_CONTEXT.get() {
cloud_agent_auth::start_refresh_coordinator(auth_context.clone(), client, ctx);
}
}
/// Converts the configured OTLP base URL into the HTTP/protobuf traces endpoint.
///
/// The configuration is treated as a base URL rather than a complete signal-specific URL, so any
/// query or fragment is discarded before appending `v1/traces`. Authenticated export requires
/// HTTPS unless the configured host is guaranteed to resolve to the local machine.
fn traces_endpoint(base_endpoint: &str) -> anyhow::Result<String> {
let mut endpoint = Url::parse(base_endpoint).context("Invalid cloud-agent OTLP endpoint")?;
match endpoint.scheme() {
"https" => {}
"http" if endpoint_host_is_loopback(&endpoint) => {}
"http" => {
return Err(anyhow!(
"Cloud-agent OTLP endpoint must use HTTPS unless its host is loopback"
));
}
_ => return Err(anyhow!("Cloud-agent OTLP endpoint must use HTTP or HTTPS")),
}
endpoint.set_query(None);
endpoint.set_fragment(None);
endpoint
.path_segments_mut()
.map_err(|_| anyhow!("Cloud-agent OTLP endpoint cannot be used as a base URL"))?
.pop_if_empty()
.extend(["v1", "traces"]);
Ok(endpoint.into())
}
/// Returns whether the endpoint host is guaranteed to resolve to the local machine.
fn endpoint_host_is_loopback(endpoint: &Url) -> bool {
match endpoint.host() {
Some(Host::Domain(domain)) => domain.eq_ignore_ascii_case("localhost"),
Some(Host::Ipv4(address)) => address.is_loopback(),
Some(Host::Ipv6(address)) => address.is_loopback(),
None => false,
}
}
/// Returns the export shutdown timeout using the standard OpenTelemetry environment variables.
fn export_timeout() -> Duration {
[
"OTEL_EXPORTER_OTLP_TRACES_TIMEOUT",
"OTEL_EXPORTER_OTLP_TIMEOUT",
]
.into_iter()
.find_map(|name| {
std::env::var(name)
.ok()
.and_then(|value| value.parse::<u64>().ok())
.map(Duration::from_millis)
})
.unwrap_or(super::DEFAULT_EXPORT_TIMEOUT)
}
/// A registry of started SDK spans used for best-effort ending before provider shutdown.
///
/// This registry belongs beside the provider in [`Initialization`]. It stores only weak references
/// so a span that ends normally can be dropped without first unregistering itself.
#[derive(Clone, Debug, Default)]
pub(super) struct ActiveSpanRegistry {
state: Arc<Mutex<ActiveSpanRegistryState>>,
}
#[derive(Debug, Default)]
struct ActiveSpanRegistryState {
/// Prevents new spans from remaining active after shutdown begins.
shutting_down: bool,
/// Weak references avoid extending the lifetime of spans that end normally.
spans: Vec<Weak<Mutex<SdkSpan>>>,
}
impl ActiveSpanRegistry {
/// Builds an SDK span, registering it before shutdown or ending it after shutdown begins.
///
/// `tracing-opentelemetry` calls this whenever it materializes an SDK span. If shutdown has
/// already begun, the newly built span is ended immediately rather than being allowed to
/// remain active.
fn build_span(
&self,
tracer: &SdkTracer,
builder: SpanBuilder,
parent_cx: &OtelContext,
) -> ShutdownAwareSpan {
let mut state = self.state.lock().unwrap_or_else(|err| err.into_inner());
let span = tracer.build_with_context(builder, parent_cx);
let span_context = span.span_context().clone();
let span = Arc::new(Mutex::new(span));
if state.shutting_down {
span.lock().unwrap_or_else(|err| err.into_inner()).end();
} else {
// Dead weak references are pruned opportunistically to avoid requiring normal span
// completion to acquire the registry lock.
state.spans.retain(|span| span.strong_count() > 0);
state.spans.push(Arc::downgrade(&span));
}
ShutdownAwareSpan {
span_context,
inner: span,
}
}
/// Ends every still-upgradeable registered span and then shuts down the provider.
///
/// The registry lock intentionally remains held through provider shutdown. This guarantees
/// that no SDK span can be built between the final end attempt and the provider becoming unable
/// to accept ended spans. It does not synchronize with the final drop of a span whose weak
/// reference can no longer be upgraded, so ending previously built spans remains best-effort.
pub(super) fn shutdown(
&self,
provider: &SdkTracerProvider,
timeout: Duration,
) -> OTelSdkResult {
let mut state = self.state.lock().unwrap_or_else(|err| err.into_inner());
state.shutting_down = true;
let spans = std::mem::take(&mut state.spans);
for span in spans {
if let Some(span) = span.upgrade() {
span.lock().unwrap_or_else(|err| err.into_inner()).end();
}
}
let result = provider.shutdown_with_timeout(timeout);
drop(state);
result
}
}
/// An [`SdkTracer`] adapter that routes spans through shutdown-aware construction.
///
/// Wrapping the tracer, rather than the span processor, is necessary because processors receive
/// only a temporary mutable reference in `on_start` and receive owned exportable data only after
/// `on_end`. A processor therefore cannot retain handles to, or end, active spans during shutdown.
#[derive(Clone, Debug)]
struct ShutdownAwareTracer {
inner: SdkTracer,
active_spans: ActiveSpanRegistry,
}
impl ShutdownAwareTracer {
fn new(inner: SdkTracer, active_spans: ActiveSpanRegistry) -> Self {
Self {
inner,
active_spans,
}
}
}
impl opentelemetry::trace::Tracer for ShutdownAwareTracer {
type Span = ShutdownAwareSpan;
fn build_with_context(&self, builder: SpanBuilder, parent_cx: &OtelContext) -> Self::Span {
self.active_spans
.build_span(&self.inner, builder, parent_cx)
}
}
/// A synchronized wrapper around an SDK span shared with [`ActiveSpanRegistry`].
///
/// The immutable [`SpanContext`] is cached outside the mutex because the OpenTelemetry
/// [`opentelemetry::trace::Span`] trait must return it by reference. All mutable SDK-span operations
/// are forwarded through the mutex, allowing shutdown to end the same underlying span. Repeated
/// end calls are harmless because SDK spans export only once.
#[derive(Debug)]
struct ShutdownAwareSpan {
span_context: SpanContext,
inner: Arc<Mutex<SdkSpan>>,
}
impl opentelemetry::trace::Span for ShutdownAwareSpan {
fn add_event_with_timestamp<T>(
&mut self,
name: T,
timestamp: SystemTime,
attributes: Vec<KeyValue>,
) where
T: Into<Cow<'static, str>>,
{
self.inner
.lock()
.unwrap_or_else(|err| err.into_inner())
.add_event_with_timestamp(name, timestamp, attributes);
}
fn span_context(&self) -> &SpanContext {
&self.span_context
}
fn is_recording(&self) -> bool {
self.inner
.lock()
.unwrap_or_else(|err| err.into_inner())
.is_recording()
}
fn set_attribute(&mut self, attribute: KeyValue) {
self.inner
.lock()
.unwrap_or_else(|err| err.into_inner())
.set_attribute(attribute);
}
fn set_status(&mut self, status: Status) {
self.inner
.lock()
.unwrap_or_else(|err| err.into_inner())
.set_status(status);
}
fn update_name<T>(&mut self, new_name: T)
where
T: Into<Cow<'static, str>>,
{
self.inner
.lock()
.unwrap_or_else(|err| err.into_inner())
.update_name(new_name);
}
fn add_link(&mut self, span_context: SpanContext, attributes: Vec<KeyValue>) {
self.inner
.lock()
.unwrap_or_else(|err| err.into_inner())
.add_link(span_context, attributes);
}
fn end_with_timestamp(&mut self, timestamp: SystemTime) {
self.inner
.lock()
.unwrap_or_else(|err| err.into_inner())
.end_with_timestamp(timestamp);
}
}
/// An exporter that restricts the shared tracing subscriber's output to explicitly marked
/// cloud-agent spans.
///
/// Filtering here preserves normal parent/context propagation inside the application while
/// ensuring unrelated application tracing is never sent to the configured cloud-agent endpoint.
/// The marker is a per-span routing attribute rather than an inherited property, so every span
/// intended for export must set it explicitly.
struct CloudAgentSpanExporter {
inner: opentelemetry_otlp::SpanExporter,
diagnostics: RateLimitedDiagnostics,
}
impl std::fmt::Debug for CloudAgentSpanExporter {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
formatter
.debug_struct("CloudAgentSpanExporter")
.finish_non_exhaustive()
}
}
impl SpanExporter for CloudAgentSpanExporter {
fn export(
&self,
batch: Vec<SpanData>,
) -> impl std::future::Future<Output = OTelSdkResult> + Send {
let batch: Vec<_> = batch
.into_iter()
.filter_map(filter_cloud_agent_span)
.collect();
async move {
if batch.is_empty() {
return Ok(());
}
let result = self.inner.export(batch).await;
if result.is_err() {
self.diagnostics.warn_export_failure();
}
result
}
}
fn shutdown_with_timeout(&self, timeout: Duration) -> OTelSdkResult {
let result = self.inner.shutdown_with_timeout(timeout);
if let Err(err) = &result {
log::warn!("Failed to shut down the cloud-agent OpenTelemetry span exporter: {err}");
}
result
}
fn force_flush(&self) -> OTelSdkResult {
let result = self.inner.force_flush();
if let Err(err) = &result {
log::warn!("Failed to flush the cloud-agent OpenTelemetry span exporter: {err}");
}
result
}
fn set_resource(&mut self, resource: &Resource) {
self.inner.set_resource(resource);
}
}
/// Rate-limits local token-free export diagnostics independently of exporter retries.
#[derive(Debug, Default)]
struct RateLimitedDiagnostics {
last_export_failure: Mutex<Option<Instant>>,
}
impl RateLimitedDiagnostics {
/// Emits at most one local export-failure warning per configured interval.
fn warn_export_failure(&self) {
let now = Instant::now();
let mut last_failure = self
.last_export_failure
.lock()
.unwrap_or_else(|err| err.into_inner());
if last_failure.is_none_or(|last| now.duration_since(last) >= EXPORT_FAILURE_LOG_INTERVAL) {
*last_failure = Some(now);
log::warn!("Failed to export cloud-agent OpenTelemetry spans");
}
}
}
/// Removes unrelated spans and strips the internal routing marker from spans and events before
/// export.
fn filter_cloud_agent_span(mut span: SpanData) -> Option<SpanData> {
let is_cloud_agent_span = span.attributes.iter().any(|attribute| {
attribute.key.as_str() == CLOUD_AGENT_MARKER && attribute.value == Value::Bool(true)
});
if !is_cloud_agent_span {
return None;
}
span.attributes
.retain(|attribute| attribute.key.as_str() != CLOUD_AGENT_MARKER);
for event in &mut span.events.events {
event
.attributes
.retain(|attribute| attribute.key.as_str() != CLOUD_AGENT_MARKER);
}
Some(span)
}