Rust AI CLI 安全加固：把用户输入清洗、模型输出校验做成一整套

Rust AI CLI 安全加固把用户输入清洗、模型输出校验做成一整套一、AI CLI 的安全边界到底在哪做 AI CLI 工具和做普通 CLI 最大的区别在于你的程序既要面对人输入的不可信数据也要面对模型输出的不可信数据。这两层不确定性叠加安全风险是乘法的。我画个数据流图把这道题讲清楚flowchart LR A[ 用户输入] --|1. 原始输入| B[ 输入清洗层] B --|2. 清洗后数据| C[ AI 模型调用] C --|3. 模型原始输出| D[ 输出校验层] D --|4. 校验后输出| E[ 命令行执行/展示] B --|被拒绝| F[❌ 拒绝并提示] D --|被拒绝| F style A fill:#ff6b6b,stroke:#333 style C fill:#ffd93d,stroke:#333 style E fill:#6bcb77,stroke:#333输入侧的风险包括注入攻击命令注入、Prompt 注入、超长输入导致 OOM、特殊字符导致解析异常。输出侧的风险包括模型幻觉产生危险命令、JSON 格式错误导致下游崩溃、输出里夹杂恶意内容。Rust 在这个场景下的最大优势不是不会被注入而是——你可以用类型系统把未清洗的数据和已清洗的数据分成两个互不兼容的类型编译器会盯着你确保你绝不会把脏数据直接喂给下游。二、输入清洗用 Newtype 模式强制类型安全的我最早的理解是加个if判断就行了。后来发现不行——项目一复杂你根本不知道哪个函数拿到的数据是洗过的、哪个是没洗过的。Rust 的Newtype 模式可以彻底解决这个问题// 输入清洗模块 /// 原始用户输入 —— 不可信任不能直接使用 pub struct RawUserInput(String); /// 经过清洗的安全输入 —— 已脱敏、截断、校验 pub struct SanitizedInput { content: String, token_count: usize, // 估算的 token 数量 } /// 清洗规则配置 pub struct SanitizeConfig { pub max_chars: usize, // 最大字符数超过就截断 pub max_lines: usize, // 最大行数 pub deny_patterns: VecString, // 禁止匹配的正则 } impl RawUserInput { /// 唯一对外暴露的安全转换入口 pub fn sanitize(self, config: SanitizeConfig) - ResultSanitizedInput, String { let mut content self.0; // 1️⃣ 截断超长输入防止 OOM if content.chars().count() config.max_chars { content content.chars().take(config.max_chars).collect(); eprintln!(⚠️ 输入过长已自动截断到 {} 字符, config.max_chars); } // 2️⃣ 行数限制 let lines: Vecstr content.lines().collect(); if lines.len() config.max_lines { return Err(format!( 输入超过最大行数限制 {}当前 {} 行, config.max_lines, lines.len() )); } // 3️⃣ 白名单 黑名单模式只保留安全字符 // 移除控制字符除了换行和制表符 content content .chars() .filter(|c| c.is_alphanumeric() || c.is_whitespace() || .,!?;:()-_/#$%^*[]{}|~\.contains(*c)) .collect(); // 4️⃣ 正则黑名单检测 for pattern in config.deny_patterns { if regex::Regex::new(pattern) .map_err(|e| format!(正则错误: {e}))? .is_match(content) { return Err(format!(输入包含禁止模式: {pattern})); } } // 5️⃣ 粗略 token 估算英文按空格分中文按字符分 let token_count content .split_whitespace() .map(|w| { if w.chars().any(|c| c as u32 0x4e00) { w.chars().count() // 中文字符大约 1 token/字 } else { w.len().div_ceil(4) // 英文约 4 字符 1 token } }) .sum(); Ok(SanitizedInput { content, token_count, }) } } // 关键设计 // SanitizedInput 没有公开的构造函数只能通过 RawUserInput::sanitize 创建 // 这意味着任何持有 SanitizedInput 的地方数据一定已经洗过了 impl SanitizedInput { /// 获取内部字符串的只读引用 pub fn as_str(self) - str { self.content } /// 估计的 token 数量 pub fn token_count(self) - usize { self.token_count } }这个设计的精髓在于下游函数只接受SanitizedInput而不是String。如果你不小心传了一个没洗过的String进去编译器直接报错不是运行时崩是编译就不让你过。// ❌ 这样写会编译失败 —— 下游只要 SanitizedInput fn call_ai_model(input: SanitizedInput) - String { // ... 模型调用 } let raw RawUserInput(rm -rf /.to_string()); // call_ai_model(raw); // 编译错误: expected SanitizedInput, found RawUserInput三、模型输出校验Schema 驱动的自动验证模型输出比用户输入更危险因为它看起来可信——JSON 结构完整、语气温和、格式漂亮。但幻觉内容能让你的系统出大事。我的方案是预先定义输出结构体用 serde 反序列化自动校验通不过的直接拒绝。use serde::{Deserialize, Serialize}; // 输出校验模块 /// AI 输出的命令 —— 必须满足这个结构才允许执行 #[derive(Debug, Deserialize)] pub struct AiCommandOutput { /// 命令类型必须在允许列表中 pub command_type: AllowedCommand, /// 要执行的命令字符串 pub command: String, /// 简短说明限制长度防止注入 #[serde(default)] pub description: String, } /// 允许的命令类型枚举 —— 之外的一律拒绝 #[derive(Debug, Deserialize, PartialEq)] #[serde(rename_all lowercase)] pub enum AllowedCommand { /// 文件操作 File, /// Git 操作 Git, /// Docker 操作 Docker, /// 系统信息查询只读 SysInfo, } /// 输出校验器 pub struct OutputValidator { /// 危险命令黑名单即使类型正确也拒绝 forbidden_patterns: VecString, } impl OutputValidator { pub fn new(forbidden_patterns: VecString) - Self { Self { forbidden_patterns } } /// 校验并转换模型输出 pub fn validate(self, raw_output: str) - ResultAiCommandOutput, VecString { let mut errors Vec::new(); // 1️⃣ 尝试 JSON 解析 let parsed: AiCommandOutput match serde_json::from_str(raw_output) { Ok(v) v, Err(e) { // 模型输出可能不是纯 JSON尝试提取 JSON 块 if let Some(json_block) extract_json_block(raw_output) { match serde_json::from_str(json_block) { Ok(v) v, Err(e2) { return Err(vec![ format!(JSON 解析失败: {e}), format!(JSON 块提取也失败: {e2}), ]); } } } else { return Err(vec![format!(JSON 解析失败: {e})]); } } }; // 2️⃣ 命令类型已在反序列化时自动校验非法值直接失败 // 3️⃣ 命令内容黑名单扫描 for pattern in self.forbidden_patterns { if let Ok(re) regex::Regex::new(pattern) { if re.is_match(parsed.command) { errors.push(format!( 命令包含危险模式: {pattern}模型输出被拒绝 )); } } } // 4️⃣ 描述长度限制 if parsed.description.len() 200 { errors.push(描述字段超长.to_string()); } if errors.is_empty() { Ok(parsed) } else { Err(errors) } } } /// 从模型输出中提取 JSON 代码块 fn extract_json_block(text: str) - OptionString { let start text.find(json)?; let content_start start 7; let end text[content_start..].find()?; Some(text[content_start..content_start end].trim().to_string()) }这里有个细节AllowedCommand是个枚举serde 在反序列化时如果遇到不认识的值会直接报错。不需要写 if-else类型系统替你做完了。四、组合起来SafeAiCli 把输入清洗和输出校验串成管道单独的清洗和校验还不够真正落地需要把它们串成一个不可绕过的管道flowchart TD Start([用户输入]) -- Raw[RawUserInput 包裹] Raw -- San[调用 sanitize 方法] San --|Err| Reject1[返回错误提示] San --|Ok: SanitizedInput| Build[构建 Prompt] Build -- Call[调用 AI 模型 API] Call -- RawOut[获取原始模型输出] RawOut -- Val[OutputValidator::validate] Val --|Err: Vec String| Reject2[记录日志 返回友好错误] Val --|Ok: AiCommandOutput| Check[安全检查: 确认命令可执行] Check --|安全| Exec[执行命令] Check --|危险| Reject2 style San fill:#ffd93d,stroke:#333 style Val fill:#ffd93d,stroke:#333 style Exec fill:#6bcb77,stroke:#333 style Reject1 fill:#ff6b6b,stroke:#333 style Reject2 fill:#ff6b6b,stroke:#333对应的 SafeAiCli 结构体实现pub struct SafeAiCli { sanitize_config: SanitizeConfig, output_validator: OutputValidator, // 这里放你实际调 AI 的客户端 // ai_client: Boxdyn AiClient, } impl SafeAiCli { pub fn new() - Self { Self { sanitize_config: SanitizeConfig { max_chars: 4000, max_lines: 50, deny_patterns: vec![ rrm\s-rf\s/.to_string(), // 毁灭性删除 r\s*/dev/sda.to_string(), // 写入磁盘设备 rcurl.*\|.*sh.to_string(), // 管道执行远程脚本 ], }, output_validator: OutputValidator::new(vec![ rrm\s-rf\s/.to_string(), rsudo\s.to_string(), // 提权操作一律禁止 rchmod\s777.to_string(), ]), } } /// 整个流程的单一入口 pub async fn ask_and_execute(self, raw: String) - ResultString, String { // 第一步清洗输入 let input RawUserInput(raw) .sanitize(self.sanitize_config) .map_err(|e| format!(输入校验失败: {e}))?; println!(✅ 输入清洗通过估算 {} tokens, input.token_count()); // 第二步调用 AI这里简化实际替换成你的 AI 调用 let prompt format!( r#你是一个安全的命令行助手。 请严格按照以下 JSON 格式回复 {{command_type: file|git|docker|sysinfo, command: 具体的命令, description: 简短说明}} 用户请求: {}#, input.as_str() ); let model_output call_ai_api(prompt).await .map_err(|e| format!(AI 调用失败: {e}))?; // 第三步校验输出 let command self.output_validator .validate(model_output) .map_err(|errors| { format!(输出校验失败:\n{}, errors.join(\n)) })?; println!(✅ 输出校验通过: {:?}, command.command_type); // 第四步用户二次确认血的教训永远别自动执行 println!(即将执行: {}, command.command); println!(说明: {}, command.description); Ok(format!(命令已确认: {}, command.command)) } } /// 模拟 AI API 调用实际替换成你的客户端 async fn call_ai_api(prompt: str) - ResultString, String { // 实际项目中会用 reqwest 调用 OpenAI / Anthropic 等 API // 这里用模拟数据占位 Ok(format!( r#json {{command_type: file, command: ls -la, description: 列出当前目录所有文件}} # )) }我在项目里就踩过这个坑上线第一周没加输出校验模型偶尔在 JSON 里塞了command: rm -rf /这种幻觉命令幸好我在执行前加了一层人工确认用户没点同意。五、总结这篇文章我把做 AI CLI 时的一套安全管道整理了出来核心是三件事Newtype 模式隔离脏数据和干净数据——不用靠命名约定或注释来标记安全性编译器帮你盯着。RawUserInput只能通过sanitize生成SanitizedInput传错类型编译不过。serde 枚举做模型输出校验——用#[derive(Deserialize)]自动校验 Schema用枚举限制允许的命令类型。非法类型在反序列化阶段就被拦住了不需要手写 if-else 判断。串成不可绕过的管道——输入→清洗→调用→校验→确认→执行中间任何一步失败都有清晰的错误提示不会悄悄吞掉异常。作为自学编程的开发者我最大的感受是Rust 的类型系统是最好的安全文档。它不会帮你写出完美的安全代码但它会确保——如果你写错了它不让你编译通过。这种编译期安全比任何代码审查都靠谱。如果你也在写 AI 工具希望这篇文章对你有帮助。有任何想法欢迎在评论区交流我还在学习中请多指教 。