pioduino/uv
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 1 Wiki Activity

feat: more rust in trampoline (#5750)

Browse Source 
## Summary

This is an experimental PR to replace more unsafe calls with more rust
while still trying to keep the binary size small enough. These changes
roughly increase the size of the trampolines to about 40kb~. This is a
alternate PR to https://github.com/astral-sh/uv/pull/5751.

The primary changes here include
* Switch to use rust path components for ease of path management
* Leverage `std::process::exit` for process exit and cleanup
* Use `std::io::Error::last_os_error` for IO Errors to remove
`FormatMessage` complexity
* Use `std::env::current_exe` to get the current executable instead of
`GetModuleFileNameA`

## Test Plan

Added one more existing test case to trampoline tests.
Still need to verify dunce::canonicalize is desired or not on
find_python_exe.

---------

Co-authored-by: konstin <konstin@mailbox.org>
This commit is contained in:
samypr100
2024-08-07 04:19:38 -04:00
committed by GitHub
parent d0b16f9018
commit 2cd63f06dc
13 changed files with 109 additions and 210 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.github/workflows/ci.yml
+1
View File
@@ -347,6 +347,7 @@ jobs:
if: matrix.target-arch == 'x86_64'
run: |
cargo xwin bloat --release --target x86_64-pc-windows-msvc | \
grep -v -i -E 'core::fmt::write|core::fmt::getcount' | \
grep -q -E 'core::fmt|std::panicking|std::backtrace_rs' && exit 1 || exit 0
env:
XWIN_ARCH: "x86_64"
crates/install-wheel-rs/src/wheel.rs
+9 -9
View File
@@ -1018,14 +1018,14 @@ mod test {
#[test]
#[cfg(all(windows, target_arch = "x86"))]
fn test_launchers_are_small() {
// At time of writing, they are 17408 bytes.
// At time of writing, they are 45kb~ bytes.
assert!(
super::LAUNCHER_I686_GUI.len() < 25 * 1024,
super::LAUNCHER_I686_GUI.len() < 45 * 1024,
"GUI launcher: {}",
super::LAUNCHER_I686_GUI.len()
);
assert!(
super::LAUNCHER_I686_CONSOLE.len() < 25 * 1024,
super::LAUNCHER_I686_CONSOLE.len() < 45 * 1024,
"CLI launcher: {}",
super::LAUNCHER_I686_CONSOLE.len()
);
@@ -1034,14 +1034,14 @@ mod test {
#[test]
#[cfg(all(windows, target_arch = "x86_64"))]
fn test_launchers_are_small() {
// At time of writing, they are 21504 and 20480 bytes.
// At time of writing, they are 45kb~ bytes.
assert!(
super::LAUNCHER_X86_64_GUI.len() < 25 * 1024,
super::LAUNCHER_X86_64_GUI.len() < 45 * 1024,
"GUI launcher: {}",
super::LAUNCHER_X86_64_GUI.len()
);
assert!(
super::LAUNCHER_X86_64_CONSOLE.len() < 25 * 1024,
super::LAUNCHER_X86_64_CONSOLE.len() < 45 * 1024,
"CLI launcher: {}",
super::LAUNCHER_X86_64_CONSOLE.len()
);
@@ -1050,14 +1050,14 @@ mod test {
#[test]
#[cfg(all(windows, target_arch = "aarch64"))]
fn test_launchers_are_small() {
// At time of writing, they are 20480 and 19456 bytes.
// At time of writing, they are 45kb~ bytes.
assert!(
super::LAUNCHER_AARCH64_GUI.len() < 25 * 1024,
super::LAUNCHER_AARCH64_GUI.len() < 45 * 1024,
"GUI launcher: {}",
super::LAUNCHER_AARCH64_GUI.len()
);
assert!(
super::LAUNCHER_AARCH64_CONSOLE.len() < 25 * 1024,
super::LAUNCHER_AARCH64_CONSOLE.len() < 45 * 1024,
"CLI launcher: {}",
super::LAUNCHER_AARCH64_CONSOLE.len()
);
crates/uv-trampoline/Cargo.lock
Generated
+1
View File
@@ -765,6 +765,7 @@ dependencies = [
"anyhow",
"assert_cmd",
"assert_fs",
"dunce",
"embed-manifest",
"fs-err",
"thiserror",
crates/uv-trampoline/Cargo.toml
+1 -7
View File
@@ -36,21 +36,15 @@ windows = { version = "0.58.0", features = [
"std",
"Win32_Foundation",
"Win32_Security",
"Win32_Storage_FileSystem",
"Win32_System_Console",
"Win32_System_Diagnostics_Debug",
"Win32_System_Environment",
"Win32_System_IO",
"Win32_System_JobObjects",
"Win32_System_JobObjects",
"Win32_System_LibraryLoader",
"Win32_System_Memory",
"Win32_System_Threading",
"Win32_System_WindowsProgramming",
"Win32_UI_WindowsAndMessaging",
] }
ufmt-write = "0.1.0"
ufmt = { version = "0.2.0", features = ["std"] }
dunce = { version = "1.0.4" }
[build-dependencies]
embed-manifest = "1.4.0"
crates/uv-trampoline/src/bounce.rs
+68 -169
View File
@@ -1,37 +1,30 @@
use std::ffi::{c_void, CStr, CString};
use std::mem::{size_of, size_of_val};
#![allow(clippy::disallowed_types)]
use std::ffi::{c_void, CString};
use std::fs::File;
use std::io::{Read, Seek, SeekFrom};
use std::mem::size_of;
use std::path::{Path, PathBuf};
use std::vec::Vec;
use windows::core::{s, PCSTR, PSTR};
use windows::core::{s, PSTR};
use windows::Win32::{
Foundation::{
CloseHandle, GetLastError, SetHandleInformation, SetLastError, BOOL,
ERROR_INSUFFICIENT_BUFFER, ERROR_SUCCESS, HANDLE, HANDLE_FLAG_INHERIT,
INVALID_HANDLE_VALUE, MAX_PATH, TRUE,
},
Storage::FileSystem::{
CreateFileA, GetFileSizeEx, ReadFile, SetFilePointerEx, FILE_ATTRIBUTE_NORMAL, FILE_BEGIN,
FILE_GENERIC_READ, FILE_SHARE_READ, OPEN_EXISTING,
CloseHandle, SetHandleInformation, BOOL, HANDLE, HANDLE_FLAG_INHERIT, INVALID_HANDLE_VALUE,
TRUE,
},
System::Console::{
GetStdHandle, SetConsoleCtrlHandler, SetStdHandle, STD_ERROR_HANDLE, STD_INPUT_HANDLE,
STD_OUTPUT_HANDLE,
},
System::Diagnostics::Debug::{
FormatMessageA, FORMAT_MESSAGE_ALLOCATE_BUFFER, FORMAT_MESSAGE_FROM_SYSTEM,
FORMAT_MESSAGE_IGNORE_INSERTS,
},
System::Environment::GetCommandLineA,
System::JobObjects::{
AssignProcessToJobObject, CreateJobObjectA, JobObjectExtendedLimitInformation,
QueryInformationJobObject, SetInformationJobObject, JOBOBJECT_EXTENDED_LIMIT_INFORMATION,
JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE, JOB_OBJECT_LIMIT_SILENT_BREAKAWAY_OK,
},
System::LibraryLoader::GetModuleFileNameA,
System::Threading::{
CreateProcessA, ExitProcess, GetExitCodeProcess, GetStartupInfoA, WaitForInputIdle,
WaitForSingleObject, INFINITE, PROCESS_CREATION_FLAGS, PROCESS_INFORMATION,
STARTF_USESTDHANDLES, STARTUPINFOA,
CreateProcessA, GetExitCodeProcess, GetStartupInfoA, WaitForInputIdle, WaitForSingleObject,
INFINITE, PROCESS_CREATION_FLAGS, PROCESS_INFORMATION, STARTF_USESTDHANDLES, STARTUPINFOA,
},
UI::WindowsAndMessaging::{
CreateWindowExA, DestroyWindow, GetMessageA, PeekMessageA, PostMessageA, HWND_MESSAGE, MSG,
@@ -47,18 +40,21 @@ const MAX_PATH_LEN: u32 = 32 * 1024;
/// Transform `<command> <arguments>` to `python <command> <arguments>`.
fn make_child_cmdline() -> CString {
let executable_name: CString = executable_filename();
let python_exe = find_python_exe(&executable_name);
let executable_name = std::env::current_exe().unwrap_or_else(|_| {
eprintln!("Failed to get executable name");
exit_with_status(1);
});
let python_exe = find_python_exe(executable_name.as_ref());
let mut child_cmdline = Vec::<u8>::new();
push_quoted_path(&python_exe, &mut child_cmdline);
push_quoted_path(python_exe.as_ref(), &mut child_cmdline);
child_cmdline.push(b' ');
// Use the full executable name because CMD only passes the name of the executable (but not the path)
// when e.g. invoking `black` instead of `<PATH_TO_VENV>/Scripts/black` and Python then fails
// to find the file. Unfortunately, this complicates things because we now need to split the executable
// from the arguments string...
push_quoted_path(&executable_name, &mut child_cmdline);
push_quoted_path(executable_name.as_ref(), &mut child_cmdline);
push_arguments(&mut child_cmdline);
@@ -67,7 +63,7 @@ fn make_child_cmdline() -> CString {
// Helpful when debugging trampoline issues
// eprintln!(
// "executable_name: '{}'\nnew_cmdline: {}",
// std::str::from_utf8(executable_name.to_bytes()).unwrap(),
// &*executable_name.to_string_lossy(),
// std::str::from_utf8(child_cmdline.as_slice()).unwrap()
// );
@@ -77,9 +73,9 @@ fn make_child_cmdline() -> CString {
})
}
fn push_quoted_path(path: &CStr, command: &mut Vec<u8>) {
fn push_quoted_path(path: &Path, command: &mut Vec<u8>) {
command.push(b'"');
for byte in path.to_bytes() {
for byte in path.as_os_str().as_encoded_bytes() {
if *byte == b'"' {
// 3 double quotes: one to end the quoted span, one to become a literal double-quote,
// and one to start a new quoted span.
@@ -91,45 +87,6 @@ fn push_quoted_path(path: &CStr, command: &mut Vec<u8>) {
command.extend(br#"""#);
}
/// Returns the full path of the executable.
/// See https://learn.microsoft.com/en-us/windows/win32/api/libloaderapi/nf-libloaderapi-getmodulefilenamea
fn executable_filename() -> CString {
// MAX_PATH is a lie, Windows paths can be longer.
// https://learn.microsoft.com/en-us/windows/win32/fileio/naming-a-file#maximum-path-length-limitation
// But it's a good first guess, usually paths are short, and we should only need a single attempt.
let mut buffer: Vec<u8> = vec![0; MAX_PATH as usize];
loop {
// Call the Windows API function to get the module file name
let len = unsafe { GetModuleFileNameA(None, &mut buffer) };
// That's the error condition because len doesn't include the trailing null byte
if len as usize == buffer.len() {
let last_error = unsafe { GetLastError() };
match last_error {
ERROR_INSUFFICIENT_BUFFER => {
unsafe { SetLastError(ERROR_SUCCESS) };
// Try again with twice the size
buffer.resize(buffer.len() * 2, 0);
}
err => {
print_last_error_and_exit(&format!(
"Failed to get executable name (code: {})",
err.0
));
}
}
} else {
buffer.truncate(len as usize + b"\0".len());
break;
}
}
CString::from_vec_with_nul(buffer).unwrap_or_else(|_| {
eprintln!("Executable name is not correctly null terminated");
exit_with_status(1);
})
}
/// Reads the executable binary from the back to find the path to the Python executable that is written
/// after the ZIP file content.
///
@@ -140,65 +97,44 @@ fn executable_filename() -> CString {
///
/// # Panics
/// If there's any IO error, or the file does not conform to the specified format.
fn find_python_exe(executable_name: &CStr) -> CString {
let file_handle = unsafe {
CreateFileA(
PCSTR::from_raw(executable_name.as_ptr() as *const _),
FILE_GENERIC_READ.0,
FILE_SHARE_READ,
None,
OPEN_EXISTING,
FILE_ATTRIBUTE_NORMAL,
None,
)
}
.unwrap_or_else(|_| {
fn find_python_exe(executable_name: &Path) -> PathBuf {
let mut file_handle = File::open(executable_name).unwrap_or_else(|_| {
print_last_error_and_exit(&format!(
"Failed to open executable '{}'",
&*executable_name.to_string_lossy(),
))
));
});
let mut file_size: i64 = 0;
// `SetFilePointerEx` supports setting the file pointer from the back, but pointing it past the file's start
// results in an error. That's why we need to know the file size to avoid ever seeking past the start of the file.
if unsafe { GetFileSizeEx(file_handle, &mut file_size) }.is_err() {
let metadata = executable_name.metadata().unwrap_or_else(|_| {
print_last_error_and_exit(&format!(
"Failed to get the size of the executable '{}'",
&*executable_name.to_string_lossy(),
));
}
});
let file_size = metadata.len();
// Start with a size of 1024 bytes which should be enough for most paths but avoids reading the
// entire file.
let mut buffer: Vec<u8> = Vec::new();
let mut bytes_to_read = 1024.min(u32::try_from(file_size).unwrap_or(u32::MAX));
let path: CString = loop {
let path: String = loop {
// SAFETY: Casting to usize is safe because we only support 64bit systems where usize is guaranteed to be larger than u32.
buffer.resize(bytes_to_read as usize, 0);
if unsafe {
SetFilePointerEx(
file_handle,
file_size - i64::from(bytes_to_read),
None,
FILE_BEGIN,
)
}
.is_err()
{
print_last_error_and_exit("Failed to set the file pointer to the end of the file");
}
file_handle
.seek(SeekFrom::Start(file_size - u64::from(bytes_to_read)))
.unwrap_or_else(|_| {
print_last_error_and_exit("Failed to set the file pointer to the end of the file");
});
let mut read_bytes = bytes_to_read;
if unsafe { ReadFile(file_handle, Some(&mut buffer), Some(&mut read_bytes), None) }.is_err()
{
// Pulls in core::fmt::{write, Write, getcount}
let read_bytes = file_handle.read(&mut buffer).unwrap_or_else(|_| {
print_last_error_and_exit("Failed to read the executable file");
}
});
// Truncate the buffer to the actual number of bytes read.
buffer.truncate(read_bytes as usize);
buffer.truncate(read_bytes);
if !buffer.ends_with(&MAGIC_NUMBER) {
eprintln!("Magic number 'UVUV' not found at the end of the file. Did you append the magic number, the length and the path to the python executable at the end of the file?");
@@ -234,10 +170,9 @@ fn find_python_exe(executable_name: &CStr) -> CString {
if let Some(path_offset) = buffer.len().checked_sub(path_len) {
buffer.drain(..path_offset);
buffer.push(b'\0');
break CString::from_vec_with_nul(buffer).unwrap_or_else(|_| {
eprintln!("Python executable path is not correctly null terminated");
break String::from_utf8(buffer).unwrap_or_else(|_| {
eprintln!("Python executable path is not a valid UTF-8 encoded path");
exit_with_status(1);
});
} else {
@@ -245,46 +180,32 @@ fn find_python_exe(executable_name: &CStr) -> CString {
// MAGIC_NUMBER is 4 bytes and PATH_LEN_SIZE is 4 bytes.
bytes_to_read = (path_len + MAGIC_NUMBER.len() + PATH_LEN_SIZE) as u32;
if i64::from(bytes_to_read) > file_size {
if u64::from(bytes_to_read) > file_size {
eprintln!("The length of the python executable path exceeds the file size. Verify that the path length is appended to the end of the launcher script as a u32 in little endian");
exit_with_status(1);
}
}
};
if unsafe { CloseHandle(file_handle) }.is_err() {
print_last_error_and_exit("Failed to close file handle");
}
if is_absolute(&path) {
let path = PathBuf::from(path);
let path = if path.is_absolute() {
path
} else {
let parent_dir = match executable_name
.to_bytes()
.rsplitn(2, |c| *c == b'\\')
.last()
{
Some(parent_dir) => parent_dir,
let parent_dir = match executable_name.parent() {
Some(parent) => parent,
None => {
eprintln!("Script path has unknown separator characters");
eprintln!("Executable path has no parent directory");
exit_with_status(1);
}
};
let final_path = [parent_dir, b"\\", path.as_bytes()].concat();
CString::new(final_path).unwrap_or_else(|_| {
eprintln!("Could not construct the absolute path to the Python executable");
exit_with_status(1);
})
}
}
parent_dir.join(path)
};
/// Returns `true` if the path is absolute.
///
/// In this context, as in the Rust standard library, c:\windows is absolute, while c:temp and
/// \temp are not.
fn is_absolute(path: &CStr) -> bool {
let path = path.to_bytes();
path.len() >= 3 && path[0].is_ascii_alphabetic() && path[1] == b':' && path[2] == b'\\'
// NOTICE: dunce adds 5kb~
dunce::canonicalize(path.as_path()).unwrap_or_else(|_| {
eprintln!("Failed to canonicalize script path");
exit_with_status(1);
})
}
fn push_arguments(output: &mut Vec<u8>) {
@@ -540,47 +461,25 @@ pub fn bounce(is_gui: bool) -> ! {
#[cold]
fn print_last_error_and_exit(message: &str) -> ! {
let err = unsafe { GetLastError() };
eprintln!("Received error code: {}", err.0);
let mut msg_ptr = PSTR::null();
let size = unsafe {
FormatMessageA(
FORMAT_MESSAGE_ALLOCATE_BUFFER
| FORMAT_MESSAGE_FROM_SYSTEM
| FORMAT_MESSAGE_IGNORE_INSERTS,
None,
err.0,
0,
// Weird calling convention: this argument is typed as *mut u8,
// but if you pass FORMAT_MESSAGE_ALLOCATE_BUFFER then you have to
// *actually* pass in a *mut *mut u8 and just lie about the type.
// Getting Rust to do this requires some convincing.
PSTR(&mut msg_ptr.0 as *mut _ as *mut _),
0,
None,
)
};
if size == 0 || msg_ptr.0.is_null() {
eprintln!(
"{}: with code {} (failed to get error message)",
message, err.0
);
} else {
let reason = unsafe {
let reason = std::slice::from_raw_parts(msg_ptr.0, size as usize);
std::str::from_utf8_unchecked(reason)
};
eprintln!("(uv internal error) {}: {}", message, reason);
}
// Note: We don't need to free the buffer here because we're going to exit anyway.
let err = std::io::Error::last_os_error();
let err_no_str = err
.raw_os_error()
.map(|raw_error| format!(" (os error {})", raw_error))
.unwrap_or_default();
// we can't access sys::os::error_string directly so err.kind().to_string()
// is the closest we can get to while avoiding bringing in a large chunk of core::fmt
eprintln!(
"(uv internal error) {}: {}.{}",
message,
err.kind().to_string(),
err_no_str
);
exit_with_status(1);
}
#[cold]
fn exit_with_status(code: u32) -> ! {
unsafe {
ExitProcess(code);
}
// ~5-10kb
// Pulls in core::fmt::{write, Write, getcount}
std::process::exit(code as _)
}
crates/uv-trampoline/src/diagnostics.rs
+10 -23
View File
@@ -1,15 +1,12 @@
use std::convert::Infallible;
use std::ffi::CString;
use std::io::Write;
use std::os::windows::io::AsRawHandle;
use std::string::String;
use ufmt_write::uWrite;
use windows::core::PCSTR;
use windows::Win32::{
Foundation::INVALID_HANDLE_VALUE,
Storage::FileSystem::WriteFile,
System::Console::{GetStdHandle, STD_ERROR_HANDLE},
UI::WindowsAndMessaging::{MessageBoxA, MESSAGEBOX_STYLE},
};
use windows::Win32::UI::WindowsAndMessaging::{MessageBoxA, MESSAGEBOX_STYLE};
#[macro_export]
macro_rules! eprintln {
@@ -41,22 +38,12 @@ impl uWrite for StringBuffer {
#[cold]
pub(crate) fn write_diagnostic(message: &str) {
let handle = unsafe { GetStdHandle(STD_ERROR_HANDLE) }.unwrap_or(INVALID_HANDLE_VALUE);
let mut written: u32 = 0;
let mut remaining = message;
while !remaining.is_empty() {
// If we get an error, it means we tried to write to an invalid handle (GUI Application)
// and we should try to write to a window instead
if unsafe { WriteFile(handle, Some(remaining.as_bytes()), Some(&mut written), None) }
.is_err()
{
let nul_terminated = unsafe { CString::new(message.as_bytes()).unwrap_unchecked() };
let pcstr_message = PCSTR::from_raw(nul_terminated.as_ptr() as *const _);
unsafe { MessageBoxA(None, pcstr_message, None, MESSAGEBOX_STYLE(0)) };
return;
}
if let Some(out) = remaining.get(written as usize..) {
remaining = out
}
let mut stderr = std::io::stderr();
if !stderr.as_raw_handle().is_null() {
let _ = stderr.write_all(message.as_bytes());
} else {
let nul_terminated = unsafe { CString::new(message.as_bytes()).unwrap_unchecked() };
let pcstr_message = PCSTR::from_raw(nul_terminated.as_ptr() as *const _);
unsafe { MessageBoxA(None, pcstr_message, None, MESSAGEBOX_STYLE(0)) };
}
}
crates/uv-trampoline/tests/harness.rs
+19 -2
View File
@@ -85,6 +85,8 @@ import sys
def main_console() -> None:
print("Hello from uv-trampoline-console.exe", file=sys.stdout)
print("Hello from uv-trampoline-console.exe", file=sys.stderr)
for arg in sys.argv[1:]:
print(arg, file=sys.stderr)
if __name__ == "__main__":
sys.argv[0] = re.sub(r"(-script\.pyw|\.exe)?$", "", sys.argv[0])
@@ -212,13 +214,28 @@ fn generate_console_launcher() -> Result<()> {
console_bin_path.path().simplified_display()
);
let stdout_predicate = "Hello from uv-trampoline-console.exe\r\n";
let stderr_predicate = "Hello from uv-trampoline-console.exe\r\n";
// Test Console Launcher
#[cfg(windows)]
Command::new(console_bin_path.path())
.assert()
.success()
.stdout("Hello from uv-trampoline-console.exe\r\n")
.stderr("Hello from uv-trampoline-console.exe\r\n");
.stdout(stdout_predicate)
.stderr(stderr_predicate);
let args_to_test = vec!["foo", "bar", "foo bar", "foo \"bar\"", "foo 'bar'"];
let stderr_predicate = format!("{}{}\r\n", stderr_predicate, args_to_test.join("\r\n"));
// Test Console Launcher (with args)
#[cfg(windows)]
Command::new(console_bin_path.path())
.args(args_to_test)
.assert()
.success()
.stdout(stdout_predicate)
.stderr(stderr_predicate);
Ok(())
}
crates/uv-trampoline/trampolines/uv-trampoline-aarch64-console.exe
BIN
View File
Binary file not shown.
crates/uv-trampoline/trampolines/uv-trampoline-aarch64-gui.exe
BIN
View File
Binary file not shown.
crates/uv-trampoline/trampolines/uv-trampoline-i686-console.exe
BIN
View File
Binary file not shown.
crates/uv-trampoline/trampolines/uv-trampoline-i686-gui.exe
BIN
View File
Binary file not shown.
crates/uv-trampoline/trampolines/uv-trampoline-x86_64-console.exe
Regular → Executable
BIN
View File
Binary file not shown.
crates/uv-trampoline/trampolines/uv-trampoline-x86_64-gui.exe
Regular → Executable
BIN
View File
Binary file not shown.