Skip to main content

vtcode_core/
zsh_exec_bridge.rs

1use anyhow::Result;
2#[cfg(not(unix))]
3use anyhow::anyhow;
4#[cfg(not(unix))]
5use hashbrown::HashMap;
6#[cfg(not(unix))]
7use std::path::Path;
8
9#[cfg(unix)]
10mod unix_impl {
11    #[derive(Debug, Clone, Serialize, Deserialize)]
12    struct WrapperExecRequest {
13        request_id: String,
14        file: String,
15        argv: Vec<String>,
16        cwd: String,
17    }
18
19    #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
20    #[serde(rename_all = "snake_case")]
21    enum WrapperExecAction {
22        Allow,
23        Deny,
24    }
25
26    #[derive(Debug, Clone, Serialize, Deserialize)]
27    struct WrapperExecResponse {
28        request_id: String,
29        action: WrapperExecAction,
30        reason: Option<String>,
31    }
32
33    const ZSH_EXEC_BRIDGE_WRAPPER_SOCKET_ENV_VAR: &str = "VTCODE_ZSH_EXEC_BRIDGE_WRAPPER_SOCKET";
34    const ZSH_EXEC_WRAPPER_MODE_ENV_VAR: &str = "VTCODE_ZSH_EXEC_WRAPPER_MODE";
35    const EXEC_WRAPPER_ENV_VAR: &str = "EXEC_WRAPPER";
36    use anyhow::{Context, Result, bail};
37    use hashbrown::HashMap;
38    use parking_lot::Mutex;
39    use serde::{Deserialize, Serialize};
40    use std::fs;
41    use std::io::{ErrorKind, Read, Write};
42    use std::os::unix::fs::PermissionsExt;
43    use std::os::unix::net::{UnixListener, UnixStream};
44    use std::path::{Path, PathBuf};
45    use std::sync::{
46        Arc,
47        atomic::{AtomicBool, Ordering},
48    };
49    use std::thread::{self, JoinHandle};
50    use std::time::Duration;
51    use tracing::warn;
52    use uuid::Uuid;
53
54    const ACCEPT_POLL_INTERVAL: Duration = Duration::from_millis(20);
55
56    pub(crate) struct ZshExecBridgeSession {
57        socket_path: PathBuf,
58        stop: Arc<AtomicBool>,
59        worker: Mutex<Option<JoinHandle<()>>>,
60    }
61
62    impl ZshExecBridgeSession {
63        pub(crate) fn spawn(allow_confirmed_dangerous: bool) -> Result<Self> {
64            let socket_path = std::env::temp_dir().join(format!("vtcode-zsh-exec-bridge-{}.sock", Uuid::new_v4()));
65
66            if socket_path.exists() {
67                fs::remove_file(&socket_path)
68                    .with_context(|| format!("remove pre-existing zsh bridge socket at {}", socket_path.display()))?;
69            }
70
71            let listener = UnixListener::bind(&socket_path)
72                .with_context(|| format!("bind zsh exec bridge socket listener at {}", socket_path.display()))?;
73            // Restrict socket to owner-only — prevents other users on the same
74            // machine from communicating with the bridge (defence in depth;
75            // the random UUID path already provides unpredictability).
76            fs::set_permissions(&socket_path, fs::Permissions::from_mode(0o700))
77                .with_context(|| format!("set permissions on zsh exec bridge socket at {}", socket_path.display()))?;
78            listener
79                .set_nonblocking(true)
80                .context("set zsh exec bridge listener to nonblocking")?;
81
82            let stop = Arc::new(AtomicBool::new(false));
83            let stop_clone = Arc::clone(&stop);
84            let cleanup_path = socket_path.clone();
85            let worker = thread::Builder::new()
86                .name("vtcode-zsh-exec-bridge".to_string())
87                .spawn(move || {
88                    run_bridge_loop(listener, stop_clone, allow_confirmed_dangerous);
89                    let _ = fs::remove_file(&cleanup_path);
90                })
91                .context("spawn zsh exec bridge listener thread")?;
92
93            Ok(Self {
94                socket_path,
95                stop,
96                worker: Mutex::new(Some(worker)),
97            })
98        }
99
100        pub(crate) fn env_vars(&self, wrapper_executable: &Path) -> HashMap<String, String> {
101            HashMap::from([
102                (ZSH_EXEC_BRIDGE_WRAPPER_SOCKET_ENV_VAR.to_string(), self.socket_path.to_string_lossy().to_string()),
103                (ZSH_EXEC_WRAPPER_MODE_ENV_VAR.to_string(), "1".to_string()),
104                (EXEC_WRAPPER_ENV_VAR.to_string(), wrapper_executable.to_string_lossy().to_string()),
105            ])
106        }
107    }
108
109    impl Drop for ZshExecBridgeSession {
110        fn drop(&mut self) {
111            self.stop.store(true, Ordering::Relaxed);
112            if let Some(worker) = self.worker.lock().take()
113                && worker.join().is_err()
114            {
115                warn!("zsh exec bridge worker thread panicked during cleanup");
116            }
117            let _ = fs::remove_file(&self.socket_path);
118        }
119    }
120
121    fn run_bridge_loop(listener: UnixListener, stop: Arc<AtomicBool>, allow_confirmed_dangerous: bool) {
122        while !stop.load(Ordering::Relaxed) {
123            match listener.accept() {
124                Ok((mut stream, _)) => {
125                    if let Err(err) = handle_wrapper_request(&mut stream, allow_confirmed_dangerous) {
126                        warn!(error = %err, "zsh exec bridge request failed");
127                    }
128                }
129                Err(err) if err.kind() == ErrorKind::WouldBlock => {
130                    thread::sleep(ACCEPT_POLL_INTERVAL);
131                }
132                Err(err) => {
133                    warn!(error = %err, "zsh exec bridge listener failed");
134                    break;
135                }
136            }
137        }
138    }
139
140    fn handle_wrapper_request(stream: &mut UnixStream, allow_confirmed_dangerous: bool) -> Result<()> {
141        let mut payload = String::new();
142        stream.read_to_string(&mut payload).context("read wrapper request payload")?;
143        let request: WrapperExecRequest =
144            serde_json::from_str(payload.trim()).context("parse wrapper request payload")?;
145
146        let (action, reason) = evaluate_wrapper_exec_request(&request, allow_confirmed_dangerous);
147        let response = WrapperExecResponse {
148            request_id: request.request_id.clone(),
149            action,
150            reason,
151        };
152        let encoded = serde_json::to_string(&response).context("serialize wrapper response")?;
153        stream.write_all(encoded.as_bytes()).context("write wrapper response payload")?;
154        stream.write_all(b"\n").context("write wrapper response newline")?;
155        stream.flush().context("flush wrapper response")?;
156        Ok(())
157    }
158
159    fn evaluate_wrapper_exec_request(
160        request: &WrapperExecRequest,
161        allow_confirmed_dangerous: bool,
162    ) -> (WrapperExecAction, Option<String>) {
163        let command = if request.argv.is_empty() {
164            vec![request.file.clone()]
165        } else {
166            request.argv.clone()
167        };
168
169        if command.is_empty() {
170            return (WrapperExecAction::Deny, Some("Rejected empty wrapped command".to_string()));
171        }
172
173        if allow_confirmed_dangerous {
174            return (WrapperExecAction::Allow, None);
175        }
176
177        let display = shell_words::join(command.iter().map(String::as_str));
178        if let Err(err) = crate::tools::validation::commands::validate_command_safety(&display) {
179            return (WrapperExecAction::Deny, Some(format!("Rejected by command safety validation: {err}")));
180        }
181        if crate::command_safety::command_might_be_dangerous(&command) {
182            return (WrapperExecAction::Deny, Some("Rejected dangerous subcommand".to_string()));
183        }
184
185        (WrapperExecAction::Allow, None)
186    }
187
188    pub(crate) fn maybe_run_zsh_exec_wrapper_mode() -> Result<bool> {
189        let wrapper_mode = std::env::var(ZSH_EXEC_WRAPPER_MODE_ENV_VAR).ok();
190        if wrapper_mode.as_deref() != Some("1") {
191            return Ok(false);
192        }
193
194        run_zsh_exec_wrapper_mode()?;
195        Ok(true)
196    }
197
198    fn run_zsh_exec_wrapper_mode() -> Result<()> {
199        let args: Vec<String> = std::env::args().collect();
200        if args.len() < 2 {
201            bail!("zsh exec wrapper mode requires target executable path");
202        }
203
204        let file = args[1].clone();
205        let argv = if args.len() > 2 {
206            args[2..].to_vec()
207        } else {
208            vec![file.clone()]
209        };
210        let cwd = std::env::current_dir()
211            .context("resolve wrapper cwd")?
212            .to_string_lossy()
213            .to_string();
214        let socket_path =
215            std::env::var(ZSH_EXEC_BRIDGE_WRAPPER_SOCKET_ENV_VAR).context("missing wrapper socket path env var")?;
216
217        let request_id = Uuid::new_v4().to_string();
218        let request = WrapperExecRequest {
219            request_id: request_id.clone(),
220            file: file.clone(),
221            argv: argv.clone(),
222            cwd,
223        };
224
225        let mut stream =
226            UnixStream::connect(&socket_path).with_context(|| format!("connect to wrapper socket at {socket_path}"))?;
227        let encoded = serde_json::to_string(&request).context("serialize wrapper request")?;
228        stream.write_all(encoded.as_bytes()).context("write wrapper request payload")?;
229        stream.write_all(b"\n").context("write wrapper request newline")?;
230        stream
231            .shutdown(std::net::Shutdown::Write)
232            .context("shutdown wrapper request writer")?;
233
234        let mut response_buf = String::new();
235        stream
236            .read_to_string(&mut response_buf)
237            .context("read wrapper response payload")?;
238        let response: WrapperExecResponse =
239            serde_json::from_str(response_buf.trim()).context("parse wrapper response payload")?;
240
241        if response.request_id != request_id {
242            bail!("wrapper response request_id mismatch: expected {request_id}, got {}", response.request_id);
243        }
244
245        if response.action == WrapperExecAction::Deny {
246            if let Some(reason) = response.reason {
247                warn!("zsh exec bridge denied execution: {reason}");
248            } else {
249                warn!("zsh exec bridge denied execution");
250            }
251            std::process::exit(1);
252        }
253
254        let mut command = std::process::Command::new(&file);
255        if argv.len() > 1 {
256            command.args(&argv[1..]);
257        }
258        command.env_remove(ZSH_EXEC_WRAPPER_MODE_ENV_VAR);
259        command.env_remove(ZSH_EXEC_BRIDGE_WRAPPER_SOCKET_ENV_VAR);
260        command.env_remove(EXEC_WRAPPER_ENV_VAR);
261        let status = command.status().context("spawn wrapped executable")?;
262        std::process::exit(status.code().unwrap_or(1));
263    }
264
265    #[cfg(test)]
266    mod tests {
267        use super::{WrapperExecAction, WrapperExecRequest, evaluate_wrapper_exec_request};
268
269        fn request(command: &[&str]) -> WrapperExecRequest {
270            let file = command.first().unwrap_or(&"/usr/bin/true").to_string();
271            WrapperExecRequest {
272                request_id: "test-request".to_string(),
273                file: file.clone(),
274                argv: command.iter().map(|s| s.to_string()).collect(),
275                cwd: "/tmp".to_string(),
276            }
277        }
278
279        #[test]
280        fn evaluate_request_denies_dangerous_when_unconfirmed() {
281            let request = request(&["rm", "-rf", "/tmp/demo"]);
282            let (action, reason) = evaluate_wrapper_exec_request(&request, false);
283            assert_eq!(action, WrapperExecAction::Deny);
284            assert!(reason.is_some());
285        }
286
287        #[test]
288        fn evaluate_request_allows_safe_when_unconfirmed() {
289            let request = request(&["/usr/bin/true"]);
290            let (action, reason) = evaluate_wrapper_exec_request(&request, false);
291            assert_eq!(action, WrapperExecAction::Allow);
292            assert!(reason.is_none());
293        }
294
295        #[test]
296        fn evaluate_request_allows_dangerous_when_confirmed() {
297            let request = request(&["rm", "-rf", "/tmp/demo"]);
298            let (action, reason) = evaluate_wrapper_exec_request(&request, true);
299            assert_eq!(action, WrapperExecAction::Allow);
300            assert!(reason.is_none());
301        }
302    }
303}
304
305#[cfg(unix)]
306pub(crate) use unix_impl::ZshExecBridgeSession;
307
308#[cfg(unix)]
309pub fn maybe_run_zsh_exec_wrapper_mode() -> Result<bool> {
310    unix_impl::maybe_run_zsh_exec_wrapper_mode()
311}
312
313#[cfg(not(unix))]
314pub(crate) struct ZshExecBridgeSession;
315
316#[cfg(not(unix))]
317impl ZshExecBridgeSession {
318    pub(crate) fn spawn(_allow_confirmed_dangerous: bool) -> Result<Self> {
319        Err(anyhow!("zsh exec bridge is only supported on Unix platforms"))
320    }
321
322    pub(crate) fn env_vars(&self, _wrapper_executable: &Path) -> HashMap<String, String> {
323        HashMap::new()
324    }
325}
326
327#[cfg(not(unix))]
328pub fn maybe_run_zsh_exec_wrapper_mode() -> Result<bool> {
329    Ok(false)
330}