diff --git a/crates/openshell-core/src/provider_credentials.rs b/crates/openshell-core/src/provider_credentials.rs index 65310b3a43..d0b7b38ad5 100644 --- a/crates/openshell-core/src/provider_credentials.rs +++ b/crates/openshell-core/src/provider_credentials.rs @@ -243,7 +243,7 @@ impl ProviderCredentialState { DEFAULT_EXPIRES_IN } else { let now = crate::time::now_ms(); - (expires_at_ms - now) / 1000 + expires_at_ms.saturating_sub(now) / 1000 } }, ); @@ -587,6 +587,27 @@ mod tests { ); } + #[test] + fn gcp_token_response_handles_already_expired_token_without_panic() { + let now_ms = i64::try_from( + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_millis(), + ) + .unwrap(); + let state = ProviderCredentialState::from_environment( + 1, + HashMap::from([("GCP_ADC_ACCESS_TOKEN".to_string(), "adc-tok".to_string())]), + HashMap::from([("GCP_ADC_ACCESS_TOKEN".to_string(), now_ms - 1_000)]), + HashMap::new(), + ); + assert!( + state.gcp_token_response().is_none(), + "expired token should be skipped rather than panic" + ); + } + #[test] fn child_env_with_gcp_resolved_resolves_vertex_vars_without_metadata_host() { let state = ProviderCredentialState::from_environment( diff --git a/crates/openshell-server/src/compute/lease.rs b/crates/openshell-server/src/compute/lease.rs index 1de4de4734..bf58fae48b 100644 --- a/crates/openshell-server/src/compute/lease.rs +++ b/crates/openshell-server/src/compute/lease.rs @@ -121,9 +121,9 @@ impl ReconcilerLease { pub async fn try_steal_expired(&self) -> Result { let record = self.read().await?.ok_or(LeaseError::NotFound)?; - let age_ms = now_ms() - record.updated_at_ms; + let current_ms = now_ms(); let ttl_ms = i64::try_from(self.ttl.as_millis()).unwrap_or(i64::MAX); - if age_ms < ttl_ms { + if !lease_is_expired(current_ms, record.updated_at_ms, ttl_ms) { return Err(LeaseError::AlreadyHeld); } @@ -252,6 +252,17 @@ pub fn replica_id() -> String { .unwrap_or_else(|_| uuid::Uuid::new_v4().to_string()) } +/// True if a lease record's age exceeds its TTL. +/// +/// Clamps at zero so a stored `updated_at_ms` in the future (clock skew) +/// is treated as age zero and the lease is considered fresh. Note that +/// `i64::saturating_sub` alone is not sufficient here: it saturates at +/// `i64::MIN`, not zero. +fn lease_is_expired(now_ms: i64, updated_at_ms: i64, ttl_ms: i64) -> bool { + let age_ms = now_ms.saturating_sub(updated_at_ms).max(0); + age_ms >= ttl_ms +} + #[cfg(test)] mod tests { use super::*; @@ -351,6 +362,28 @@ mod tests { assert_eq!(record.resource_version, guard.resource_version); } + #[test] + fn future_updated_at_is_treated_as_age_zero() { + let now = 1_000_000; + let future_updated = now + 86_400_000; // 1 day in the future + assert!( + !lease_is_expired(now, future_updated, 1_000), + "future updated_at_ms should not be stealable with positive TTL" + ); + assert!( + lease_is_expired(now, future_updated, 0), + "TTL zero always expires regardless of timestamp" + ); + assert!( + !lease_is_expired(now, now, 1_000), + "fresh lease with positive TTL should not be expired" + ); + assert!( + lease_is_expired(now, now, 0), + "present timestamp with TTL zero should be expired" + ); + } + #[tokio::test] async fn release_allows_immediate_reacquire() { let store = test_store().await; diff --git a/crates/openshell-server/src/inference.rs b/crates/openshell-server/src/inference.rs index 97217c2565..39d6afe2fd 100644 --- a/crates/openshell-server/src/inference.rs +++ b/crates/openshell-server/src/inference.rs @@ -627,11 +627,11 @@ fn resolve_vertex_ai_route( // protocol and path contract, but only for the OpenAI-compatible Vertex surface. // Anthropic-on-Vertex needs model-path shaping and body adaptation that a fully // caller-controlled URL cannot safely preserve. - if let Some(base_url) = config - .get(profile.base_url_config_keys[0]) - .or_else(|| config.get(profile.base_url_config_keys[1])) + if let Some(base_url) = profile + .base_url_config_keys + .iter() + .find_map(|key| config.get(*key).filter(|v| !v.trim().is_empty())) .map(String::as_str) - .filter(|v| !v.trim().is_empty()) { if is_anthropic { return Err(Status::invalid_argument( @@ -1133,6 +1133,65 @@ mod tests { } } + #[test] + fn resolve_vertex_ai_route_handles_empty_base_url_keys() { + let config = HashMap::new(); + let profile = openshell_core::inference::InferenceProviderProfile { + provider_type: "google_vertex_ai", + default_base_url: "https://example.com", + protocols: &[], + credential_key_names: &[], + base_url_config_keys: &[], + auth: openshell_core::inference::AuthHeader::Bearer, + default_headers: &[], + passthrough_headers: &[], + }; + let result = resolve_vertex_ai_route(&config, "model-id", "route", "api-key", &profile); + // Empty base_url_config_keys must not panic. The route still errors because + // the minimal Vertex config is missing, so we only assert reachability. + assert!( + result.is_err(), + "expected missing config to error, got {result:?}" + ); + } + + #[test] + fn resolve_vertex_ai_route_skips_blank_preferred_for_fallback() { + let mut config = HashMap::new(); + config.insert("GOOGLE_VERTEX_AI_BASE_URL".to_string(), " ".to_string()); + config.insert( + "VERTEX_AI_BASE_URL".to_string(), + "https://us-central1-aiplatform.googleapis.com".to_string(), + ); + config.insert( + VERTEX_AI_PROJECT_ID_KEY.to_string(), + "my-project".to_string(), + ); + config.insert(VERTEX_AI_REGION_KEY.to_string(), "us-central1".to_string()); + + let profile = openshell_core::inference::InferenceProviderProfile { + provider_type: "google-vertex-ai", + default_base_url: "", + protocols: &[], + credential_key_names: &[], + base_url_config_keys: &["GOOGLE_VERTEX_AI_BASE_URL", "VERTEX_AI_BASE_URL"], + auth: openshell_core::inference::AuthHeader::Bearer, + default_headers: &[], + passthrough_headers: &[], + }; + + let result = resolve_vertex_ai_route(&config, "model-id", "route", "api-key", &profile); + assert!( + result.is_ok(), + "blank preferred key should fall back to valid alias: {result:?}" + ); + let route = result.unwrap(); + assert_eq!( + route.endpoint, + "https://us-central1-aiplatform.googleapis.com" + ); + } + #[test] fn inference_bundle_requires_sandbox_principal() { let sandbox = test_sandbox_principal(); diff --git a/crates/openshell-server/src/provider_refresh.rs b/crates/openshell-server/src/provider_refresh.rs index f8fa867c1b..a51ec53373 100644 --- a/crates/openshell-server/src/provider_refresh.rs +++ b/crates/openshell-server/src/provider_refresh.rs @@ -769,6 +769,7 @@ async fn mint_aws_sts_assume_role( DEFAULT_MAX_LIFETIME_SECONDS }; let max_lifetime = i32::try_from(max_lifetime_i64.min(i64::from(i32::MAX))).unwrap_or(i32::MAX); + let max_lifetime_ms = i64::from(max_lifetime).saturating_mul(1000); let mut req = client .assume_role() @@ -794,11 +795,8 @@ async fn mint_aws_sts_assume_role( let session_token = creds.session_token().to_string(); let now_ms = current_time_ms(); - let expires_at_ms = creds - .expiration() - .to_millis() - .unwrap_or_else(|_| now_ms + max_lifetime_i64 * 1000); - let max_expires = now_ms + max_lifetime_i64 * 1000; + let max_expires = now_ms.saturating_add(max_lifetime_ms); + let expires_at_ms = creds.expiration().to_millis().unwrap_or(max_expires); let expires_at_ms = expires_at_ms.min(max_expires); // Map STS response fields to the env keys the profile bound to each semantic @@ -1519,6 +1517,16 @@ mod tests { ); } + #[test] + fn aws_sts_max_expires_saturates_on_saturated_clock() { + assert_eq!(i64::MAX.saturating_add(3_600_000), i64::MAX); + let now_ms = i64::MAX - 1_000; + let max_lifetime_ms = 3_600_000; + let max_expires = now_ms.saturating_add(max_lifetime_ms); + assert_eq!(max_expires, i64::MAX); + assert!(max_expires >= now_ms); + } + #[tokio::test] async fn aws_sts_assume_role_mints_three_credentials_from_mock_endpoint() { let mock_server = MockServer::start().await;