feat(extensions): unify chat and gateway auth paths, clean up secrets on remove#677
feat(extensions): unify chat and gateway auth paths, clean up secrets on remove#677henrypark133 wants to merge 1 commit intomainfrom
Conversation
github-actions
bot
added
scope: extensions
Summary of ChangesHello, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request significantly enhances the robustness and consistency of extension management, particularly for chat channels. It addresses issues with lingering secrets after extension removal and unifies the authentication process, making it more reliable and extensible. The changes ensure that authentication flows are standardized and that sensitive data is properly cleaned up, improving the overall security and user experience of extensions. Highlights
Changelog
Activity
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces significant improvements to the extension authentication and removal process, unifying the auth paths for chat and the Extensions tab for greater consistency. However, it introduces a High-severity SSRF vulnerability in the new secret validation logic and misses path traversal validation in the save_setup_secrets method. Additionally, there are suggestions to improve the robustness of the new error handling, enhance performance by using a more efficient data structure, and address code duplication for better maintainability.
| let resp = reqwest::Client::builder() | ||
| .timeout(std::time::Duration::from_secs(10)) | ||
| .build() | ||
| .map_err(|e| ExtensionError::Other(e.to_string()))? | ||
| .get(&url) | ||
| .send() | ||
| .await | ||
| .map_err(|e| { |
There was a problem hiding this comment.
The save_setup_secrets method uses a validation_endpoint URL template from the extension's capabilities file to perform a request from the host process. This URL is not validated against an allowlist, allowing a malicious extension to trigger SSRF requests to internal services or leak secrets to external servers. Implement a strict allowlist for validation_endpoint URLs or ensure they match the extension's declared HTTP allowlist. At a minimum, block requests to loopback and private IP ranges.
References
- This rule emphasizes the importance of resolving hostnames to IP addresses and validating all resolved IPs against restricted ranges to prevent DNS-based SSRF vulnerabilities, aligning with the need to block requests to loopback and private IP ranges.
src/extensions/manager.rs
Outdated
| let msg = e.to_string(); | ||
| if msg.contains("Invalid token") | ||
| || msg.contains("API returned") | ||
| || msg.contains("validate token") | ||
| { |
There was a problem hiding this comment.
Relying on string matching (.contains()) for error handling is brittle and can lead to bugs if the error messages change. It would be more robust to introduce a specific error variant for token validation failures.
For example, you could add a TokenValidationFailed(String) variant to your ExtensionError enum. Then, save_setup_secrets can return Err(ExtensionError::TokenValidationFailed(...)) and this code can handle it explicitly with if let ExtensionError::TokenValidationFailed(msg) = e or a match statement. This would make the error handling type-safe.
References
- This rule advocates for creating specific error variants for different failure modes to provide semantically correct and clear error messages, which aligns with the suggestion to use a specific error variant for token validation failures instead of brittle string matching.
| if kind == ExtensionKind::WasmChannel { | ||
| let cap_path = self | ||
| .wasm_channels_dir | ||
| .join(format!("{}.capabilities.json", name)); |
There was a problem hiding this comment.
The save_setup_secrets method uses the name parameter to construct file paths without calling validate_extension_name. This allows an attacker to provide a name with path traversal sequences (e.g., ..) to interact with files outside the intended extensions directory. Although the impact is limited by the forced .capabilities.json extension, it is recommended to call Self::validate_extension_name(name)?; at the beginning of the method to ensure the name is safe.
References
- This rule requires validation of path traversal characters in extension names at the public API boundary to prevent security vulnerabilities, directly supporting the recommendation to validate the
nameparameter.
src/extensions/manager.rs
Outdated
| let mut still_missing = Vec::new(); | ||
| for s in required_secrets { | ||
| if s.optional { | ||
| continue; | ||
| } | ||
| if !self | ||
| .secrets | ||
| .exists(&self.user_id, &s.name) | ||
| .await | ||
| .unwrap_or(false) | ||
| { | ||
| still_missing.push(s); | ||
| } | ||
| } |
There was a problem hiding this comment.
This block of code to find missing secrets is a duplicate of the logic at lines 2115-2127. To improve maintainability and reduce duplication, consider extracting this logic into a private helper function, for example:
async fn get_missing_secrets<'a>(
&self,
required_secrets: &'a [crate::channels::wasm::SecretSetupSchema],
) -> Vec<&'a crate::channels::wasm::SecretSetupSchema> {
// ... implementation ...
}You could then call self.get_missing_secrets(required_secrets).await in both places.
src/extensions/manager.rs
Outdated
| /// 1. Read secret names from the capabilities file (while it still exists) | ||
| /// 2. Fallback: list all secrets where `provider == extension_name` | ||
| async fn cleanup_extension_secrets(&self, name: &str, kind: ExtensionKind) { | ||
| let mut deleted_names: Vec<String> = Vec::new(); |
There was a problem hiding this comment.
For better performance, consider using a HashSet<String> for deleted_names instead of a Vec<String>. The contains() check on line 3247 has O(n) time complexity for a Vec, but would be O(1) on average for a HashSet. This is more efficient, especially if an extension has many secrets.
To make this change, you'll also need to replace all calls to .push() on deleted_names with .insert().
| let mut deleted_names: Vec<String> = Vec::new(); | |
| let mut deleted_names: std::collections::HashSet<String> = std::collections::HashSet::new(); |
There was a problem hiding this comment.
Pull request overview
This PR addresses inconsistent authentication behavior between chat-driven setup and the Extensions UI, and fixes stale secrets persisting across uninstall/reinstall cycles for extensions.
Changes:
- Add secret cleanup during
ExtensionManager::remove()using capabilities-based secret enumeration plus a provider-tag fallback. - Unify chat auth for WASM channels by routing token entry through
save_setup_secrets()(enabling validation, auto-generation, and activation). - Replace Telegram-specific token validation with a generic
validation_endpointdeclared in channel capabilities (Telegram updated accordingly).
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 4 comments.
| File | Description |
|---|---|
src/extensions/manager.rs |
Implements secret cleanup on remove, unifies WASM channel auth path via save_setup_secrets, and adds capabilities-driven token validation. |
channels-src/telegram/telegram.capabilities.json |
Adds validation_endpoint for Telegram bot token validation via the Telegram API. |
Comments suppressed due to low confidence (1)
src/extensions/manager.rs:2988
save_setup_secretsreads and parses the channel capabilities file multiple times (allowed-secret collection, validation_endpoint, and auto-generate blocks). This adds redundant I/O/JSON parsing and makes the logic harder to keep consistent. Consider loadingcap_fileonce for the channel path and reusing it for allowed names, validation, and auto-generation.
// Load allowed secret names from the extension's capabilities file
let allowed: std::collections::HashSet<String> = match kind {
ExtensionKind::WasmChannel => {
let cap_path = self
.wasm_channels_dir
.join(format!("{}.capabilities.json", name));
if !cap_path.exists() {
return Err(ExtensionError::Other(format!(
"Capabilities file not found for '{}'",
name
)));
}
let cap_bytes = tokio::fs::read(&cap_path)
.await
.map_err(|e| ExtensionError::Other(e.to_string()))?;
let cap_file =
crate::channels::wasm::ChannelCapabilitiesFile::from_bytes(&cap_bytes)
.map_err(|e| ExtensionError::Other(e.to_string()))?;
cap_file
.setup
.required_secrets
.iter()
.map(|s| s.name.clone())
.collect()
}
ExtensionKind::WasmTool => {
let cap_file = self.load_tool_capabilities(name).await.ok_or_else(|| {
ExtensionError::Other(format!("Capabilities file not found for '{}'", name))
})?;
match cap_file.setup {
Some(s) => s.required_secrets.iter().map(|s| s.name.clone()).collect(),
None => {
return Err(ExtensionError::Other(format!(
"Tool '{}' has no setup schema — no secrets to configure",
name
)));
}
}
}
_ => {
return Err(ExtensionError::Other(
"Setup is only supported for WASM channels and tools".to_string(),
));
}
};
// Validate secrets against the validation_endpoint if declared in capabilities.
// The endpoint URL template uses {secret_name} placeholders that are
// substituted with the provided secret value before making the request.
if kind == ExtensionKind::WasmChannel {
let cap_path = self
.wasm_channels_dir
.join(format!("{}.capabilities.json", name));
if let Ok(cap_bytes) = tokio::fs::read(&cap_path).await
&& let Ok(cap_file) =
crate::channels::wasm::ChannelCapabilitiesFile::from_bytes(&cap_bytes)
&& let Some(ref endpoint_template) = cap_file.setup.validation_endpoint
&& let Some(secret_def) = cap_file
.setup
.required_secrets
.iter()
.find(|s| !s.optional && secrets.contains_key(&s.name))
&& let Some(token_value) = secrets.get(&secret_def.name)
{
let token = token_value.trim();
if !token.is_empty() {
let encoded =
url::form_urlencoded::byte_serialize(token.as_bytes()).collect::<String>();
let url =
endpoint_template.replace(&format!("{{{}}}", secret_def.name), &encoded);
let resp = reqwest::Client::builder()
.timeout(std::time::Duration::from_secs(10))
.build()
.map_err(|e| ExtensionError::Other(e.to_string()))?
.get(&url)
.send()
.await
.map_err(|e| {
ExtensionError::Other(format!("Failed to validate token: {}", e))
})?;
if !resp.status().is_success() {
return Err(ExtensionError::Other(format!(
"Invalid token (API returned {})",
resp.status()
)));
}
}
}
}
// Validate and store each submitted secret
for (secret_name, secret_value) in secrets {
if !allowed.contains(secret_name.as_str()) {
return Err(ExtensionError::Other(format!(
"Unknown secret '{}' for extension '{}'",
secret_name, name
)));
}
if secret_value.trim().is_empty() {
continue;
}
let params =
CreateSecretParams::new(secret_name, secret_value).with_provider(name.to_string());
self.secrets
.create(&self.user_id, params)
.await
.map_err(|e| ExtensionError::AuthFailed(e.to_string()))?;
}
// Auto-generate any missing secrets (channel-only feature)
if kind == ExtensionKind::WasmChannel {
let cap_path = self
.wasm_channels_dir
.join(format!("{}.capabilities.json", name));
if let Ok(cap_bytes) = tokio::fs::read(&cap_path).await
&& let Ok(cap_file) =
crate::channels::wasm::ChannelCapabilitiesFile::from_bytes(&cap_bytes)
{
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
src/extensions/manager.rs
Outdated
| ExtensionKind::McpServer => { | ||
| for suffix in ["access_token", "refresh_token"] { | ||
| let secret_name = format!("mcp_{}_{}", name, suffix); | ||
| if self.secrets.delete(&self.user_id, &secret_name).await.ok() == Some(true) { | ||
| deleted_names.push(secret_name); | ||
| } | ||
| } | ||
| } | ||
| } | ||
|
|
||
| // Fallback: delete any secrets tagged with this provider | ||
| if let Ok(all_refs) = self.secrets.list(&self.user_id).await { | ||
| for secret_ref in all_refs { | ||
| if secret_ref.provider.as_deref() == Some(name) | ||
| && !deleted_names.contains(&secret_ref.name) | ||
| { | ||
| let _ = self.secrets.delete(&self.user_id, &secret_ref.name).await; | ||
| deleted_names.push(secret_ref.name); | ||
| } |
There was a problem hiding this comment.
cleanup_extension_secrets doesn’t fully clean up MCP server secrets: MCP auth stores a DCR client id under mcp_{name}_client_id (see McpServerConfig::client_id_secret_name() usage), but the removal path only deletes access/refresh tokens. Also, MCP secrets are tagged with provider mcp:{name}, so the provider-tag fallback (provider == name) won’t catch any remaining MCP-tagged secrets. Consider deleting the mcp_{name}_client_id secret explicitly and extending the provider fallback to include mcp:{name} (or a mcp: prefix match).
src/extensions/manager.rs
Outdated
| Err(e) => { | ||
| // Token validation errors (bad token, API rejected) should | ||
| // re-prompt the user. Infrastructure errors (DB down, file | ||
| // read failure) should propagate to the caller. | ||
| let msg = e.to_string(); | ||
| if msg.contains("Invalid token") | ||
| || msg.contains("API returned") | ||
| || msg.contains("validate token") | ||
| { | ||
| return Ok(AuthResult::awaiting_token( | ||
| name, | ||
| ExtensionKind::WasmChannel, | ||
| format!( | ||
| "Authentication failed: {}. Please provide a valid token.", | ||
| e | ||
| ), | ||
| cap_file.setup.setup_url.clone(), | ||
| )); | ||
| } | ||
| return Err(e); | ||
| } |
There was a problem hiding this comment.
The error handling here classifies failures by substring matching on e.to_string(), and will treat infrastructure failures as “token validation” in some cases (e.g., save_setup_secrets returns Other("Failed to validate token: ...") on network/timeout, which matches "validate token" and triggers a user re-prompt). Also, the subsequent exists(...).unwrap_or(false) re-check will silently convert SecretsStore errors into “missing secret”, again re-prompting instead of surfacing an infra problem. Consider adding a structured error variant for validation failures (vs infra) and propagating SecretsStore read errors rather than swallowing them.
| if let Ok(cap_bytes) = tokio::fs::read(&cap_path).await | ||
| && let Ok(cap_file) = | ||
| crate::channels::wasm::ChannelCapabilitiesFile::from_bytes(&cap_bytes) | ||
| && let Some(ref endpoint_template) = cap_file.setup.validation_endpoint | ||
| && let Some(secret_def) = cap_file | ||
| .setup | ||
| .required_secrets | ||
| .iter() | ||
| .find(|s| !s.optional && secrets.contains_key(&s.name)) | ||
| && let Some(token_value) = secrets.get(&secret_def.name) | ||
| { | ||
| let token = token_value.trim(); | ||
| if !token.is_empty() { | ||
| let encoded = | ||
| url::form_urlencoded::byte_serialize(token.as_bytes()).collect::<String>(); | ||
| let url = | ||
| endpoint_template.replace(&format!("{{{}}}", secret_def.name), &encoded); | ||
| let resp = reqwest::Client::builder() | ||
| .timeout(std::time::Duration::from_secs(10)) | ||
| .build() | ||
| .map_err(|e| ExtensionError::Other(e.to_string()))? | ||
| .get(&url) | ||
| .send() | ||
| .await | ||
| .map_err(|e| { | ||
| ExtensionError::Other(format!("Failed to validate token: {}", e)) | ||
| })?; | ||
| if !resp.status().is_success() { | ||
| return Err(ExtensionError::Other(format!( | ||
| "Invalid token (API returned {})", | ||
| resp.status() | ||
| ))); |
There was a problem hiding this comment.
validation_endpoint substitution/selection logic is too narrow for the stated “capabilities-driven” behavior: it picks the first non-optional secret present in the submitted map and only replaces that one {secret_name} placeholder. If the template references a different secret (or multiple secrets), placeholders can remain unexpanded and the request may hit an invalid URL, potentially blocking unrelated secret updates. Consider (1) replacing placeholders for all provided secrets (and/or validating that every placeholder in the template is satisfied) and (2) failing fast with a clear error when the template can’t be fully rendered.
| if let Ok(cap_bytes) = tokio::fs::read(&cap_path).await | |
| && let Ok(cap_file) = | |
| crate::channels::wasm::ChannelCapabilitiesFile::from_bytes(&cap_bytes) | |
| && let Some(ref endpoint_template) = cap_file.setup.validation_endpoint | |
| && let Some(secret_def) = cap_file | |
| .setup | |
| .required_secrets | |
| .iter() | |
| .find(|s| !s.optional && secrets.contains_key(&s.name)) | |
| && let Some(token_value) = secrets.get(&secret_def.name) | |
| { | |
| let token = token_value.trim(); | |
| if !token.is_empty() { | |
| let encoded = | |
| url::form_urlencoded::byte_serialize(token.as_bytes()).collect::<String>(); | |
| let url = | |
| endpoint_template.replace(&format!("{{{}}}", secret_def.name), &encoded); | |
| let resp = reqwest::Client::builder() | |
| .timeout(std::time::Duration::from_secs(10)) | |
| .build() | |
| .map_err(|e| ExtensionError::Other(e.to_string()))? | |
| .get(&url) | |
| .send() | |
| .await | |
| .map_err(|e| { | |
| ExtensionError::Other(format!("Failed to validate token: {}", e)) | |
| })?; | |
| if !resp.status().is_success() { | |
| return Err(ExtensionError::Other(format!( | |
| "Invalid token (API returned {})", | |
| resp.status() | |
| ))); | |
| if let Ok(cap_bytes) = tokio::fs::read(&cap_path).await { | |
| if let Ok(cap_file) = | |
| crate::channels::wasm::ChannelCapabilitiesFile::from_bytes(&cap_bytes) | |
| { | |
| if let Some(ref endpoint_template) = cap_file.setup.validation_endpoint { | |
| // Render the validation endpoint URL by substituting all relevant | |
| // non-optional required secrets, and ensure no placeholders remain. | |
| let mut url = endpoint_template.clone(); | |
| let mut any_replaced = false; | |
| for secret_def in &cap_file.setup.required_secrets { | |
| if secret_def.optional { | |
| continue; | |
| } | |
| if let Some(raw_value) = secrets.get(&secret_def.name) { | |
| let token = raw_value.trim(); | |
| if token.is_empty() { | |
| continue; | |
| } | |
| let encoded = url::form_urlencoded::byte_serialize( | |
| token.as_bytes(), | |
| ) | |
| .collect::<String>(); | |
| let placeholder = format!("{{{}}}", secret_def.name); | |
| if url.contains(&placeholder) { | |
| url = url.replace(&placeholder, &encoded); | |
| any_replaced = true; | |
| } | |
| } | |
| } | |
| // If at least one substitution occurred, verify that no placeholders remain. | |
| if any_replaced { | |
| if url.contains('{') && url.contains('}') { | |
| return Err(ExtensionError::Other( | |
| "Validation endpoint template contains unresolved placeholders; \ | |
| ensure all required secrets are provided." | |
| .to_string(), | |
| )); | |
| } | |
| let resp = reqwest::Client::builder() | |
| .timeout(std::time::Duration::from_secs(10)) | |
| .build() | |
| .map_err(|e| ExtensionError::Other(e.to_string()))? | |
| .get(&url) | |
| .send() | |
| .await | |
| .map_err(|e| { | |
| ExtensionError::Other(format!( | |
| "Failed to validate token: {}", | |
| e | |
| )) | |
| })?; | |
| if !resp.status().is_success() { | |
| return Err(ExtensionError::Other(format!( | |
| "Invalid token (API returned {})", | |
| resp.status() | |
| ))); | |
| } | |
| } |
src/extensions/manager.rs
Outdated
| let kind = self.determine_installed_kind(name).await?; | ||
|
|
||
| // Clean up secrets before deleting capabilities files | ||
| self.cleanup_extension_secrets(name, kind).await; |
There was a problem hiding this comment.
remove() calls cleanup_extension_secrets() but discards any failures inside that cleanup (most operations are ok()/ignored). If the secrets backend is unavailable, remove will still succeed while leaving stale secrets behind—the exact class of issue this PR is trying to prevent. Consider returning Result<()> from cleanup_extension_secrets and either failing remove() on cleanup errors or at least logging warnings when deletions/listing fail.
| self.cleanup_extension_secrets(name, kind).await; | |
| if let Err(err) = self.cleanup_extension_secrets(name, kind).await { | |
| eprintln!( | |
| "Warning: failed to clean up secrets for extension '{}': {}", | |
| name, err | |
| ); | |
| } |
Refactors the extension lifecycle to eliminate the divergence between chat and gateway paths that caused Telegram setup via chat to fail (missing webhook secret auto-generation, no token validation). Key changes: - Rename save_setup_secrets() → configure(): single entrypoint for providing secrets to any extension (WasmChannel, WasmTool, MCP). Validates, stores, auto-generates, and activates. - Add configure_token(): convenience wrapper for single-token callers (chat auth card, WebSocket, agent auth mode). - Refactor auth() to pure status check: remove token parameter, delete token-storing branches from auth_mcp/auth_wasm_tool, rename auth_wasm_channel → auth_wasm_channel_status. - Add ConfigureResult/MissingSecret types for structured responses. - Replace hardcoded Telegram token validation with generic validation_endpoint from capabilities.json. - Update all callers (9 files) to use the new interface. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
henrypark133
force-pushed
the
unify-extension-lifecycle
branch
from
64a72a0 to
974d916
Compare
github-actions
bot
added
size: L
Summary
remove()now deletes all associated secrets (setup secrets, webhook secrets, OAuth tokens) using capabilities-file enumeration + provider-tag fallback. Fixes stale secrets surviving uninstall/reinstall cycles.auth_wasm_channel()delegates tosave_setup_secrets()when a token is provided, giving chat the same behavior as the Extensions tab: token validation, webhook secret auto-generation, auto-activation.if name == "telegram"token validation with capabilities-drivenvalidation_endpointURL template. Any channel can declare validation in its capabilities.json.Problem
When setting up Telegram via chat (
tool_auth), old invalid bot tokens survivedtool_removeand causedauth()to report "authenticated" while activation failed. The Extensions tab worked because it used a completely different code path (save_setup_secrets) with validation and auto-generation.Test plan
test_remove_channel_cleans_up_secrets— bot token + webhook secret deleted on removetest_remove_cleans_secrets_by_provider_fallback— provider-tagged secrets caught by fallbacktest_auth_after_reinstall_detects_missing_secrets— stale secrets don't survive remove+reinstallcargo clippy --all --all-features— zero warnings🤖 Generated with Claude Code