Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add Windows path updates for uv tool #5029

Merged
merged 1 commit into from
Jul 13, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

1 change: 1 addition & 0 deletions Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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" }
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Already in our dependency tree. I'd slightly rather use winapi if possible but rustup uses winreg and it's just a pain to translate it.

wiremock = { version = "0.6.0" }
zip = { version = "0.6.6", default-features = false, features = ["deflate"] }

Expand Down
5 changes: 5 additions & 0 deletions crates/uv-shell/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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 }
6 changes: 5 additions & 1 deletion crates/uv-shell/src/lib.rs
Original file line number Diff line number Diff line change
@@ -1,3 +1,5 @@
pub mod windows;

use std::path::{Path, PathBuf};
use uv_fs::Simplified;

Expand Down Expand Up @@ -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![],
}
}
Expand Down
128 changes: 128 additions & 0 deletions crates/uv-shell/src/windows.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,128 @@
//! Windows-specific utilities for manipulating the environment.
//!
//! Based on rustup's Windows implementation: <https://github.com/rust-lang/rustup/blob/fede22fea7b160868cece632bd213e6d72f8912f/src/cli/self_update/windows.rs>

#![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<bool> {
// 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<u16>) -> 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", &reg_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<Option<Vec<u16>>> {
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(&reg_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<u16>, path: Vec<u16>) -> Option<Vec<u16>> {
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<u16>) -> Vec<u8> {
value.push(0);
#[allow(unsafe_code)]
unsafe {
slice::from_raw_parts(value.as_ptr().cast::<u8>(), 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<Vec<u16>> {
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::<u16>(), val.bytes.len() / 2)
.to_owned()
};
while words.last() == Some(&0) {
words.pop();
}
Some(words)
}
_ => None,
}
}
22 changes: 22 additions & 0 deletions crates/uv/src/commands/tool/update_shell.rs
Original file line number Diff line number Diff line change
@@ -1,3 +1,5 @@
#![cfg_attr(windows, allow(unreachable_code))]

use std::fmt::Write;

use anyhow::Result;
Expand Down Expand Up @@ -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(),
Expand Down
Loading