fix: close all Chutes TEE model validation bypass vulnerabilities

echobt · echobt · commit f6ed270e4e80 · 2026-02-03T13:35:12.000Z
- Add CHUTES_ALLOWED_MODELS whitelist for defense-in-depth validation
- validate_chutes_model() now checks both suffix and whitelist
- provider_allows_custom_models() uses case-insensitive comparison
- ProviderManager::new() validates model from config for Chutes provider
- set_provider() resets model to default TEE model when switching to Chutes
- set_model() uses case-insensitive provider check
- ensure_client() adds defense-in-depth validation before API calls
- Modal handlers now check set_model() result and show error toasts
- handle_set_value() validates model changes through provider manager
- handle_interactive_selection() checks validation results
- Updated unit tests to verify whitelist behavior
diff --git a/src/cortex-common/src/model_presets/mod.rs b/src/cortex-common/src/model_presets/mod.rs
@@ -18,7 +18,7 @@ pub use types::{ModelAlias, ModelPreset, ModelResolution};
 pub use constants::{DEFAULT_MODEL, DEFAULT_MODELS, DEFAULT_PROVIDER};
 
 // Re-export preset data and helpers
-pub use presets::{MODEL_PRESETS, get_model_preset, get_models_for_provider, validate_chutes_model, provider_allows_custom_models};
+pub use presets::{MODEL_PRESETS, get_model_preset, get_models_for_provider, validate_chutes_model, provider_allows_custom_models, CHUTES_ALLOWED_MODELS};
 
 // Re-export alias data and helpers
 pub use aliases::{MODEL_ALIASES, list_model_aliases, resolve_model_alias};
diff --git a/src/cortex-common/src/model_presets/presets.rs b/src/cortex-common/src/model_presets/presets.rs
@@ -2,6 +2,10 @@
 
 use super::types::ModelPreset;
 
