From cff3dc430927783f4195dfcace66ec1558b7ec05 Mon Sep 17 00:00:00 2001 From: Charlie Marsh Date: Fri, 12 Jul 2024 20:56:53 -0400 Subject: [PATCH] Add Windows path updates for uv tool --- Cargo.lock | 3 + Cargo.toml | 1 + crates/uv-shell/Cargo.toml | 5 + crates/uv-shell/src/lib.rs | 6 +- crates/uv-shell/src/windows.rs | 128 ++++++++++++++++++++ crates/uv/src/commands/tool/update_shell.rs | 22 ++++ 6 files changed, 164 insertions(+), 1 deletion(-) create mode 100644 crates/uv-shell/src/windows.rs diff --git a/Cargo.lock b/Cargo.lock index 9ef483377ca6..8132dcb6bce2 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5089,9 +5089,12 @@ dependencies = [ name = "uv-shell" version = "0.0.1" dependencies = [ + "anyhow", "home", "same-file", + "tracing", "uv-fs", + "winreg", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index 129e4ee1e867..accd7f8d0c2f 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -151,6 +151,7 @@ urlencoding = { version = "2.1.3" } walkdir = { version = "2.5.0" } which = { version = "6.0.0" } winapi = { version = "0.3.9", features = ["fileapi", "handleapi", "ioapiset", "winbase", "winioctl", "winnt"] } +winreg = { version = "0.52.0" } wiremock = { version = "0.6.0" } zip = { version = "0.6.6", default-features = false, features = ["deflate"] } diff --git a/crates/uv-shell/Cargo.toml b/crates/uv-shell/Cargo.toml index 4cee54ce71f9..062bf74e8aa9 100644 --- a/crates/uv-shell/Cargo.toml +++ b/crates/uv-shell/Cargo.toml @@ -10,5 +10,10 @@ workspace = true [dependencies] uv-fs = { workspace = true } +anyhow = { workspace = true } home = { workspace = true } same-file = { workspace = true } +tracing = { workspace = true } + +[target.'cfg(windows)'.dependencies] +winreg = { workspace = true } diff --git a/crates/uv-shell/src/lib.rs b/crates/uv-shell/src/lib.rs index b547c9271bc8..06a05517d057 100644 --- a/crates/uv-shell/src/lib.rs +++ b/crates/uv-shell/src/lib.rs @@ -1,3 +1,5 @@ +pub mod windows; + use std::path::{Path, PathBuf}; use uv_fs::Simplified; @@ -150,9 +152,11 @@ impl Shell { // On Csh, we need to update both `.cshrc` and `.login`, like Bash. vec![home_dir.join(".cshrc"), home_dir.join(".login")] } - // TODO(charlie): Add support for Nushell, PowerShell, and Cmd. + // TODO(charlie): Add support for Nushell. Shell::Nushell => vec![], + // See: [`crate::windows::prepend_path`]. Shell::Powershell => vec![], + // See: [`crate::windows::prepend_path`]. Shell::Cmd => vec![], } } diff --git a/crates/uv-shell/src/windows.rs b/crates/uv-shell/src/windows.rs new file mode 100644 index 000000000000..8e6e6b124abf --- /dev/null +++ b/crates/uv-shell/src/windows.rs @@ -0,0 +1,128 @@ +//! Windows-specific utilities for manipulating the environment. +//! +//! Based on rustup's Windows implementation: + +#![cfg(windows)] + +use std::ffi::OsString; +use std::io; +use std::os::windows::ffi::OsStrExt; +use std::path::Path; +use std::slice; + +use anyhow::Context; +use winreg::enums::{RegType, HKEY_CURRENT_USER, KEY_READ, KEY_WRITE}; +use winreg::{RegKey, RegValue}; + +/// Append the given [`Path`] to the `PATH` environment variable in the Windows registry. +/// +/// Returns `Ok(true)` if the path was successfully appended, and `Ok(false)` if the path was +/// already in `PATH`. +pub fn prepend_path(path: &Path) -> anyhow::Result { + // Get the existing `PATH` variable from the registry. + let windows_path = get_windows_path_var()?; + + // Add the new path to the existing `PATH` variable. + let windows_path = windows_path.and_then(|windows_path| { + prepend_to_path(windows_path, OsString::from(path).encode_wide().collect()) + }); + + // If the path didn't change, then we don't need to do anything. + let Some(windows_path) = windows_path else { + return Ok(false); + }; + + // Set the `PATH` variable in the registry. + apply_windows_path_var(windows_path)?; + + Ok(true) +} + +/// Set the windows `PATH` variable in the registry. +fn apply_windows_path_var(path: Vec) -> anyhow::Result<()> { + let root = RegKey::predef(HKEY_CURRENT_USER); + let environment = root.open_subkey_with_flags("Environment", KEY_READ | KEY_WRITE)?; + + if path.is_empty() { + environment.delete_value("PATH")?; + } else { + let reg_value = RegValue { + bytes: to_winreg_bytes(path), + vtype: RegType::REG_EXPAND_SZ, + }; + environment.set_raw_value("PATH", ®_value)?; + } + + Ok(()) +} + +/// Retrieve the windows `PATH` variable from the registry. +/// +/// Returns `Ok(None)` if the `PATH` variable is not a string. +fn get_windows_path_var() -> anyhow::Result>> { + let root = RegKey::predef(HKEY_CURRENT_USER); + let environment = root + .open_subkey_with_flags("Environment", KEY_READ | KEY_WRITE) + .context("Failed to open `Environment` key")?; + + let reg_value = environment.get_raw_value("PATH"); + match reg_value { + Ok(reg_value) => { + if let Some(reg_value) = from_winreg_value(®_value) { + Ok(Some(reg_value)) + } else { + tracing::warn!("`HKEY_CURRENT_USER\\Environment\\PATH` is a non-string"); + Ok(None) + } + } + Err(ref err) if err.kind() == io::ErrorKind::NotFound => Ok(Some(Vec::new())), + Err(err) => Err(err.into()), + } +} + +/// Prepend a path to the `PATH` variable in the Windows registry. +/// +/// Returns `Ok(None)` if the given path is already in `PATH`. +fn prepend_to_path(existing_path: Vec, path: Vec) -> Option> { + if existing_path.is_empty() { + Some(path) + } else if existing_path.windows(path.len()).any(|p| p == path) { + None + } else { + let mut new_path = path; + new_path.push(u16::from(b';')); + new_path.extend(existing_path); + Some(new_path) + } +} + +/// Convert a vector UCS-2 chars to a null-terminated UCS-2 string in bytes. +fn to_winreg_bytes(mut value: Vec) -> Vec { + value.push(0); + #[allow(unsafe_code)] + unsafe { + slice::from_raw_parts(value.as_ptr().cast::(), value.len() * 2).to_vec() + } +} + +/// Decode the `HKCU\Environment\PATH` value. +/// +/// If the key is not `REG_SZ` or `REG_EXPAND_SZ`, returns `None`. +/// The `winreg` library itself does a lossy unicode conversion. +fn from_winreg_value(val: &RegValue) -> Option> { + match val.vtype { + RegType::REG_SZ | RegType::REG_EXPAND_SZ => { + #[allow(unsafe_code)] + let mut words = unsafe { + #[allow(clippy::cast_ptr_alignment)] + slice::from_raw_parts(val.bytes.as_ptr().cast::(), val.bytes.len() / 2) + .to_owned() + }; + while words.last() == Some(&0) { + words.pop(); + } + Some(words) + } + _ => None, + } +} diff --git a/crates/uv/src/commands/tool/update_shell.rs b/crates/uv/src/commands/tool/update_shell.rs index 6f8481d61bd2..999279fc1248 100644 --- a/crates/uv/src/commands/tool/update_shell.rs +++ b/crates/uv/src/commands/tool/update_shell.rs @@ -1,3 +1,5 @@ +#![cfg_attr(windows, allow(unreachable_code))] + use std::fmt::Write; use anyhow::Result; @@ -26,6 +28,26 @@ pub(crate) async fn update_shell(preview: PreviewMode, printer: Printer) -> Resu executable_directory.simplified_display() ); + #[cfg(windows)] + { + if uv_shell::windows::prepend_path(&executable_directory)? { + writeln!( + printer.stderr(), + "Updated PATH to include executable directory {}", + executable_directory.simplified_display().cyan() + )?; + writeln!(printer.stderr(), "Restart your shell to apply changes.")?; + } else { + writeln!( + printer.stderr(), + "Executable directory {} is already in PATH", + executable_directory.simplified_display().cyan() + )?; + } + + return Ok(ExitStatus::Success); + } + if Shell::contains_path(&executable_directory) { writeln!( printer.stderr(),