Skip to main content

upstream_rs/services/integration/
shell_manager.rs

1use anyhow::{Context, Result};
2use std::fs;
3use std::path::Path;
4use std::sync::{Mutex, OnceLock};
5
6/// Process-global lock used to serialize reads/writes to the shared PATH file.
7fn paths_file_lock() -> &'static Mutex<()> {
8    static LOCK: OnceLock<Mutex<()>> = OnceLock::new();
9    LOCK.get_or_init(|| Mutex::new(()))
10}
11
12#[cfg(windows)]
13fn normalize_windows_path(path: &str) -> String {
14    let mut normalized = path.replace('/', "\\").trim().to_ascii_lowercase();
15    while normalized.ends_with('\\') {
16        normalized.pop();
17    }
18    normalized
19}
20
21pub struct ShellManager<'a> {
22    paths_file: &'a Path,
23}
24
25impl<'a> ShellManager<'a> {
26    pub fn new(paths_file: &'a Path) -> Self {
27        Self { paths_file }
28    }
29
30    /// Adds a package's installation path to PATH
31    pub fn add_to_paths(&self, install_path: &Path) -> Result<()> {
32        if !install_path.is_dir() {
33            anyhow::bail!(
34                "Package install directory not found: {}",
35                install_path.to_string_lossy()
36            );
37        }
38
39        #[cfg(unix)]
40        {
41            let _guard = paths_file_lock()
42                .lock()
43                .ok()
44                .ok_or_else(|| anyhow::anyhow!("Failed to lock PATH file for writing"))?;
45            let mut content =
46                fs::read_to_string(self.paths_file).context("Failed to read paths file")?;
47            let escaped = install_path
48                .to_string_lossy()
49                .replace('$', "\\$")
50                .replace('"', "\\\"");
51            let export_line = format!("export PATH=\"{escaped}:$PATH\"");
52
53            if !content.contains(&export_line) {
54                content.push_str(&format!("{export_line}\n"));
55                fs::write(self.paths_file, &content).context("Failed to write paths file")?;
56            }
57        }
58
59        #[cfg(windows)]
60        {
61            self.add_to_windows_registry(install_path)?;
62        }
63
64        Ok(())
65    }
66
67    /// Removes a package's PATH entry
68    pub fn remove_from_paths(&self, install_path: &Path) -> Result<()> {
69        #[cfg(unix)]
70        {
71            let _guard = paths_file_lock()
72                .lock()
73                .ok()
74                .ok_or_else(|| anyhow::anyhow!("Failed to lock PATH file for writing"))?;
75            let mut content =
76                fs::read_to_string(self.paths_file).context("Failed to read paths file")?;
77            let escaped = install_path
78                .to_string_lossy()
79                .replace('$', "\\$")
80                .replace('"', "\\\"");
81            let export_line = format!("export PATH=\"{escaped}:$PATH\"");
82
83            content = content.replace(&format!("{export_line}\n"), "");
84            content = content.replace(&export_line, "");
85            fs::write(self.paths_file, content).context("Failed to write paths file")?;
86        }
87
88        #[cfg(windows)]
89        {
90            self.remove_from_windows_registry(install_path)?;
91        }
92
93        Ok(())
94    }
95
96    #[cfg(windows)]
97    fn add_to_windows_registry(&self, path: &Path) -> Result<()> {
98        use winreg::RegKey;
99        use winreg::enums::*;
100
101        let hkcu = RegKey::predef(HKEY_CURRENT_USER);
102        let env_key = hkcu
103            .open_subkey_with_flags("Environment", KEY_READ | KEY_WRITE)
104            .context("Failed to open registry Environment key")?;
105
106        let target_path = path.to_string_lossy().to_string();
107        let target_norm = normalize_windows_path(&target_path);
108
109        // Get current PATH
110        let current_path: String = env_key.get_value("Path").unwrap_or_else(|_| String::new());
111
112        // Check if path is already in PATH
113        let path_entries: Vec<&str> = current_path.split(';').collect();
114        if path_entries
115            .iter()
116            .any(|&p| normalize_windows_path(p) == target_norm)
117        {
118            return Ok(()); // Already in PATH
119        }
120
121        // Add path to the beginning
122        let new_path = if current_path.is_empty() {
123            target_path
124        } else {
125            format!("{};{}", target_path, current_path)
126        };
127
128        env_key
129            .set_value("Path", &new_path)
130            .context("Failed to set PATH in registry")?;
131
132        // Broadcast environment change
133        Self::broadcast_environment_change();
134
135        Ok(())
136    }
137
138    #[cfg(windows)]
139    fn remove_from_windows_registry(&self, path: &Path) -> Result<()> {
140        use winreg::RegKey;
141        use winreg::enums::*;
142
143        let hkcu = RegKey::predef(HKEY_CURRENT_USER);
144        let env_key = hkcu
145            .open_subkey_with_flags("Environment", KEY_READ | KEY_WRITE)
146            .context("Failed to open registry Environment key")?;
147
148        let target_path = path.to_string_lossy().to_string();
149        let target_norm = normalize_windows_path(&target_path);
150
151        // Get current PATH
152        let current_path: String = env_key.get_value("Path").unwrap_or_else(|_| String::new());
153
154        // Remove target path from PATH
155        let path_entries: Vec<&str> = current_path
156            .split(';')
157            .filter(|&p| normalize_windows_path(p) != target_norm)
158            .collect();
159
160        let new_path = path_entries.join(";");
161
162        env_key
163            .set_value("Path", &new_path)
164            .context("Failed to set PATH in registry")?;
165
166        // Broadcast environment change
167        Self::broadcast_environment_change();
168
169        Ok(())
170    }
171
172    #[cfg(windows)]
173    fn broadcast_environment_change() {
174        use std::ptr;
175        use winapi::shared::minwindef::LPARAM;
176        use winapi::um::winuser::{
177            HWND_BROADCAST, SMTO_ABORTIFHUNG, SendMessageTimeoutW, WM_SETTINGCHANGE,
178        };
179
180        unsafe {
181            let env_string: Vec<u16> = "Environment\0".encode_utf16().collect();
182            SendMessageTimeoutW(
183                HWND_BROADCAST,
184                WM_SETTINGCHANGE,
185                0,
186                env_string.as_ptr() as LPARAM,
187                SMTO_ABORTIFHUNG,
188                5000,
189                ptr::null_mut(),
190            );
191        }
192    }
193}
194
195#[cfg(test)]
196mod tests {
197    use super::ShellManager;
198    use std::path::{Path, PathBuf};
199    use std::time::{SystemTime, UNIX_EPOCH};
200    use std::{fs, io};
201
202    fn temp_root(name: &str) -> PathBuf {
203        let nanos = SystemTime::now()
204            .duration_since(UNIX_EPOCH)
205            .map(|d| d.as_nanos())
206            .unwrap_or(0);
207        std::env::temp_dir().join(format!("upstream-shell-test-{name}-{nanos}"))
208    }
209
210    fn cleanup(path: &Path) -> io::Result<()> {
211        fs::remove_dir_all(path)
212    }
213
214    #[cfg(unix)]
215    #[test]
216    fn add_to_paths_is_idempotent_and_escapes_special_characters() {
217        let root = temp_root("add-idempotent");
218        let install_path = root.join("tool\"dir$");
219        let paths_file = root.join("paths.sh");
220        fs::create_dir_all(&install_path).expect("create install dir");
221        fs::write(&paths_file, "#!/usr/bin/env sh\n").expect("create paths file");
222        let manager = ShellManager::new(&paths_file);
223
224        manager.add_to_paths(&install_path).expect("first add");
225        manager.add_to_paths(&install_path).expect("second add");
226
227        let content = fs::read_to_string(&paths_file).expect("read paths file");
228        assert_eq!(content.matches("export PATH=").count(), 1);
229        assert!(content.contains("\\\""));
230        assert!(content.contains("\\$"));
231
232        cleanup(&root).expect("cleanup");
233    }
234
235    #[cfg(unix)]
236    #[test]
237    fn remove_from_paths_removes_existing_export_line() {
238        let root = temp_root("remove");
239        let install_path = root.join("pkg/bin");
240        let paths_file = root.join("paths.sh");
241        fs::create_dir_all(&install_path).expect("create install dir");
242        fs::write(&paths_file, "").expect("create paths file");
243        let manager = ShellManager::new(&paths_file);
244
245        manager.add_to_paths(&install_path).expect("add path");
246        manager
247            .remove_from_paths(&install_path)
248            .expect("remove path");
249
250        let content = fs::read_to_string(&paths_file).expect("read paths file");
251        assert!(!content.contains("export PATH="));
252
253        cleanup(&root).expect("cleanup");
254    }
255}