+/// Allowed models for Chutes provider (whitelist approach for security).
+/// Only these models can be used with the Chutes provider.
+pub const CHUTES_ALLOWED_MODELS: &[&str] = &["moonshotai/Kimi-K2.5-TEE"];
+
 /// Available model presets.
 pub const MODEL_PRESETS: &[ModelPreset] = &[
     ModelPreset {
@@ -837,21 +841,41 @@ pub fn get_models_for_provider(provider: &str) -> Vec<&'static ModelPreset> {
 /// Returns Ok(()) if valid, Err with message if invalid.
 pub fn validate_chutes_model(model: &str) -> Result<(), String> {
     let model = model.trim();
+    
+    // Check for empty model
+    if model.is_empty() {
+        return Err("Model name cannot be empty for Chutes provider".to_string());
+    }
+    
+    // Check suffix (case-insensitive)
     if !model.to_uppercase().ends_with("-TEE") {
         return Err(format!(
             "Chutes provider only allows TEE models (models ending with '-TEE'). \
-             Model '{}' is not a TEE model. Available TEE models: moonshotai/Kimi-K2.5-TEE",
-            model
+             Model '{}' is not a TEE model. Available TEE models: {}",
+            model,
+            CHUTES_ALLOWED_MODELS.join(", ")
         ));
     }
+    
+    // SECURITY: Also verify model is in the allowed whitelist
+    if !CHUTES_ALLOWED_MODELS.iter().any(|&allowed| allowed.eq_ignore_ascii_case(model)) {
+        return Err(format!(
+            "Model '{}' is not in the allowed Chutes models list. \
+             Available models: {}",
+            model,
+            CHUTES_ALLOWED_MODELS.join(", ")
+        ));
+    }
+    
     Ok(())
 }
 
 /// Checks if a provider restricts custom models.
 /// Chutes only allows predefined TEE models, no custom models.
 pub fn provider_allows_custom_models(provider: &str) -> bool {
     // Chutes does NOT allow custom models - only predefined TEE models
-    provider != "chutes"
+    // Use case-insensitive comparison
+    !provider.eq_ignore_ascii_case("chutes")
 }
 
 #[cfg(test)]
@@ -860,31 +884,52 @@ mod tests {
 
     #[test]
     fn test_validate_chutes_model_valid() {
-        // Valid TEE models
+        // Valid TEE models in whitelist
         assert!(validate_chutes_model("moonshotai/Kimi-K2.5-TEE").is_ok());
-        assert!(validate_chutes_model("some-model-TEE").is_ok());
         // Case insensitive
-        assert!(validate_chutes_model("model-tee").is_ok());
-        assert!(validate_chutes_model("model-Tee").is_ok());
+        assert!(validate_chutes_model("moonshotai/kimi-k2.5-tee").is_ok());
+        assert!(validate_chutes_model("MOONSHOTAI/KIMI-K2.5-TEE").is_ok());
         // Whitespace handling
-        assert!(validate_chutes_model("  model-TEE  ").is_ok());
+        assert!(validate_chutes_model("  moonshotai/Kimi-K2.5-TEE  ").is_ok());
     }
 
     #[test]
     fn test_validate_chutes_model_invalid() {
+        // Not a TEE model
         assert!(validate_chutes_model("gpt-4").is_err());
         assert!(validate_chutes_model("claude-3").is_err());
+        
+        // Has TEE suffix but NOT in whitelist
+        assert!(validate_chutes_model("fake-model-TEE").is_err());
+        assert!(validate_chutes_model("some-model-TEE").is_err());
+        
+        // TEE in wrong position
         assert!(validate_chutes_model("model-TEE-v2").is_err());
-        assert!(validate_chutes_model("").is_err());
+        
+        // Empty string
+        let result = validate_chutes_model("");
+        assert!(result.is_err());
+        assert!(result.unwrap_err().contains("cannot be empty"));
     }
 
     #[test]
     fn test_provider_allows_custom_models() {
-        // Chutes does NOT allow custom models
+        // Chutes does NOT allow custom models - case insensitive
         assert!(!provider_allows_custom_models("chutes"));
+        assert!(!provider_allows_custom_models("Chutes"));
+        assert!(!provider_allows_custom_models("CHUTES"));
+        
         // Other providers allow custom models
         assert!(provider_allows_custom_models("cortex"));
         assert!(provider_allows_custom_models("openai"));
         assert!(provider_allows_custom_models("anthropic"));
     }
+    
+    #[test]
+    fn test_chutes_allowed_models_list() {
+        // Verify the whitelist contains expected model
+        assert!(CHUTES_ALLOWED_MODELS.contains(&"moonshotai/Kimi-K2.5-TEE"));
+        // Whitelist should not be empty
+        assert!(!CHUTES_ALLOWED_MODELS.is_empty());
+    }
 }
diff --git a/src/cortex-tui/src/providers/manager.rs b/src/cortex-tui/src/providers/manager.rs
@@ -11,7 +11,7 @@ use cortex_engine::client::{
 
 use super::config::CortexConfig;
 use super::models::{ModelInfo, get_models_for_provider, get_popular_models};
-use cortex_common::model_presets::validate_chutes_model;
+use cortex_common::model_presets::{validate_chutes_model, CHUTES_ALLOWED_MODELS};
 
 // ============================================================
 // PROVIDER MANAGER
@@ -40,7 +40,22 @@ impl ProviderManager {
     pub fn new(config: CortexConfig) -> Self {
         let (current_provider, current_model) =
             if let Some((provider, model)) = config.get_last_model() {
-                (provider.to_string(), model.to_string())
+                // SECURITY: Validate model when loading for Chutes provider
+                if provider.eq_ignore_ascii_case("chutes") {
+                    if validate_chutes_model(model).is_err() {
+                        // Fall back to first allowed Chutes model if invalid
+                        tracing::warn!(
+                            "Invalid TEE model '{}' loaded from config for Chutes provider, \
+                             resetting to default",
+                            model
+                        );
+                        (provider.to_string(), CHUTES_ALLOWED_MODELS[0].to_string())
+                    } else {
+                        (provider.to_string(), model.to_string())
+                    }
+                } else {
+                    (provider.to_string(), model.to_string())
+                }
             } else {
                 (
                     config.default_provider.clone(),
@@ -212,6 +227,19 @@ impl ProviderManager {
     /// Sets the current provider.
     pub fn set_provider(&mut self, provider: &str) -> Result<()> {
         self.current_provider = provider.to_string();
+        
+        // SECURITY: Validate current model when switching to Chutes provider
+        if provider.eq_ignore_ascii_case("chutes") {
+            if validate_chutes_model(&self.current_model).is_err() {
+                tracing::warn!(
+                    "Current model '{}' is not a valid TEE model for Chutes provider, \
+                     switching to default",
+                    self.current_model
+                );
+                self.current_model = CHUTES_ALLOWED_MODELS[0].to_string();
+            }
+        }
+        
         self.client = None;
         Ok(())
     }
@@ -221,7 +249,8 @@ impl ProviderManager {
         let resolved = self.config.resolve_alias(model);
 
         // Validate model for Chutes provider (TEE-only security requirement)
-        if self.current_provider == "chutes" {
+        // Use case-insensitive comparison for provider check
+        if self.current_provider.eq_ignore_ascii_case("chutes") {
             validate_chutes_model(&resolved)
                 .map_err(|e| anyhow::anyhow!(e))?;
         }
@@ -369,6 +398,13 @@ impl ProviderManager {
         if self.client.is_some() {
             return Ok(());
         }
+        
+        // SECURITY: Defense-in-depth - validate TEE model before API calls
+        if self.current_provider.eq_ignore_ascii_case("chutes") {
+            validate_chutes_model(&self.current_model)
+                .map_err(|e| anyhow::anyhow!("{}", e))?;
+        }
+        
         let token = self.get_token()?;
         self.client = Some(create_client("cortex", &self.current_model, &token, None)?);
         Ok(())
diff --git a/src/cortex-tui/src/runner/event_loop/commands.rs b/src/cortex-tui/src/runner/event_loop/commands.rs
@@ -548,22 +548,56 @@ impl EventLoop {
     pub(super) fn handle_set_value(&mut self, key: &str, value: &str) {
         match key {
             "model" => {
-                self.app_state.model = value.to_string();
-                if let Some(pm) = &self.provider_manager
+                // Check validation result first, storing any error
+                let validation_result: Result<(), String> = if let Some(pm) = &self.provider_manager
                     && let Ok(mut manager) = pm.try_write()
                 {
-                    let _ = manager.set_model(value);
+                    manager.set_model(value).map_err(|e| e.to_string())
+                } else {
+                    Ok(())
+                };
+                
+                // Handle validation result after releasing the borrow
+                if let Err(e) = validation_result {
+                    self.add_system_message(&format!("Cannot set model: {}", e));
+                    return;
                 }
+                
+                // Only update state if validation passed
+                self.app_state.model = value.to_string();
                 self.update_session_model(value);
                 self.add_system_message(&format!("Model set to: {}", value));
             }
             "provider" => {
-                self.app_state.provider = value.to_string();
-                if let Some(pm) = &self.provider_manager
+                // Check validation result first, storing any error and new model
+                let validation_result: Result<Option<String>, String> = if let Some(pm) = &self.provider_manager
                     && let Ok(mut manager) = pm.try_write()
                 {
-                    let _ = manager.set_provider(value);
+                    match manager.set_provider(value) {
+                        Ok(()) => {
+                            // Get updated model after provider switch
+                            Ok(Some(manager.current_model().to_string()))
+                        }
+                        Err(e) => Err(e.to_string()),
+                    }
+                } else {
+                    Ok(None)
+                };
+                
+                // Handle validation result after releasing the borrow
+                match validation_result {
+                    Err(e) => {
+                        self.add_system_message(&format!("Cannot set provider: {}", e));
+                        return;
+                    }
+                    Ok(Some(new_model)) => {
+                        // Update model to reflect any changes made during provider switch
+                        self.app_state.model = new_model;
+                    }
+                    Ok(None) => {}
                 }
+                
+                self.app_state.provider = value.to_string();
                 self.add_system_message(&format!("Provider set to: {}", value));
             }
             "session_name" => {
diff --git a/src/cortex-tui/src/runner/event_loop/modal.rs b/src/cortex-tui/src/runner/event_loop/modal.rs
@@ -37,14 +37,21 @@ impl EventLoop {
                     KeyCode::Enter => {
                         if let Some(model) = self.app_state.model_picker.selected_model() {
                             let model_id = model.id.clone();
-                            self.app_state.model = model_id.clone();
 
+                            // Validate through manager instead of direct assignment
                             if let Some(pm) = &self.provider_manager
                                 && let Ok(mut manager) = pm.try_write()
                             {
-                                let _ = manager.set_model(&model_id);
+                                if let Err(e) = manager.set_model(&model_id) {
+                                    self.app_state.toasts.error(format!("Cannot use model: {}", e));
+                                    self.app_state.close_modal();
+                                    return Ok(true);
+                                }
                             }
 
+                            // Only update if validation passed
+                            self.app_state.model = model_id.clone();
+
                             if let Ok(mut config) = crate::providers::config::CortexConfig::load() {
                                 let _ = config.save_last_model(&self.app_state.provider, &model_id);
                             }
@@ -216,11 +223,17 @@ impl EventLoop {
     pub(super) async fn process_modal_action(&mut self, action: ModalAction) {
         match action {
             ModalAction::SelectModel(model_id) => {
-                self.app_state.model = model_id.clone();
+                // Don't set model directly - let set_model handle validation
                 if let Some(pm) = &self.provider_manager
                     && let Ok(mut manager) = pm.try_write()
                 {
-                    let _ = manager.set_model(&model_id);
+                    // Check result of set_model instead of ignoring it
+                    if let Err(e) = manager.set_model(&model_id) {
+                        self.app_state.toasts.error(format!("Cannot use model: {}", e));
+                        return;
+                    }
+                    // Only update app state if validation passed
+                    self.app_state.model = model_id.clone();
                 }
                 if let Ok(mut config) = crate::providers::config::CortexConfig::load() {
                     let _ = config.save_last_model(&self.app_state.provider, &model_id);
@@ -498,7 +511,12 @@ impl EventLoop {
                 if let Some(pm) = &self.provider_manager
                     && let Ok(mut manager) = pm.try_write()
                 {
-                    let _ = manager.set_provider(&item_id);
+                    if let Err(e) = manager.set_provider(&item_id) {
+                        self.app_state.toasts.error(format!("Cannot switch provider: {}", e));
+                        return false;
+                    }
+                    // Update model to reflect any changes made during provider switch
+                    self.app_state.model = manager.current_model().to_string();
                     self.app_state.provider = item_id.clone();
                 }
                 return false;
@@ -507,7 +525,11 @@ impl EventLoop {
                 if let Some(pm) = &self.provider_manager
                     && let Ok(mut manager) = pm.try_write()
                 {
-                    let _ = manager.set_model(&item_id);
+                    // Check validation result
+                    if let Err(e) = manager.set_model(&item_id) {
+                        self.app_state.toasts.error(format!("Cannot use model: {}", e));
+                        return false;
+                    }
                     self.app_state.model = item_id.clone();
                 }
                 self.update_session_model(&item_id);
