diff --git a/Cargo.lock b/Cargo.lock index 79e71f5..5d127f9 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -499,6 +499,7 @@ dependencies = [ "uu_wc", "uu_yes", "uucore", + "wild", "windows-sys 0.59.0", "winresource", ] diff --git a/Cargo.toml b/Cargo.toml index 5e97da5..1937185 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -22,6 +22,8 @@ uucore = { path = "deps/coreutils/src/uucore" } uucore_procs = { path = "deps/coreutils/src/uucore_procs" } [dependencies] +# Dependency versions match those used in deps/coreutils. +wild = "2.2.1" clap = { version = "4.5", features = ["wrap_help", "cargo", "color"] } itertools = "0.14" phf = "0.13" @@ -130,8 +132,14 @@ ntfind = { package = "find", path = "deps/ntfind" } [dependencies.windows-sys] version = "*" features = [ + "Win32_Foundation", + "Win32_Security_Cryptography", + "Win32_Storage_FileSystem", "Win32_System_Console", + "Win32_System_Pipes", "Win32_System_Registry", + "Win32_UI_Shell", + "Win32_UI_WindowsAndMessaging", ] [build-dependencies] diff --git a/PSScriptAnalyzerSettings.psd1 b/PSScriptAnalyzerSettings.psd1 new file mode 100644 index 0000000..3efab24 --- /dev/null +++ b/PSScriptAnalyzerSettings.psd1 @@ -0,0 +1,15 @@ +@{ + Rules = @{ + PSUseCompatibleCommands = @{ + Enable = $true + TargetProfiles = @('win-8_x64_10.0.17763.0_5.1.17763.316_x64_4.0.30319.42000_framework') + } + PSUseCompatibleSyntax = @{ + Enable = $true + TargetVersions = @('5.1', '7.0') + } + PSAvoidOverwritingBuiltInCmdlets = @{ Enable = $false } + PSUseShouldProcessForStateChangingFunctions = @{ Enable = $false } + PSUseSingularNouns = @{ Enable = $false } + } +} diff --git a/README.md b/README.md index 2f2df72..8691ec1 100644 --- a/README.md +++ b/README.md @@ -43,6 +43,10 @@ Or grab the latest build from our [Release Page](https://github.com/microsoft/co > [!NOTE] > Any command not mentioned is included in this suite. The following only lists conflicts. +> [!NOTE] +> You can remove additional utilities using `coreutils-manager disable `. +> See `coreutils-manager --help` for other management commands. + > [!WARNING] > PowerShell 7.4 or later is required. > PowerShell 7.6 or later is recommended for `~` support. diff --git a/build.rs b/build.rs index 4752313..8879eb3 100644 --- a/build.rs +++ b/build.rs @@ -61,6 +61,11 @@ fn generate_uutils_map() { entries.push(("sort".into(), "(sort_uumain, sort_uu_app)".into())); } + entries.push(( + "coreutils-manager".into(), + "(manager::uumain, manager::uu_app)".into(), + )); + entries.sort(); let mut phf_map = phf_codegen::OrderedMap::new(); @@ -74,7 +79,7 @@ type UtilityMap = phf::OrderedMap<&'static str, (fn(T) -> i32, fn() -> Comman #[allow(clippy::too_many_lines)] #[allow(clippy::unreadable_literal)] -fn util_map() -> UtilityMap {{ +const fn util_map() -> UtilityMap {{ {} }} ", diff --git a/coreutils.iss b/coreutils.iss index 7af6385..b1373aa 100644 --- a/coreutils.iss +++ b/coreutils.iss @@ -93,25 +93,24 @@ end; procedure CreateHardlinks; var Output: TExecOutput; - Name: String; - ResultCode, I: Integer; + Detail: String; + ResultCode: Integer; begin ForceDirectories(g_AppBinDirPath); ForceDirectories(g_AppCmdDirPath); - if (not ExecAndCaptureOutput(g_CoreutilsExePath, '--list', '', SW_SHOWNORMAL, ewWaitUntilTerminated, ResultCode, Output)) or (ResultCode <> 0) then - RaiseException('Failed to execute coreutils.exe --list'); - - for I := 0 to GetArrayLength(Output.StdOut) - 1 do + if (not ExecAndCaptureOutput(g_CoreutilsExePath, 'coreutils-manager refresh', '', SW_SHOWNORMAL, ewWaitUntilTerminated, ResultCode, Output)) or (ResultCode <> 0) then begin - Name := Trim(Output.StdOut[I]); - if (Name <> '') and (Name <> '[') then - begin - if not CreateHardLink(g_AppBinDirPath + Name + '.exe', g_CoreutilsExePath, 0) then - RaiseException('Failed to create hardlink for ' + Name); - if not CreateHardLink(g_AppCmdDirPath + Name + '.cmd', g_CoreutilsExePath, 0) then - RaiseException('Failed to create hardlink for ' + Name); - end; + Detail := ''; + if GetArrayLength(Output.StdErr) > 0 then + Detail := Output.StdErr[0] + else if GetArrayLength(Output.StdOut) > 0 then + Detail := Output.StdOut[0]; + + if Detail <> '' then + RaiseException('Failed to refresh coreutils links: ' + Detail) + else + RaiseException('Failed to refresh coreutils links'); end; end; diff --git a/src/main.rs b/src/main.rs index 42355e5..caa51fb 100644 --- a/src/main.rs +++ b/src/main.rs @@ -5,6 +5,7 @@ // Microsoft-authored changes, which Microsoft makes available to uutils // under the uutils MIT License for upstream incorporation. See NOTICE.md. +mod manager; mod nthelpers; use std::borrow::Cow; @@ -16,9 +17,10 @@ use std::sync::atomic::AtomicU32; use clap::Command; use itertools::Itertools as _; +use uucore::Args; use uucore::display::Quotable as _; use uucore::windows_sys::Win32::System::Threading::GetCurrentProcess; -use uucore::{Args, error::strip_errno, locale}; +use uucore::{error::strip_errno, locale}; use windows_sys::Win32::Globalization::CP_UTF8; use windows_sys::Win32::System::Console; use windows_sys::Win32::System::Threading::TerminateProcess; @@ -28,10 +30,14 @@ unsafe extern "C" { unsafe fn ntsort_main(argc: i32, argv: *const *const u8) -> i32; } -const VERSION: &str = env!("CARGO_PKG_VERSION"); - include!(concat!(env!("OUT_DIR"), "/uutils_map.rs")); +// While uucore::Args is a trait, this is the actual type it'll resolve to in the end. +type ArgsType = std::iter::Chain, wild::ArgsOs>; + +const VERSION: &str = env!("CARGO_PKG_VERSION"); +const UTIL_MAP: UtilityMap = util_map(); + fn usage(utils: &UtilityMap, name: &str) { let display_list = utils.keys().copied().join(", "); let width = cmp::min(textwrap::termwidth(), 100) - 8; // (opinion/heuristic) max 100 chars wide with 4 character side indentions @@ -70,18 +76,19 @@ fn main() { // ntsort just hardcodes to CP_OEMCP, so it also isn't affected. set_console_modes(); - let utils = util_map(); - let mut args = uucore::args_os(); + // `wild::args_os()` is what `uucore::args_os()` uses under the hood. + // By using it directly we avoid duplicating all arg strings. + let mut args = wild::args_os(); let binary = binary_path(&mut args); let binary_as_util = name(&binary).unwrap_or_else(|| { - usage(&utils, ""); + usage(&UTIL_MAP, ""); exit(0); }); // binary name ends with util name? let is_coreutils = binary_as_util.ends_with("utils"); - let matched_util = utils + let matched_util = UTIL_MAP .keys() .filter(|&&u| binary_as_util.ends_with(u) && !is_coreutils) .max_by_key(|u| u.len()); //Prefer stty more than tty. *utils is not ls @@ -110,7 +117,7 @@ fn main() { exit(1); } let mut out = io::stdout().lock(); - for util in utils.keys() { + for util in UTIL_MAP.keys() { if let Err(e) = writeln!(out, "{util}") && e.kind() != io::ErrorKind::BrokenPipe { @@ -133,7 +140,7 @@ fn main() { _ => {} } - match utils.get(util) { + match UTIL_MAP.get(util) { Some(&(uumain, _)) => { // TODO: plug the deactivation of the translation // and load the English strings directly at compilation time in the @@ -151,7 +158,7 @@ fn main() { not_found(&util_os) }; - match utils.get(util) { + match UTIL_MAP.get(util) { Some(&(uumain, _)) => { setup_localization_or_exit(util); let code = uumain( @@ -165,7 +172,7 @@ fn main() { None => not_found(&util_os), } } - usage(&utils, binary_as_util); + usage(&UTIL_MAP, binary_as_util); exit(0); } else if util.starts_with('-') { // Argument looks like an option but wasn't recognized @@ -179,7 +186,7 @@ fn main() { // GNU just fails, but busybox tests needs usage // todo: patch the test suite instead if binary_as_util.ends_with("box") { - usage(&utils, binary_as_util); + usage(&UTIL_MAP, binary_as_util); } else { let _ = writeln!(io::stderr(), "coreutils: missing argument"); } diff --git a/src/manager.rs b/src/manager.rs new file mode 100644 index 0000000..aeb7fe1 --- /dev/null +++ b/src/manager.rs @@ -0,0 +1,609 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +// TODO: It would be cleaner to implement the entire PowerShell logic here in Rust. +// FWIW we loose 2 things this way: +// * Access to `$PROFILE.CurrentUserCurrentHost` +// * .NET's text decoding/encoding logic (including support for ACP encoding) + +use std::collections::BTreeSet; +use std::ffi::{OsStr, OsString}; +use std::fmt::Write as _; +use std::fs; +use std::io::{self, Write as _}; +use std::os::windows::ffi::OsStrExt as _; +use std::path::{Path, PathBuf}; +use std::process; +use std::ptr; + +use clap::builder::{NonEmptyStringValueParser, PathBufValueParser}; +use clap::{Arg, ArgMatches, Command}; +use uucore::Args; +use uucore::error::{FromIo as _, UError, UResult, USimpleError}; +use windows_sys::Win32::Foundation; +use windows_sys::Win32::Security::Cryptography; +use windows_sys::Win32::Storage::FileSystem; +use windows_sys::Win32::System::Console; +use windows_sys::Win32::System::Pipes; +use windows_sys::Win32::System::Registry; +use windows_sys::Win32::UI::Shell; +use windows_sys::Win32::UI::WindowsAndMessaging; +use windows_sys::w; + +const REG_PATH: *const u16 = w!(r"SOFTWARE\Microsoft\coreutils"); +const REG_DISABLED_UTILITIES: *const u16 = w!("DisabledUtilities"); + +#[uucore::main(no_signals)] +pub fn uumain(args: T) -> UResult<()> { + let matches = uucore::clap_localization::handle_clap_result_with_exit_code(uu_app(), args, 2)?; + + // Redirect our stdout/stderr into the given named pipe if needed. + // This is used for self-elevation. + let _stdout_pipe = if let Some(path) = matches.get_one::("stdout-pipe") { + Some(connect_stdout_pipe(path)?) + } else { + None + }; + + let utilities = utility_names(); + let mut disabled = read_disabled_utilities()?; + + match matches.subcommand() { + Some((action @ ("enable" | "disable"), matches)) => { + let requested: BTreeSet = matches + .get_many::("utilities") + .expect("utility is a required parameter") + .cloned() + .collect(); + + // Update the disabled set + let remove = action == "enable"; + let mut unknown = Vec::new(); + for name in &requested { + if !utilities.contains(name.as_str()) { + unknown.push(name.clone()); + continue; + } + if remove { + disabled.remove(name); + } else { + disabled.insert(name.clone()); + } + } + if !unknown.is_empty() { + return Err(USimpleError::new( + 1, + format!("unknown coreutils utility: {}", unknown.join(", ")), + )); + } + + if !ensure_elevated(matches, action, &requested)? { + return Ok(()); + } + + write_disabled_utilities(&disabled)?; + sync_install("refresh", &utilities, &disabled) + } + Some(("refresh", _)) => sync_install("", &utilities, &disabled), + Some(("status", _)) => { + for utility_name in utilities { + let status = if disabled.contains(utility_name) { + "disabled" + } else { + "enabled" + }; + println!("{utility_name:16}{status}"); + } + Ok(()) + } + _ => unreachable!("clap enforces a known subcommand"), + } +} + +pub fn uu_app() -> Command { + Command::new("coreutils-manager") + .version(env!("CARGO_PKG_VERSION")) + .about("Manage coreutils utilities and PowerShell profiles") + .arg( + Arg::new("stdout-pipe") + .long("stdout-pipe") + .hide(true) + .global(true) + .value_parser(PathBufValueParser::new()), + ) + .arg( + Arg::new("no-elevate") + .long("no-elevate") + .hide(true) + .global(true) + .action(clap::ArgAction::SetTrue), + ) + .subcommand_required(true) + .subcommand( + Command::new("enable") + .about("Enable one or more utilities") + .arg( + Arg::new("utilities") + .help("Utility names to enable") + .num_args(1..) + .required(true) + .trailing_var_arg(true) + .value_parser(NonEmptyStringValueParser::new()), + ), + ) + .subcommand( + Command::new("disable") + .about("Disable one or more utilities") + .arg( + Arg::new("utilities") + .help("Utility names to disable") + .num_args(1..) + .required(true) + .trailing_var_arg(true) + .value_parser(NonEmptyStringValueParser::new()), + ), + ) + .subcommand(Command::new("refresh").hide(true)) + .subcommand(Command::new("status").about("List all utilities with their status")) +} + +fn utility_names() -> BTreeSet<&'static str> { + super::UTIL_MAP + .keys() + .copied() + .filter(|&n| n != "[" && n != "coreutils-manager") + .collect() +} + +fn connect_stdout_pipe(path: &Path) -> UResult { + let path = wide_null(path); + let handle = unsafe { + FileSystem::CreateFileW( + path.as_ptr(), + FileSystem::FILE_GENERIC_WRITE, + 0, + ptr::null(), + FileSystem::OPEN_EXISTING, + FileSystem::FILE_ATTRIBUTE_NORMAL, + ptr::null_mut(), + ) + }; + if handle == Foundation::INVALID_HANDLE_VALUE { + return Err(last_os_error("failed to connect stdout pipe")); + } + + if unsafe { Console::SetStdHandle(Console::STD_OUTPUT_HANDLE, handle) } == 0 + || unsafe { Console::SetStdHandle(Console::STD_ERROR_HANDLE, handle) } == 0 + { + unsafe { + Foundation::CloseHandle(handle); + } + return Err(last_os_error("failed to redirect stdout/stderr")); + } + + Ok(OwnedHandle(handle)) +} + +fn ensure_elevated( + matches: &ArgMatches, + action: &str, + utilities: &BTreeSet, +) -> UResult { + if unsafe { Shell::IsUserAnAdmin() != 0 } { + Ok(true) + } else { + if matches.get_flag("no-elevate") { + return Err(USimpleError::new( + 1, + "administrator privileges are required".to_string(), + )); + } + elevate(action, utilities)?; + Ok(false) + } +} + +fn elevate(action: &str, utilities: &BTreeSet) -> UResult<()> { + let (pipe_path, pipe) = NamedPipe::with_random_pipe_path("coreutils-manager")?; + let exe = std::env::current_exe()?; + + let mut command_line = OsString::new(); + if !exe + .file_stem() + .is_some_and(|stem| stem == "coreutils-manager") + { + command_line.push("coreutils-manager "); + } + _ = write!(command_line, "{action} --no-elevate --stdout-pipe "); + command_line.push(pipe_path); + for name in utilities { + _ = write!(command_line, " {name}"); + } + + let exe = wide_null(&exe); + let parameters = wide_null(&command_line); + let result = unsafe { + Shell::ShellExecuteW( + ptr::null_mut(), + w!("runas"), + exe.as_ptr(), + parameters.as_ptr(), + ptr::null(), + WindowsAndMessaging::SW_HIDE, + ) + }; + if result as usize <= 32 { + return Err(last_os_error("failed to elevate coreutils-manager")); + } + + pipe.connect()?; + pipe.copy_to_stdout() +} + +fn sync_install( + pwsh_install_action: &str, + utility_names: &BTreeSet<&str>, + disabled: &BTreeSet, +) -> UResult<()> { + let app_dir = app_dir()?; + let bin_dir = app_dir.join("bin"); + let cmd_dir = app_dir.join("cmd"); + let coreutils_exe = app_dir.join("coreutils.exe"); + + // Synchronize hardlinks + { + sync_hardlinks(utility_names, disabled, &coreutils_exe, &bin_dir, "exe")?; + sync_hardlinks(utility_names, disabled, &coreutils_exe, &cmd_dir, "cmd")?; + sync_hardlink( + "coreutils-manager", + disabled, + &coreutils_exe, + &bin_dir, + "exe", + )?; + } + + // Refresh PowerShell profiles + if !pwsh_install_action.is_empty() { + let script = app_dir.join("pwsh-install.ps1"); + if !script.is_file() { + return Ok(()); + } + + let status = match process::Command::new("pwsh.exe") + .arg("-NoProfile") + .arg("-NonInteractive") + .arg("-ExecutionPolicy") + .arg("Bypass") + .arg("-File") + .arg(&script) + .arg("-Action") + .arg(pwsh_install_action) + .arg("-CmdDir") + .arg(cmd_dir) + .status() + { + Ok(status) => status, + Err(err) if err.kind() == io::ErrorKind::NotFound => return Ok(()), + Err(err) => { + return Err(USimpleError::new( + 1, + format!("failed to start pwsh.exe: {err}"), + )); + } + }; + if !status.success() { + return Err(USimpleError::new( + 1, + format!("failed to refresh PowerShell profiles: {status}"), + )); + } + } + + Ok(()) +} + +fn app_dir() -> UResult { + let exe = std::env::current_exe()?; + // Navigate from... + // * C:\Program Files\coreutils\bin\utility.exe + // * C:\Program Files\coreutils\bin (check if the suffix is bin/cmd) + // * C:\Program Files\coreutils + // Or from... + // * C:\Program Files\coreutils\corutils.exe + // * C:\Program Files\coreutils + exe.parent() + .and_then(|p| { + if matches!(p.file_name(), Some(name) if name == "bin" || name == "cmd") { + p.parent() + } else { + Some(p) + } + }) + .map(Path::to_path_buf) + .ok_or_else(|| { + USimpleError::new( + 1, + format!("cannot determine install directory from {}", exe.display()), + ) + }) +} + +fn sync_hardlinks( + utility_names: &BTreeSet<&str>, + disabled: &BTreeSet, + coreutils_path: &Path, + dest_dir: &Path, + suffix: &str, +) -> UResult<()> { + fs::create_dir_all(dest_dir)?; + for &utility_name in utility_names { + sync_hardlink(utility_name, disabled, coreutils_path, dest_dir, suffix)?; + } + Ok(()) +} + +fn sync_hardlink( + utility_name: &str, + disabled: &BTreeSet, + coreutils_path: &Path, + dest_dir: &Path, + suffix: &str, +) -> UResult<()> { + let link = dest_dir.join(format!("{utility_name}.{suffix}")); + let (res, kind_ok) = if disabled.contains(utility_name) { + (fs::remove_file(&link), io::ErrorKind::NotFound) + } else { + ( + fs::hard_link(coreutils_path, &link), + io::ErrorKind::AlreadyExists, + ) + }; + if let Err(err) = res + && err.kind() != kind_ok + { + return Err(err.map_err_context(|| "failed to synchronize utility links".to_string())); + } + Ok(()) +} + +fn read_disabled_utilities() -> UResult> { + OwnedHKEY::get_multi_sz_value( + Registry::HKEY_LOCAL_MACHINE, + REG_PATH, + REG_DISABLED_UTILITIES, + ) + .map(|data| parse_multi_sz(&data)) +} + +fn write_disabled_utilities(disabled: &BTreeSet) -> UResult<()> { + let key = OwnedHKEY::create_key(Registry::HKEY_LOCAL_MACHINE, REG_PATH)?; + if disabled.is_empty() { + let ret = unsafe { Registry::RegDeleteValueW(key.get(), REG_DISABLED_UTILITIES) }; + if ret != 0 && ret != Foundation::ERROR_FILE_NOT_FOUND { + return Err(last_os_error( + "failed to clear disabled aliases from registry", + )); + } + Ok(()) + } else { + let data = make_multi_sz(disabled.iter()); + reg_check_result("failed to write disabled aliases to registry", unsafe { + Registry::RegSetValueExW( + key.get(), + REG_DISABLED_UTILITIES, + 0, + Registry::REG_MULTI_SZ, + data.as_ptr().cast(), + (data.len() * size_of::()) as u32, + ) + })?; + Ok(()) + } +} + +fn parse_multi_sz(data: &[u16]) -> BTreeSet { + data.split(|&ch| ch == 0) + .filter(|&slice| !slice.is_empty()) + .map(String::from_utf16_lossy) + .collect() +} + +fn make_multi_sz(values: impl Iterator>) -> Vec { + let mut data = Vec::new(); + for value in values { + data.extend(value.as_ref().encode_wide()); + data.push(0); + } + data.push(0); + data +} + +struct OwnedHandle(Foundation::HANDLE); + +impl OwnedHandle { + fn get(&self) -> Foundation::HANDLE { + self.0 + } +} + +impl Drop for OwnedHandle { + fn drop(&mut self) { + if !self.0.is_null() && self.0 != Foundation::INVALID_HANDLE_VALUE { + unsafe { Foundation::CloseHandle(self.0) }; + } + } +} + +struct OwnedHKEY(Registry::HKEY); + +impl OwnedHKEY { + fn create_key(key: Registry::HKEY, subkey: *const u16) -> UResult { + let mut result = ptr::null_mut(); + reg_check_result("failed to create registry key", unsafe { + Registry::RegCreateKeyW(key, subkey, &mut result) + })?; + Ok(Self(result)) + } + + fn get_multi_sz_value( + key: Registry::HKEY, + subkey: *const u16, + value: *const u16, + ) -> UResult> { + let mut bytes = 0u32; + let ret = unsafe { + Registry::RegGetValueW( + key, + subkey, + value, + Registry::RRF_RT_REG_MULTI_SZ, + ptr::null_mut(), + ptr::null_mut(), + &mut bytes, + ) + }; + if ret == Foundation::ERROR_FILE_NOT_FOUND || ret == Foundation::ERROR_PATH_NOT_FOUND { + return Ok(Vec::new()); + } + reg_check_result("failed to read from registry", ret)?; + if bytes == 0 { + return Ok(Vec::new()); + } + + let mut data: Vec = Vec::with_capacity(bytes as usize / 2 + 128); + bytes = data.capacity() as u32 * 2; + + reg_check_result("failed to read from registry", unsafe { + Registry::RegGetValueW( + key, + subkey, + value, + Registry::RRF_RT_REG_MULTI_SZ, + ptr::null_mut(), + data.as_mut_ptr().cast(), + &mut bytes, + ) + })?; + + unsafe { data.set_len(bytes as usize / 2) }; + Ok(data) + } + + fn get(&self) -> Registry::HKEY { + self.0 + } +} + +impl Drop for OwnedHKEY { + fn drop(&mut self) { + if !self.0.is_null() { + unsafe { Registry::RegCloseKey(self.0) }; + } + } +} + +struct NamedPipe(OwnedHandle); + +impl NamedPipe { + fn with_random_pipe_path(prefix: &str) -> UResult<(PathBuf, Self)> { + let mut random = [0u8; 16]; + unsafe { Cryptography::ProcessPrng(random.as_mut_ptr(), random.len()) }; + + let mut path = String::with_capacity(random.len() * 2 + prefix.len() + 10); + _ = write!(path, "\\\\.\\pipe\\{prefix}-"); + for byte in random { + _ = write!(path, "{byte:02x}"); + } + + let path = PathBuf::from(path); + let pipe = Self::create(&path)?; + Ok((path, pipe)) + } + + fn create(path: &Path) -> UResult { + let path = wide_null(path); + let handle = unsafe { + Pipes::CreateNamedPipeW( + path.as_ptr(), + FileSystem::PIPE_ACCESS_INBOUND, + Pipes::PIPE_TYPE_BYTE | Pipes::PIPE_WAIT, + 1, + 4 * 1024, + 4 * 1024, + 0, + ptr::null(), + ) + }; + if handle == Foundation::INVALID_HANDLE_VALUE { + return Err(last_os_error("failed to create stdout pipe")); + } + Ok(Self(OwnedHandle(handle))) + } + + fn connect(&self) -> UResult<()> { + if unsafe { Pipes::ConnectNamedPipe(self.0.get(), ptr::null_mut()) } != 0 { + return Ok(()); + } + + let error = unsafe { Foundation::GetLastError() }; + if error == Foundation::ERROR_PIPE_CONNECTED { + return Ok(()); + } + + Err(last_os_error("failed to connect stdout pipe")) + } + + fn copy_to_stdout(&self) -> UResult<()> { + let mut stdout = io::stdout().lock(); + let mut buffer = [0u8; 4 * 1024]; + + loop { + let mut read = 0u32; + if unsafe { + FileSystem::ReadFile( + self.0.get(), + buffer.as_mut_ptr(), + buffer.len() as u32, + &mut read, + ptr::null_mut(), + ) + } == 0 + { + let error = unsafe { Foundation::GetLastError() }; + if error == Foundation::ERROR_BROKEN_PIPE { + _ = stdout.flush(); + return Ok(()); + } + return Err(last_os_error("failed to read from stdout pipe")); + } + + if read == 0 { + return Ok(()); + } + + if stdout.write_all(&buffer[..read as usize]).is_err() { + return Ok(()); + } + } + } +} + +fn wide_null(value: impl AsRef) -> Vec { + fn encode_wide(value: &OsStr) -> Vec { + value.encode_wide().chain(std::iter::once(0)).collect() + } + encode_wide(value.as_ref()) +} + +fn last_os_error(context: &str) -> Box { + io::Error::last_os_error().map_err_context(|| context.to_string()) +} + +fn reg_check_result(context: &str, result: u32) -> UResult<()> { + if result == 0 { + Ok(()) + } else { + Err(io::Error::from_raw_os_error(result as i32).map_err_context(|| context.to_string())) + } +} diff --git a/src/pwsh-install-template.ps1 b/src/pwsh-install-template.ps1 index 335124f..a64f875 100644 --- a/src/pwsh-install-template.ps1 +++ b/src/pwsh-install-template.ps1 @@ -1,24 +1,7 @@ # vvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvv # Inlining the template into the profile shaves off ~10ms (25%). $script:__COREUTILS__ = [System.Collections.Generic.HashSet[string]]::new( - [string[]]@( - 'arch', 'b2sum', 'base32', 'base64', 'basename', - 'basenc', 'cat', 'cksum', 'comm', 'cp', - 'csplit', 'cut', 'date', 'df', 'dirname', - 'du', 'echo', 'env', 'expr', 'factor', - 'false', 'find', 'fmt', 'fold', 'grep', - 'head', 'hostname', 'join', 'la', 'link', - 'ln', 'ls', 'md5sum', 'mkdir', 'mktemp', - 'mv', 'nl', 'nproc', 'numfmt', 'od', - 'pathchk', 'pr', 'printenv', 'printf', 'ptx', - 'pwd', 'readlink', 'realpath', 'rm', 'rmdir', - 'seq', 'sha1sum', 'sha224sum', 'sha256sum', 'sha384sum', - 'sha512sum', 'shuf', 'sleep', 'sort', 'split', - 'stat', 'sum', 'tac', 'tail', 'tee', - 'test', 'touch', 'tr', 'true', 'truncate', - 'tsort', 'unexpand', 'uniq', 'unlink', 'uptime', - 'wc', 'xargs', 'yes' - ), + [string[]]@('!!COREUTILS!!'), [System.StringComparer]::OrdinalIgnoreCase ) diff --git a/src/pwsh-install.ps1 b/src/pwsh-install.ps1 index 7bc34d5..a165f3e 100644 --- a/src/pwsh-install.ps1 +++ b/src/pwsh-install.ps1 @@ -6,7 +6,7 @@ param( [Parameter(Mandatory = $true)] - [ValidateSet('Install', 'Uninstall')] + [ValidateSet('Install', 'Uninstall', 'Refresh')] [string]$Action, [ValidateSet('AllUsers', 'CurrentUser')] [string]$Scope = 'AllUsers', @@ -24,6 +24,9 @@ $SectionMarker = '60b36fc6-2d59-49df-be51-28dd2f4c3c9a' $MarkerLine = "# DO NOT MODIFY -- coreutils -- $SectionMarker" # Earliest PowerShell that supports PSNativeCommandPreserveBytePipe. $MinPwshVersion = [version]'7.4.0' +$CoreutilsRegPath = 'HKLM:\SOFTWARE\Microsoft\coreutils' +# A REG_MULTI_SZ list of utility names (e.g. "ls") that the user has disabled. +$DisabledUtilitiesRegName = 'DisabledUtilities' # Contains SID --> Microsoft.PowerShell_profile.ps1 mappins, # such that we can clean them up on uninstall. $ProfilesRegPath = 'HKLM:\SOFTWARE\Microsoft\coreutils\PowerShellProfiles' @@ -37,18 +40,46 @@ function Remove-FileIfExists([string]$Path) { } } +function Get-EnabledCoreutilsAliases([string]$CmdDir) { + [string[]]$disabled = @() + if ($props = Get-ItemProperty -LiteralPath $CoreutilsRegPath -Name $DisabledUtilitiesRegName -ErrorAction Ignore) { + $disabled = $props.$DisabledUtilitiesRegName + } + $disabled = [System.Collections.Generic.HashSet[string]]::new($disabled, [System.StringComparer]::OrdinalIgnoreCase) + + $aliases = [System.Collections.Generic.HashSet[string]]::new([System.StringComparer]::OrdinalIgnoreCase) + foreach ($file in Get-ChildItem -LiteralPath $CmdDir -Filter '*.cmd' -File -ErrorAction Ignore) { + $alias = [System.IO.Path]::GetFileNameWithoutExtension($file.Name) + if ($alias -and !$disabled.Contains($alias)) { + [void]$aliases.Add($alias) + } + } + + # coreutils-manager is a hidden utility. Theoretically it should not be in the cmd + # directory at all (it's only exposed as coreutils-manager.exe), but just in case... + [void]$aliases.Remove('coreutils-manager') + + if ($aliases.Contains('ls')) { + [void]$aliases.Add('la') + } + + return @($aliases | Sort-Object) +} + function Get-InjectedSection([string]$CmdDir) { $templatePath = Join-Path $PSScriptRoot 'pwsh-install-template.ps1' $template = Get-Content -LiteralPath $templatePath -Raw $cmdDir = [System.IO.Path]::GetFullPath($CmdDir).TrimEnd('\') + '\' - $template = $template.Replace('!!CMDDIR!!', $cmdDir) + $aliases = (Get-EnabledCoreutilsAliases $cmdDir | ForEach-Object { "'$_'" }) -join ',' + $template = $template.Replace("'!!COREUTILS!!'", $aliases) + $template = $template.Replace("!!CMDDIR!!", $cmdDir.Replace("'", "''")) $body = $template.TrimEnd("`r", "`n") return "$MarkerLine`r`n$body`r`n$MarkerLine" } -function Update-PowerShellProfile([string]$Path, [bool] $Install, [bool] $UseBom, [string]$Section) { +function Update-PowerShellProfile([string]$Path, [bool] $Install, [bool] $UseBom, [string]$Section, [bool]$RefreshOnly) { $parent = Split-Path -LiteralPath $Path - if ($Install) { + if ($Install -and !$RefreshOnly) { [void](New-Item -Path $parent -ItemType Directory -Force) } elseif (!(Test-Path -LiteralPath $Path)) { @@ -67,6 +98,9 @@ function Update-PowerShellProfile([string]$Path, [bool] $Install, [bool] $UseBom if ($markerCount -ne 0 -and $markerCount -ne 2) { throw "Invalid coreutils section markers in PowerShell profile: $Path" } + if ($RefreshOnly -and $markerCount -eq 0) { + return + } # Strip the existing section (markers + content + any surrounding blank lines) in one shot. if ($markerCount -eq 2) { @@ -175,7 +209,7 @@ function Get-RecordedProfiles { function Save-RecordedProfile([string]$Sid, [string]$Value) { [void](New-Item -Path $ProfilesRegPath -Force) - Set-ItemProperty -LiteralPath $ProfilesRegPath -Name $Sid -Value $Value -Type String + [void](New-ItemProperty -LiteralPath $ProfilesRegPath -Name $Sid -Value $Value -PropertyType String -Force) } function Remove-RecordedProfile([string]$Sid) { @@ -240,9 +274,45 @@ function Get-ProfilePlan([bool] $Install, [string]$Scope) { return $plan.Values } +function Get-RefreshProfilePlan { + $plan = @{} + + function Add([string]$Path) { + if (!$Path) { + return + } + if ($plan[$Path]) { + return + } + + $plan[$Path] = [PSCustomObject]@{ + Path = $Path + Install = $true + RecordSid = $null + RecordValue = $null + } + } + + foreach ($profilePath in Get-MsiPwshProfilePaths) { + Add $profilePath + } + foreach ($r in Get-RecordedProfiles) { + Add $r.Path + } + Add $PROFILE.CurrentUserCurrentHost + + return $plan.Values +} + $install = $Action -eq 'Install' -$plan = @(Get-ProfilePlan $install $Scope) -$section = if ($install) { +$refresh = $Action -eq 'Refresh' +$plan = if ($refresh) { + @(Get-RefreshProfilePlan) +} +else { + @(Get-ProfilePlan $install $Scope) +} +$section = if ($install -or $refresh) { Get-InjectedSection $CmdDir } else { @@ -250,20 +320,22 @@ else { } foreach ($entry in $plan) { - Update-PowerShellProfile $entry.Path $entry.Install $false $section + Update-PowerShellProfile -Path $entry.Path -Install $entry.Install -UseBom $false -Section $section -RefreshOnly $refresh } # Only adjust records once every Update succeeded. A failure mid-loop leaves # the old records intact so a re-run re-discovers the same paths and retries # the cleanup/install. -foreach ($entry in $plan) { - if (!$entry.RecordSid) { - continue - } - if ($entry.Install) { - Save-RecordedProfile $entry.RecordSid $entry.RecordValue - } - else { - Remove-RecordedProfile $entry.RecordSid +if (!$refresh) { + foreach ($entry in $plan) { + if (!$entry.RecordSid) { + continue + } + if ($entry.Install) { + Save-RecordedProfile $entry.RecordSid $entry.RecordValue + } + else { + Remove-RecordedProfile $entry.RecordSid + } } }