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