返回 CodeWhale
shell_denial_acceptance.rs
根目录 / crates / tui / tests / integration / shell_denial_acceptance.rs
1 //! Real CLI/Engine denial boundary with a loopback-only mock provider.
2 #![cfg(unix)]
3
4 use serde_json::{Value, json};
5 use std::io::Read;
6 use std::process::{Command, Stdio};
7 use std::sync::{
8 Arc,
9 atomic::{AtomicUsize, Ordering},
10 };
11 use std::time::Duration;
12 use tempfile::TempDir;
13 use wait_timeout::ChildExt;
14 use wiremock::matchers::{method, path};
15 use wiremock::{Mock, MockServer, Request, Respond, ResponseTemplate};
16
17 const MODEL: &str = "shell-denial-fixture";
18 const COMMAND: &str = "printf fixture > denial-canary.txt";
19
20 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
21 async fn denied_bash_cannot_be_reached_through_task_search_and_start() {
22 scenario("task_shell_start", true).await;
23 }
24
25 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
26 async fn denied_bash_cannot_be_reached_through_tasks_gate_action() {
27 scenario("tasks", true).await;
28 }
29
30 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
31 async fn unrestricted_task_search_and_start_still_executes() {
32 scenario("task_shell_start", false).await;
33 }
34
35 #[derive(Clone)]
36 struct Script {
37 calls: Arc<Vec<(&'static str, &'static str, Value)>>,
38 count: Arc<AtomicUsize>,
39 }
40 impl Respond for Script {
41 fn respond(&self, request: &Request) -> ResponseTemplate {
42 assert_eq!(request.body_json::<Value>().unwrap()["model"], MODEL);
43 let sequence = self.count.fetch_add(1, Ordering::SeqCst);
44 let (delta, finish) = match self.calls.get(sequence) {
45 Some((id, tool, input)) => (
46 json!({"tool_calls":[{"index":0,"id":id,"type":"function","function":{"name":tool,"arguments":input.to_string()}}]}),
47 "tool_calls",
48 ),
49 None => (json!({"content":"PERMISSION_FIXTURE_FINISHED"}), "stop"),
50 };
51 let mut response = String::new();
52 for (delta, finish) in [(delta, None), (json!({}), Some(finish))] {
53 let chunk = json!({"id":"fixture","model":MODEL,"choices":[{"index":0,"delta":delta,"finish_reason":finish}]});
54 response.push_str(&format!("data: {chunk}\n\n"));
55 }
56 response.push_str("data: [DONE]\n\n");
57 ResponseTemplate::new(200)
58 .insert_header("content-type", "text/event-stream")
59 .set_body_string(response)
60 }
61 }
62
63 async fn scenario(tool: &'static str, deny: bool) {
64 let workspace = TempDir::new().unwrap();
65 let home = TempDir::new().unwrap();
66 let server = MockServer::start().await;
67 let mut calls = Vec::new();
68 if deny {
69 calls.push(("direct", "Bash", json!({"command":COMMAND})));
70 }
71 calls.push((
72 "search",
73 "tool_search",
74 json!({"query":tool, "max_results":8}),
75 ));
76 let input = if tool == "tasks" {
77 json!({"action":"gate_run", "gate":"custom", "command":COMMAND})
78 } else {
79 json!({"command":COMMAND})
80 };
81 calls.push(("route", tool, input.clone()));
82 // Repeat the attempted call even when catalog shaping hid it. A guessed
83 // name, a cached schema, and deferred hydration must not grant execution.
84 calls.push(("route-retry", tool, input));
85 let expected_calls = calls.len() + 1;
86 let count = Arc::new(AtomicUsize::new(0));
87 Mock::given(method("POST"))
88 .and(path("/v1/chat/completions"))
89 .respond_with(Script {
90 calls: Arc::new(calls),
91 count: count.clone(),
92 })
93 .mount(&server)
94 .await;
95 let config = home.path().join(".codewhale/config.toml");
96 std::fs::create_dir_all(config.parent().unwrap()).unwrap();
97 std::fs::write(
98 &config,
99 "allow_shell = true\ntelemetry = false\n[retry]\nenabled = false\n",
100 )
101 .unwrap();
102 let mut command = Command::new(env!("CARGO_BIN_EXE_codewhale-tui"));
103 command.env_clear();
104 for key in ["PATH", "LANG", "TMPDIR"] {
105 if let Some(value) = std::env::var_os(key) {
106 command.env(key, value);
107 }
108 }
109 command
110 .current_dir(workspace.path())
111 .arg("--workspace")
112 .arg(workspace.path())
113 .args([
114 "--no-project-config",
115 "exec",
116 "--auto",
117 "--provider",
118 "deepseek",
119 "--model",
120 MODEL,
121 "--output-format",
122 "stream-json",
123 ])
124 .env("HOME", home.path())
125 .env("CODEWHALE_HOME", config.parent().unwrap())
126 .env("CODEWHALE_CONFIG_PATH", &config)
127 .env("DEEPSEEK_API_KEY", "fixture-key-not-real")
128 .env("DEEPSEEK_BASE_URL", server.uri())
129 .env("CODEWHALE_BASE_URL", server.uri())
130 .env("DEEPSEEK_MODEL", MODEL)
131 .env("CODEWHALE_MODEL", MODEL)
132 .env("CODEWHALE_TELEMETRY", "0")
133 .env("RUST_LOG", "warn")
134 .stdout(Stdio::piped())
135 .stderr(Stdio::piped());
136 if deny {
137 command.args(["--disallowed-tools", "Bash"]);
138 }
139 command.arg("Execute the local permission fixture.");
140 let mut child = command.spawn().unwrap();
141 let drain = |mut pipe: Box<dyn Read + Send>| {
142 std::thread::spawn(move || {
143 let mut bytes = Vec::new();
144 pipe.read_to_end(&mut bytes).unwrap();
145 bytes
146 })
147 };
148 let stdout = drain(Box::new(child.stdout.take().unwrap()));
149 let stderr = drain(Box::new(child.stderr.take().unwrap()));
150 let status = child
151 .wait_timeout(Duration::from_secs(60))
152 .unwrap()
153 .unwrap_or_else(|| {
154 child.kill().ok();
155 child.wait().ok();
156 panic!("permission fixture timed out")
157 });
158 let stdout = String::from_utf8_lossy(&stdout.join().unwrap()).into_owned();
159 let stderr = String::from_utf8_lossy(&stderr.join().unwrap()).into_owned();
160 assert!(status.success(), "{stdout}\n{stderr}");
161 assert_eq!(count.load(Ordering::SeqCst), expected_calls);
162 assert_eq!(
163 workspace.path().join("denial-canary.txt").exists(),
164 !deny,
165 "{stdout}\n{stderr}"
166 );
167 let requests = server.received_requests().await.unwrap();
168 let final_request: Value = requests.last().unwrap().body_json().unwrap();
169 let results: Vec<_> = final_request["messages"]
170 .as_array()
171 .unwrap()
172 .iter()
173 .filter(|message| message["role"] == "tool")
174 .collect();
175 if deny {
176 for id in ["direct", "route", "route-retry"] {
177 let receipt = results
178 .iter()
179 .find(|message| message["tool_call_id"] == id)
180 .unwrap();
181 assert!(
182 receipt["content"]
183 .as_str()
184 .unwrap()
185 .contains("disallowed-tools"),
186 "{receipt}"
187 );
188 }
189 } else {
190 assert!(
191 results
192 .iter()
193 .any(|message| message["tool_call_id"] == "route-retry"
194 && message["content"].as_str().unwrap().contains("task_id"))
195 );
196 }
197 }
198
198 lines RUST