diff --git a/crates/openshell-driver-vault/src/lib.rs b/crates/openshell-driver-vault/src/lib.rs index b21524ae2..5ba60d50c 100644 --- a/crates/openshell-driver-vault/src/lib.rs +++ b/crates/openshell-driver-vault/src/lib.rs @@ -131,13 +131,18 @@ impl VaultCredentialDriver { &self, request: StoreCredentialRequest, ) -> Result { - let token = self.auth_token().await?; let logical_path = if let Some(existing_handle) = request.existing_handle.as_ref() { Self::logical_path_from_handle(existing_handle)? } else { managed_secret_path(&request.provider_name, &request.credential_key) }; validate_secret_path(&logical_path).map_err(Status::invalid_argument)?; + validate_managed_secret_path( + &request.provider_name, + &request.credential_key, + &logical_path, + )?; + let token = self.auth_token().await?; let reference = VaultSecretReference { api_path: api_path_for_reference( &self.settings.mount, @@ -157,9 +162,14 @@ impl VaultCredentialDriver { } pub async fn delete_credential(&self, request: DeleteCredentialRequest) -> Result<(), Status> { - let token = self.auth_token().await?; let handle = Self::handle_from_request("delete", request.handle)?; let logical_path = Self::logical_path_from_handle(&handle)?; + validate_managed_secret_path( + &request.provider_name, + &request.credential_key, + &logical_path, + )?; + let token = self.auth_token().await?; let api_path = delete_api_path_for_reference( &self.settings.mount, self.settings.kv_version, @@ -172,11 +182,16 @@ impl VaultCredentialDriver { &self, requests: Vec, ) -> Result, Status> { - let token = self.auth_token().await?; let mut responses = Vec::with_capacity(requests.len()); + let mut resolved_requests = Vec::with_capacity(requests.len()); for request in requests { let handle = Self::handle_from_request(&request.request_id, request.handle)?; let logical_path = Self::logical_path_from_handle(&handle)?; + validate_managed_secret_path( + &request.provider_name, + &request.credential_key, + &logical_path, + )?; let reference = VaultSecretReference { api_path: api_path_for_reference( &self.settings.mount, @@ -186,9 +201,14 @@ impl VaultCredentialDriver { key: STORED_VALUE_KEY.to_string(), kv_version: self.settings.kv_version, }; + resolved_requests.push((request.request_id, reference)); + } + + let token = self.auth_token().await?; + for (request_id, reference) in resolved_requests { let value = self.resolve_secret_value(&reference, &token).await?; responses.push(ResolvedCredential { - request_id: request.request_id, + request_id, value, expires_at_ms: 0, }); @@ -690,6 +710,20 @@ fn managed_secret_path(provider_name: &str, credential_key: &str) -> String { format!("openshell/provider-credentials/{}", &hex[..40]) } +fn validate_managed_secret_path( + provider_name: &str, + credential_key: &str, + logical_path: &str, +) -> Result<(), Status> { + let expected = managed_secret_path(provider_name, credential_key); + if logical_path == expected { + return Ok(()); + } + Err(Status::invalid_argument(format!( + "vault credential handle path does not match the managed path for provider credential '{credential_key}'" + ))) +} + async fn read_secret_file(path: &Path, description: &str) -> Result { let contents = tokio::fs::read_to_string(path).await.map_err(|err| { Status::unauthenticated(format!( @@ -920,6 +954,19 @@ mod tests { assert!(err.message().contains("path segments")); } + #[test] + fn handle_rejects_unexpected_managed_path() { + let err = validate_managed_secret_path( + "nvidia-prod", + "NVIDIA_API_KEY", + "openshell/provider-credentials/other", + ) + .unwrap_err(); + + assert_eq!(err.code(), Code::InvalidArgument); + assert!(err.message().contains("managed path")); + } + #[tokio::test] async fn store_and_resolve_token_file_kv2_secret() { let mock_server = MockServer::start().await; @@ -985,10 +1032,9 @@ mod tests { #[tokio::test] async fn store_with_existing_handle_reuses_logical_path() { let mock_server = MockServer::start().await; + let logical_path = managed_secret_path("nvidia-prod", "NVIDIA_API_KEY"); Mock::given(method("POST")) - .and(path( - "/v1/secret/data/openshell/provider-credentials/existing", - )) + .and(path(format!("/v1/secret/data/{logical_path}"))) .and(header("x-vault-token", "dev-token")) .and(body_string_contains("updated-secret")) .respond_with(ResponseTemplate::new(200)) @@ -1010,21 +1056,20 @@ mod tests { provider_name: "nvidia-prod".to_string(), credential_key: "NVIDIA_API_KEY".to_string(), value: "updated-secret".to_string(), - existing_handle: Some(handle("v1:openshell/provider-credentials/existing")), + existing_handle: Some(handle(&format!("v1:{logical_path}"))), }) .await .unwrap(); - assert_eq!(stored.handle, "v1:openshell/provider-credentials/existing"); + assert_eq!(stored.handle, format!("v1:{logical_path}")); } #[tokio::test] async fn delete_token_file_kv2_secret() { let mock_server = MockServer::start().await; + let logical_path = managed_secret_path("nvidia-prod", "NVIDIA_API_KEY"); Mock::given(method("DELETE")) - .and(path( - "/v1/secret/metadata/openshell/provider-credentials/existing", - )) + .and(path(format!("/v1/secret/metadata/{logical_path}"))) .and(header("x-vault-token", "dev-token")) .respond_with(ResponseTemplate::new(204)) .mount(&mock_server) @@ -1044,7 +1089,7 @@ mod tests { .delete_credential(DeleteCredentialRequest { provider_name: "nvidia-prod".to_string(), credential_key: "NVIDIA_API_KEY".to_string(), - handle: Some(handle("v1:openshell/provider-credentials/existing")), + handle: Some(handle(&format!("v1:{logical_path}"))), }) .await .unwrap(); @@ -1105,10 +1150,9 @@ mod tests { #[tokio::test] async fn resolve_maps_missing_key() { let mock_server = MockServer::start().await; + let logical_path = managed_secret_path("nvidia-prod", "NVIDIA_API_KEY"); Mock::given(method("GET")) - .and(path( - "/v1/secret/data/openshell/provider-credentials/missing-key", - )) + .and(path(format!("/v1/secret/data/{logical_path}"))) .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ "data": { "data": {} @@ -1132,7 +1176,7 @@ mod tests { request_id: "credential-0".to_string(), provider_name: "nvidia-prod".to_string(), credential_key: "NVIDIA_API_KEY".to_string(), - handle: Some(handle("v1:openshell/provider-credentials/missing-key")), + handle: Some(handle(&format!("v1:{logical_path}"))), }]) .await .unwrap_err(); @@ -1144,10 +1188,9 @@ mod tests { #[tokio::test] async fn resolve_maps_permission_denied() { let mock_server = MockServer::start().await; + let logical_path = managed_secret_path("nvidia-prod", "NVIDIA_API_KEY"); Mock::given(method("GET")) - .and(path( - "/v1/secret/data/openshell/provider-credentials/denied", - )) + .and(path(format!("/v1/secret/data/{logical_path}"))) .respond_with(ResponseTemplate::new(403)) .mount(&mock_server) .await; @@ -1167,7 +1210,7 @@ mod tests { request_id: "credential-0".to_string(), provider_name: "nvidia-prod".to_string(), credential_key: "NVIDIA_API_KEY".to_string(), - handle: Some(handle("v1:openshell/provider-credentials/denied")), + handle: Some(handle(&format!("v1:{logical_path}"))), }]) .await .unwrap_err(); diff --git a/crates/openshell-server/src/credentials.rs b/crates/openshell-server/src/credentials.rs index 596410ffa..ccb3f5e19 100644 --- a/crates/openshell-server/src/credentials.rs +++ b/crates/openshell-server/src/credentials.rs @@ -193,6 +193,10 @@ impl CredentialRuntime { self.drivers.contains_key(&driver_name) } + pub fn storage_owns_handle(&self, handle: &CredentialHandle) -> bool { + normalize_driver_name(&handle.driver) == self.registry.storage_owner_name() + } + pub async fn store_provider_credentials( &self, provider_name: &str, diff --git a/crates/openshell-server/src/grpc/provider.rs b/crates/openshell-server/src/grpc/provider.rs index 13ed1622a..5535cd658 100644 --- a/crates/openshell-server/src/grpc/provider.rs +++ b/crates/openshell-server/src/grpc/provider.rs @@ -287,9 +287,6 @@ async fn update_provider_record_validating( candidate.credential_expires_at_ms, provider.credential_expires_at_ms, ); - for key in updated_credential_values.keys() { - candidate.credential_handles.remove(key); - } // Validate BEFORE writing to prevent persisting invalid state. // Validate only the mutable fields (credentials/config) plus metadata and @@ -298,25 +295,42 @@ async fn update_provider_record_validating( // strand legacy records whose stored type predates current limits. See // #1347. super::validation::validate_object_metadata(candidate.metadata.as_ref(), "provider")?; - validate_provider_mutable_fields(&candidate)?; - delete_removed_provider_credentials( + let credential_update = prepare_provider_credential_update( credentials, candidate.object_name(), - &removed_credential_handles, - ) - .await?; - for key in removed_credential_handles.keys() { - candidate.credential_handles.remove(key); - } - store_provider_credentials_if_configured( - credentials, - &mut candidate, &updated_credential_values, &existing_handles, ) .await?; - validate_provider_mutable_fields(&candidate)?; - validate_provider_update_against_attached_sandboxes(store, &candidate).await?; + if credentials.is_some_and(crate::credentials::CredentialRuntime::stores_provider_credentials) { + for key in updated_credential_values.keys() { + candidate.credentials.remove(key); + } + } + for key in removed_credential_handles.keys() { + candidate.credential_handles.remove(key); + } + candidate + .credential_handles + .extend(credential_update.pre_stored_handles.clone()); + if let Err(err) = validate_provider_mutable_fields(&candidate) { + cleanup_pre_stored_provider_credentials( + credentials, + candidate.object_name(), + &credential_update.pre_stored_handles, + ) + .await; + return Err(err); + } + if let Err(err) = validate_provider_update_against_attached_sandboxes(store, &candidate).await { + cleanup_pre_stored_provider_credentials( + credentials, + candidate.object_name(), + &credential_update.pre_stored_handles, + ) + .await; + return Err(err); + } // Serialize labels for storage let labels_map = candidate.object_labels(); @@ -326,14 +340,22 @@ async fn update_provider_record_validating( { None } else { - Some( - serde_json::to_string(&labels_map) - .map_err(|e| Status::internal(format!("serialize labels failed: {e}")))?, - ) + match serde_json::to_string(&labels_map) { + Ok(labels_json) => Some(labels_json), + Err(e) => { + cleanup_pre_stored_provider_credentials( + credentials, + candidate.object_name(), + &credential_update.pre_stored_handles, + ) + .await; + return Err(Status::internal(format!("serialize labels failed: {e}"))); + } + } }; // Write validated candidate with CAS condition - let result = store + let write_result = store .put_if( Provider::object_type(), candidate.object_id(), @@ -342,22 +364,29 @@ async fn update_provider_record_validating( labels_json.as_deref(), WriteCondition::MatchResourceVersion(cas_version), ) - .await - .map_err(|e| { - if matches!(e, crate::persistence::PersistenceError::Conflict { .. }) { - Status::aborted(format!( - "provider was modified concurrently (current resource_version: {})", - match e { - crate::persistence::PersistenceError::Conflict { - current_resource_version, - } => current_resource_version.unwrap_or(0), - _ => 0, - } - )) - } else { - Status::internal(format!("update provider failed: {e}")) - } - })?; + .await; + + let result = match write_result { + Ok(result) => result, + Err(e) => { + cleanup_pre_stored_provider_credentials( + credentials, + candidate.object_name(), + &credential_update.pre_stored_handles, + ) + .await; + return Err(provider_update_persistence_error_to_status(e)); + } + }; + + finish_provider_credential_update( + credentials, + candidate.object_name(), + credential_update, + &removed_credential_handles, + &existing_handles, + ) + .await?; // Update resource_version from successful write if let Some(metadata) = candidate.metadata.as_mut() { @@ -564,22 +593,132 @@ fn credential_handles_removed_by_update( .collect() } -async fn delete_removed_provider_credentials( +#[derive(Debug, Clone, Default)] +struct ProviderCredentialUpdate { + pre_stored_handles: std::collections::HashMap, + deferred_store_values: std::collections::HashMap, + replaced_handles: std::collections::HashMap, +} + +async fn prepare_provider_credential_update( credentials: Option<&crate::credentials::CredentialRuntime>, provider_name: &str, + updated_values: &std::collections::HashMap, + existing_handles: &std::collections::HashMap, +) -> Result { + let Some(credentials) = credentials else { + return Ok(ProviderCredentialUpdate::default()); + }; + if !credentials.stores_provider_credentials() || updated_values.is_empty() { + return Ok(ProviderCredentialUpdate::default()); + } + + let mut update = ProviderCredentialUpdate::default(); + let mut values_requiring_new_handles = std::collections::HashMap::new(); + for (credential_key, value) in updated_values { + match existing_handles.get(credential_key) { + Some(existing_handle) if credentials.storage_owns_handle(existing_handle) => { + update + .deferred_store_values + .insert(credential_key.clone(), value.clone()); + } + Some(replaced_handle) => { + values_requiring_new_handles.insert(credential_key.clone(), value.clone()); + update + .replaced_handles + .insert(credential_key.clone(), replaced_handle.clone()); + } + None => { + values_requiring_new_handles.insert(credential_key.clone(), value.clone()); + } + } + } + + if !values_requiring_new_handles.is_empty() { + update.pre_stored_handles = credentials + .store_provider_credentials( + provider_name, + &values_requiring_new_handles, + &std::collections::HashMap::new(), + ) + .await?; + } + + Ok(update) +} + +async fn finish_provider_credential_update( + credentials: Option<&crate::credentials::CredentialRuntime>, + provider_name: &str, + update: ProviderCredentialUpdate, removed_handles: &std::collections::HashMap, + existing_handles: &std::collections::HashMap, ) -> Result<(), Status> { - if removed_handles.is_empty() { + let Some(credentials) = credentials else { + return Ok(()); + }; + if !credentials.stores_provider_credentials() { return Ok(()); } - let credentials = credentials.ok_or_else(|| { - Status::failed_precondition( - "provider credential handles require a configured credential runtime", - ) - })?; - credentials - .delete_provider_credential_handles(provider_name, removed_handles) + + if !update.deferred_store_values.is_empty() { + credentials + .store_provider_credentials( + provider_name, + &update.deferred_store_values, + existing_handles, + ) + .await?; + } + + let mut handles_to_delete = removed_handles.clone(); + handles_to_delete.extend(update.replaced_handles); + if !handles_to_delete.is_empty() { + credentials + .delete_provider_credential_handles(provider_name, &handles_to_delete) + .await?; + } + + Ok(()) +} + +async fn cleanup_pre_stored_provider_credentials( + credentials: Option<&crate::credentials::CredentialRuntime>, + provider_name: &str, + handles: &std::collections::HashMap, +) { + if handles.is_empty() { + return; + } + let Some(credentials) = credentials else { + return; + }; + if let Err(err) = credentials + .delete_provider_credential_handles(provider_name, handles) .await + { + warn!( + provider_name = %provider_name, + error = %err, + "failed to clean up staged provider credentials after provider update failure" + ); + } +} + +fn provider_update_persistence_error_to_status( + err: crate::persistence::PersistenceError, +) -> Status { + if let crate::persistence::PersistenceError::Conflict { + current_resource_version, + } = err + { + Status::aborted(format!( + "provider was modified concurrently (current resource_version: {})", + current_resource_version.unwrap_or(0) + )) + } else { + Status::internal(format!("update provider failed: {err}")) + } } async fn store_provider_credentials_if_configured( @@ -6714,6 +6853,69 @@ mod tests { assert!(!unchanged.credential_handles.contains_key("NEW_KEY")); } + #[tokio::test] + async fn update_provider_stale_version_does_not_overwrite_stored_credential() { + let mut state = test_server_state().await; + let config = state + .config + .clone() + .with_credential_drivers(["test-static"]); + let credentials = crate::credentials::CredentialRuntime::from_config(&config).unwrap(); + let state_mut = Arc::get_mut(&mut state).unwrap(); + state_mut.config = config; + state_mut.credentials = credentials; + + handle_create_provider( + &state, + Request::new(CreateProviderRequest { + provider: Some(provider_with_credential_value( + "openai-local", + "openai", + "OPENAI_API_KEY", + "sk-first", + )), + }), + ) + .await + .unwrap(); + + let current = state + .store + .get_message_by_name::("openai-local") + .await + .unwrap() + .unwrap(); + let mut stale_provider = current.clone(); + stale_provider.credential_handles.clear(); + stale_provider + .credentials + .insert("OPENAI_API_KEY".to_string(), "sk-stale".to_string()); + stale_provider.metadata.as_mut().unwrap().resource_version = 99; + + let err = handle_update_provider( + &state, + Request::new(UpdateProviderRequest { + provider: Some(stale_provider), + credential_expires_at_ms: HashMap::new(), + }), + ) + .await + .unwrap_err(); + assert_eq!(err.code(), Code::Aborted); + + let resolved = resolve_provider_environment_with_credentials( + state.store.as_ref(), + &["openai-local".to_string()], + &state.credentials, + ) + .await + .unwrap(); + assert_eq!( + resolved.get("OPENAI_API_KEY"), + Some(&"sk-first".to_string()) + ); + } + #[tokio::test] async fn update_provider_concurrent_updates_with_stale_versions() { use std::sync::Arc;