From db50fb17c8e73d6d79ceff0af7c495bb618a92f3 Mon Sep 17 00:00:00 2001 From: Illia Polosukhin Date: Fri, 27 Mar 2026 19:50:00 -0700 Subject: [PATCH] fix: review fixes for custom LLM provider PR - Add server-side validation of custom provider ID format (lowercase alphanumeric + hyphens, 1-64 chars) to match frontend regex - Tighten is_nearai_private_endpoint to exact-match private.near.ai or *.private.near.ai, rejecting lookalikes like private-evil.near.ai - Fix misleading priority doc comments in config/mod.rs and settings.rs to reflect the split model: LLM uses DB > env, others use env > DB - Clean up #1581 artifacts: remove TOML file creation from persist_selected_model (DB is sufficient), update stale priority comments in commands.rs, fix contradictory test assertions - Add 18 new tests for provider ID validation, adapter validation, and nearai private endpoint matching Co-Authored-By: Claude Opus 4.6 (1M context) --- src/agent/commands.rs | 36 +++----- src/channels/web/handlers/settings.rs | 118 +++++++++++++++++++++++++- src/channels/web/server.rs | 41 ++++++++- src/config/mod.rs | 19 +++-- src/settings.rs | 74 ++++++++-------- 5 files changed, 214 insertions(+), 74 deletions(-) diff --git a/src/agent/commands.rs b/src/agent/commands.rs index 643d8c7c..04a1022a 100644 --- a/src/agent/commands.rs +++ b/src/agent/commands.rs @@ -947,6 +947,12 @@ impl Agent { /// Best-effort: logs warnings on failure but does not propagate errors, /// since the in-memory model switch already succeeded. /// + /// The DB setting is the primary persistence layer. For LLM settings the + /// resolution priority is `DB > env > TOML > default`, so writing to DB + /// is sufficient for the change to survive restarts. The `.env` and TOML + /// files are only updated as a courtesy when they already contain a model + /// var, to avoid user confusion. + /// /// In multi-tenant mode, only the per-user DB setting is written — global /// .env and TOML files are shared across users and must not be mutated. async fn persist_selected_model(&self, tenant: &crate::tenant::TenantCtx, model: &str) { @@ -972,22 +978,18 @@ impl Agent { return; } - // 3. Update .env and TOML config file (sync I/O in spawn_blocking). + // 3. Best-effort update of .env and TOML if they already contain a + // model var. DB is authoritative (DB > env > TOML), but keeping + // these in sync avoids confusion when users inspect the files. let model_owned = model.to_string(); let backend = self.deps.llm_backend.clone(); if let Err(e) = tokio::task::spawn_blocking(move || { - // 2a. Update the backend-specific model env var in ~/.ironclaw/.env. - // - // Env vars have the HIGHEST priority in LlmConfig::resolve_model() - // (env var > TOML > DB > default). If the .env file has e.g. - // NEARAI_MODEL=old-model, it shadows everything else. We must - // update this var or the /model change is invisible on restart. + // 3a. Update the backend-specific model env var in ~/.ironclaw/.env + // only if the var already exists (don't inject new vars). let registry = crate::llm::ProviderRegistry::load(); let model_env = registry.model_env_var(&backend); let env_var_prefix = format!("{}=", model_env); - // Only update the .env file if the var is actually set there - // (avoid injecting new vars the user never configured). let env_path = crate::bootstrap::ironclaw_env_path(); let env_has_var = std::fs::read_to_string(&env_path) .ok() @@ -1005,10 +1007,8 @@ impl Agent { } } - // 2b. Update (or create) the TOML config file. - // - // The TOML overlay has higher priority than DB settings on - // startup, so it MUST stay in sync with the DB. + // 3b. Update TOML config file if it already exists. + // Don't create a new one — DB persistence is sufficient. let toml_path = crate::settings::Settings::default_toml_path(); match crate::settings::Settings::load_toml(&toml_path) { Ok(Some(mut settings)) => { @@ -1018,15 +1018,7 @@ impl Agent { } } Ok(None) => { - // No config file yet — create one so the model choice - // survives restarts even when the DB is unavailable. - let settings = crate::settings::Settings { - selected_model: Some(model_owned), - ..Default::default() - }; - if let Err(e) = settings.save_toml(&toml_path) { - tracing::warn!("Failed to create config.toml for model persistence: {}", e); - } + // No config file on disk; DB persistence is sufficient. } Err(e) => { tracing::warn!("Failed to load config.toml for model persistence: {}", e); diff --git a/src/channels/web/handlers/settings.rs b/src/channels/web/handlers/settings.rs index d308e359..1ca40952 100644 --- a/src/channels/web/handlers/settings.rs +++ b/src/channels/web/handlers/settings.rs @@ -116,7 +116,7 @@ pub async fn settings_set_handler( // Guard: cannot remove a custom provider that is currently active. if key == "llm_custom_providers" { guard_active_provider_not_removed(store, &user.user_id, &body.value).await?; - validate_custom_providers_adapters(&body.value)?; + validate_custom_providers(&body.value)?; } // Extract API keys from LLM settings and vault them in the secrets store. @@ -144,6 +144,34 @@ pub async fn settings_set_handler( const VALID_ADAPTERS: &[&str] = &["open_ai_completions", "anthropic", "ollama"]; +/// Valid provider ID: lowercase alphanumeric and hyphens, 1-64 chars. +fn is_valid_provider_id(id: &str) -> bool { + !id.is_empty() + && id.len() <= 64 + && id + .bytes() + .all(|b| b.is_ascii_lowercase() || b.is_ascii_digit() || b == b'-') +} + +/// Returns `Err(422)` if any provider has an invalid ID or unrecognised adapter. +fn validate_custom_providers(value: &serde_json::Value) -> Result<(), StatusCode> { + let providers = match value.as_array() { + Some(arr) => arr, + None => return Ok(()), + }; + for p in providers { + let id = p.get("id").and_then(|v| v.as_str()).unwrap_or(""); + if !is_valid_provider_id(id) { + tracing::warn!( + id = %id, + "Rejected custom provider with invalid ID (must be lowercase alphanumeric/hyphens, 1-64 chars)" + ); + return Err(StatusCode::UNPROCESSABLE_ENTITY); + } + } + validate_custom_providers_adapters(value) +} + /// Returns `Err(422)` if any provider in the incoming list has an unrecognised adapter. fn validate_custom_providers_adapters(value: &serde_json::Value) -> Result<(), StatusCode> { let providers = match value.as_array() { @@ -846,4 +874,92 @@ mod tests { .unwrap_err(); assert_eq!(err, StatusCode::SERVICE_UNAVAILABLE); } + + // --- Provider ID validation tests --- + + #[test] + fn test_valid_provider_ids() { + assert!(is_valid_provider_id("my-llm")); + assert!(is_valid_provider_id("openai")); + assert!(is_valid_provider_id("custom-provider-123")); + assert!(is_valid_provider_id("a")); + } + + #[test] + fn test_invalid_provider_ids() { + assert!(!is_valid_provider_id(""), "empty ID"); + assert!(!is_valid_provider_id("My-LLM"), "uppercase"); + assert!(!is_valid_provider_id("my llm"), "spaces"); + assert!(!is_valid_provider_id("my_llm"), "underscores"); + assert!(!is_valid_provider_id("../../etc"), "path traversal"); + assert!(!is_valid_provider_id("a.b"), "dots"); + assert!( + !is_valid_provider_id(&"a".repeat(65)), + "exceeds 64 char limit" + ); + } + + #[test] + fn test_validate_custom_providers_rejects_bad_id() { + let input = serde_json::json!([ + { "id": "UPPER-CASE", "adapter": "open_ai_completions" } + ]); + assert_eq!( + validate_custom_providers(&input).unwrap_err(), + StatusCode::UNPROCESSABLE_ENTITY, + ); + } + + #[test] + fn test_validate_custom_providers_accepts_valid() { + let input = serde_json::json!([ + { "id": "my-llm", "adapter": "open_ai_completions" }, + { "id": "local-ollama", "adapter": "ollama" } + ]); + assert!(validate_custom_providers(&input).is_ok()); + } + + // --- Adapter validation tests --- + + #[test] + fn test_validate_adapters_rejects_unknown() { + let input = serde_json::json!([ + { "id": "test", "adapter": "not_a_real_adapter" } + ]); + assert_eq!( + validate_custom_providers_adapters(&input).unwrap_err(), + StatusCode::UNPROCESSABLE_ENTITY, + ); + } + + #[test] + fn test_validate_adapters_rejects_missing() { + let input = serde_json::json!([ + { "id": "test" } + ]); + assert_eq!( + validate_custom_providers_adapters(&input).unwrap_err(), + StatusCode::UNPROCESSABLE_ENTITY, + ); + } + + #[test] + fn test_validate_adapters_accepts_all_valid() { + for adapter in VALID_ADAPTERS { + let input = serde_json::json!([ + { "id": "test", "adapter": adapter } + ]); + assert!( + validate_custom_providers_adapters(&input).is_ok(), + "adapter '{}' should be accepted", + adapter + ); + } + } + + #[test] + fn test_validate_adapters_non_array_is_ok() { + let input = serde_json::json!("not-an-array"); + assert!(validate_custom_providers_adapters(&input).is_ok()); + } } diff --git a/src/channels/web/server.rs b/src/channels/web/server.rs index 051b6774..7fcd0025 100644 --- a/src/channels/web/server.rs +++ b/src/channels/web/server.rs @@ -2747,14 +2747,16 @@ async fn resolve_api_key_from_secrets( } } -/// Check if a base URL belongs to a NEAR AI private endpoint by verifying the -/// hostname ends with `.near.ai` and contains "private". This prevents an -/// attacker from crafting `https://evil.com/private/...` to match. +/// Check if a base URL belongs to a NEAR AI private endpoint. +/// +/// Matches `private.near.ai` exactly or any subdomain of it +/// (e.g. `us.private.near.ai`). Rejects lookalikes like +/// `private-evil.near.ai` or `myprivate.near.ai`. fn is_nearai_private_endpoint(base_url: &str) -> bool { url::Url::parse(base_url) .ok() .and_then(|u| u.host_str().map(|h| h.to_lowercase())) - .is_some_and(|host| host.ends_with(".near.ai") && host.contains("private")) + .is_some_and(|host| host == "private.near.ai" || host.ends_with(".private.near.ai")) } async fn test_provider_connection(req: TestConnectionRequest) -> TestConnectionResponse { @@ -4772,4 +4774,35 @@ mod tests { assert!(!is_local_origin("not-a-url")); assert!(!is_local_origin("")); } + + // --- is_nearai_private_endpoint tests --- + + #[test] + fn test_nearai_private_exact_match() { + assert!(is_nearai_private_endpoint("https://private.near.ai/v1")); + } + + #[test] + fn test_nearai_private_subdomain() { + assert!(is_nearai_private_endpoint("https://us.private.near.ai/v1")); + } + + #[test] + fn test_nearai_public_endpoint_not_private() { + assert!(!is_nearai_private_endpoint("https://cloud-api.near.ai/v1")); + } + + #[test] + fn test_nearai_private_lookalike_rejected() { + // "private" appears in the hostname but not as the correct domain + assert!(!is_nearai_private_endpoint( + "https://private-evil.near.ai/v1" + )); + assert!(!is_nearai_private_endpoint("https://myprivate.near.ai/v1")); + } + + #[test] + fn test_nearai_private_non_near_ai_rejected() { + assert!(!is_nearai_private_endpoint("https://private.evil.com/v1")); + } } diff --git a/src/config/mod.rs b/src/config/mod.rs index d5a1df71..241eacb8 100644 --- a/src/config/mod.rs +++ b/src/config/mod.rs @@ -1,9 +1,13 @@ //! Configuration for IronClaw. //! -//! Settings are loaded with priority: database > env var > default. +//! Settings are loaded from env vars, the DB settings table, TOML config, +//! and built-in defaults. Priority varies by subsystem: +//! +//! - **LLM settings** (backend, model, api_key, base_url): DB > env > default +//! - **Most other settings** (agent, channels, tunnel, …): env > DB > default +//! //! `DATABASE_URL` lives in `~/.ironclaw/.env` (loaded via dotenvy early -//! in startup). Everything else comes from env vars, the DB settings -//! table, or auto-detection. +//! in startup). mod agent; mod builder; @@ -186,8 +190,9 @@ impl Config { /// Load configuration from environment variables and the database. /// - /// Priority: DB settings > env var > TOML config file > default. - /// This is the primary way to load config after DB is connected. + /// TOML is loaded first as a base, then DB values are merged on top + /// (DB wins over TOML). Individual subsystem resolvers then apply + /// their own env-vs-DB priority — see module docs for details. pub async fn from_db( store: &(dyn crate::db::SettingsStore + Sync), user_id: &str, @@ -197,9 +202,9 @@ impl Config { /// Load from DB with an optional TOML config file overlay. /// - /// Priority: DB settings > env var > TOML config file > default. /// TOML is loaded first as a base, then DB values are merged on top - /// so that DB always wins over TOML. + /// (DB wins over TOML). Per-subsystem resolvers then decide whether + /// env vars or DB values take final precedence — see module docs. pub async fn from_db_with_toml( store: &(dyn crate::db::SettingsStore + Sync), user_id: &str, diff --git a/src/settings.rs b/src/settings.rs index f0039d8e..46390f9d 100644 --- a/src/settings.rs +++ b/src/settings.rs @@ -902,7 +902,8 @@ impl Settings { let content = format!( "# IronClaw configuration file.\n\ #\n\ - # Priority: database settings > env var > this file > defaults.\n\ + # Priority varies by subsystem. LLM: DB > env > this file > defaults.\n\ + # Most others: env > DB > this file > defaults.\n\ # Uncomment and edit values to override defaults.\n\ # Run `ironclaw config init` to regenerate this file.\n\ #\n\ @@ -1386,56 +1387,53 @@ mod tests { ); } - /// Regression: TOML overlay must not clobber a DB-persisted selected_model - /// when the TOML file matches the DB. This is the normal case after /model - /// successfully writes to both DB and TOML. + /// TOML is loaded as a base, then DB is merged on top (DB wins). + /// When both agree, the result matches. #[test] - fn toml_overlay_preserves_matching_model() { - // DB settings with new model from /model command. - let mut db_settings = Settings { + fn toml_and_db_matching_model_preserved() { + // from_db_with_toml: TOML base, then DB merged on top. + let mut toml_base = Settings { + selected_model: Some("new-model".to_string()), + ..Default::default() + }; + + let db_overlay = Settings { llm_backend: Some("nearai".to_string()), selected_model: Some("new-model".to_string()), ..Default::default() }; - // TOML also updated by /model command to the same value. - let toml_settings = Settings { - selected_model: Some("new-model".to_string()), - ..Default::default() - }; - - db_settings.merge_from(&toml_settings); + toml_base.merge_from(&db_overlay); assert_eq!( - db_settings.selected_model, + toml_base.selected_model, Some("new-model".to_string()), - "TOML overlay must not clobber matching model" + "matching values: result should be the shared value" ); } - /// Regression: when /model updates DB but TOML write fails, a stale TOML - /// file would overwrite the DB value. This test documents the priority: - /// TOML > DB (by design). persist_selected_model MUST update the TOML. + /// Regression: when TOML has a stale model but DB has been updated via + /// /model command, DB must win. This matches from_db_with_toml where + /// TOML is loaded first as base, then DB is merged on top. #[test] - fn stale_toml_overwrites_db_model() { - // DB has the new model from /model. - let mut db_settings = Settings { - selected_model: Some("new-model".to_string()), - ..Default::default() - }; - - // TOML still has the old model (write failed or was not attempted). - let stale_toml = Settings { + fn db_model_wins_over_stale_toml() { + // TOML base with old model. + let mut toml_base = Settings { selected_model: Some("old-model".to_string()), ..Default::default() }; - db_settings.merge_from(&stale_toml); - // This documents the current priority: TOML wins over DB. - // The fix in persist_selected_model ensures TOML is always updated. + // DB has the new model from /model command. + let db_overlay = Settings { + selected_model: Some("new-model".to_string()), + ..Default::default() + }; + + // from_db_with_toml: TOML first, then DB merged on top. + toml_base.merge_from(&db_overlay); assert_eq!( - db_settings.selected_model, - Some("old-model".to_string()), - "TOML overlay has higher priority than DB (by design)" + toml_base.selected_model, + Some("new-model".to_string()), + "DB selected_model must win over stale TOML value" ); } @@ -1464,24 +1462,20 @@ mod tests { assert_eq!(reloaded.selected_model, Some("new-model".to_string())); } - /// Regression: /model must create config.toml when it doesn't exist, so the - /// model survives restarts. Previously the Ok(None) case was a no-op. + /// save_toml / load_toml round-trip for selected_model. #[test] - fn toml_created_when_missing_for_model_persist() { + fn toml_save_and_load_round_trip() { let dir = tempfile::tempdir().unwrap(); let path = dir.path().join("config.toml"); - // No config.toml yet (fresh install, no wizard). assert!(Settings::load_toml(&path).unwrap().is_none()); - // Simulate what persist_selected_model now does for the Ok(None) case. let settings = Settings { selected_model: Some("new-model".to_string()), ..Default::default() }; settings.save_toml(&path).unwrap(); - // Verify the model survived. let loaded = Settings::load_toml(&path).unwrap().unwrap(); assert_eq!(loaded.selected_model, Some("new-model".to_string())); }