Skip to content

Commit

Permalink
Use trampolines for Python executables on Windows
Browse files Browse the repository at this point in the history
  • Loading branch information
zanieb committed Oct 31, 2024
1 parent b52f229 commit 6d494bb
Show file tree
Hide file tree
Showing 11 changed files with 238 additions and 44 deletions.
2 changes: 1 addition & 1 deletion .github/workflows/ci.yml
Original file line number Diff line number Diff line change
Expand Up @@ -296,7 +296,7 @@ jobs:
# See https://github.com/astral-sh/uv/issues/6940
UV_LINK_MODE: copy
run: |
cargo nextest run --no-default-features --features python,pypi --workspace --status-level skip --failure-output immediate-final --no-fail-fast -j 20 --final-status-level slow
cargo nextest run --no-default-features --features python,pypi,python-managed --workspace --status-level skip --failure-output immediate-final --no-fail-fast -j 20 --final-status-level slow
- name: "Smoke test"
working-directory: ${{ env.UV_WORKSPACE }}
Expand Down
1 change: 1 addition & 0 deletions Cargo.lock

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

13 changes: 7 additions & 6 deletions crates/uv-fs/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -104,21 +104,22 @@ pub fn remove_symlink(path: impl AsRef<Path>) -> std::io::Result<()> {
fs_err::remove_file(path.as_ref())
}

/// Create a symlink at `dst` pointing to `src` or, on Windows, copy `src` to `dst`.
/// Create a symlink at `dst` pointing to `src` on Unix or copy `src` to `dst` on Windows
///
/// This does not replace an existing symlink or file at `dst`.
///
/// This does not fallback to copying on Unix.
///
/// This function should only be used for files. If targeting a directory, use [`replace_symlink`]
/// instead; it will use a junction on Windows, which is more performant.
pub fn symlink_copy_fallback_file(
src: impl AsRef<Path>,
dst: impl AsRef<Path>,
) -> std::io::Result<()> {
pub fn symlink_or_copy_file(src: impl AsRef<Path>, dst: impl AsRef<Path>) -> std::io::Result<()> {
#[cfg(windows)]
{
fs_err::copy(src.as_ref(), dst.as_ref())?;
}
#[cfg(unix)]
{
std::os::unix::fs::symlink(src.as_ref(), dst.as_ref())?;
fs_err::os::unix::fs::symlink(src.as_ref(), dst.as_ref())?;
}

Ok(())
Expand Down
1 change: 1 addition & 0 deletions crates/uv-python/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,7 @@ uv-platform-tags = { workspace = true }
uv-pypi-types = { workspace = true }
uv-state = { workspace = true }
uv-static = { workspace = true }
uv-trampoline-builder = { workspace = true }
uv-warnings = { workspace = true }

anyhow = { workspace = true }
Expand Down
74 changes: 55 additions & 19 deletions crates/uv-python/src/managed.rs
Original file line number Diff line number Diff line change
@@ -1,15 +1,20 @@
use core::fmt;
use fs_err as fs;
use itertools::Itertools;
use std::cmp::Reverse;
use std::ffi::OsStr;
use std::io::{self, Write};
use std::path::{Path, PathBuf};
use std::str::FromStr;

use fs_err::{self as fs, File};
use itertools::Itertools;
use same_file::is_same_file;
use thiserror::Error;
use tracing::{debug, warn};

use uv_fs::{symlink_or_copy_file, LockedFile, Simplified};
use uv_state::{StateBucket, StateStore};
use uv_static::EnvVars;
use uv_trampoline_builder::{windows_python_launcher, Launcher};

use crate::downloads::Error as DownloadError;
use crate::implementation::{
Expand All @@ -21,9 +26,6 @@ use crate::platform::Error as PlatformError;
use crate::platform::{Arch, Libc, Os};
use crate::python_version::PythonVersion;
use crate::{PythonRequest, PythonVariant};
use uv_fs::{LockedFile, Simplified};
use uv_static::EnvVars;

#[derive(Error, Debug)]
pub enum Error {
#[error(transparent)]
Expand Down Expand Up @@ -74,6 +76,8 @@ pub enum Error {
},
#[error("Failed to find a directory to install executables into")]
NoExecutableDirectory,
#[error(transparent)]
LauncherError(#[from] uv_trampoline_builder::Error),
#[error("Failed to read managed Python directory name: {0}")]
NameError(String),
#[error("Failed to construct absolute path to managed Python directory: {}", _0.user_display())]
Expand Down Expand Up @@ -425,7 +429,7 @@ impl ManagedPythonInstallation {
continue;
}

match uv_fs::symlink_copy_fallback_file(&python, &executable) {
match uv_fs::symlink_or_copy_file(&python, &executable) {
Ok(()) => {
debug!(
"Created link {} -> {}",
Expand Down Expand Up @@ -475,28 +479,60 @@ impl ManagedPythonInstallation {
Ok(())
}

/// Create a link to the Python executable in the given `bin` directory.
pub fn create_bin_link(&self, bin: &Path) -> Result<PathBuf, Error> {
/// Create a link to the managed Python executable.
pub fn create_bin_link(&self, target: &Path) -> Result<(), Error> {
let python = self.executable();

let bin = target.parent().ok_or(Error::NoExecutableDirectory)?;
fs_err::create_dir_all(bin).map_err(|err| Error::ExecutableDirectory {
to: bin.to_path_buf(),
err,
})?;

// TODO(zanieb): Add support for a "default" which
let python_in_bin = bin.join(self.key.versioned_executable_name());
if cfg!(unix) {
// Note this will never copy on Unix — we use it here to allow compilation on Windows
match symlink_or_copy_file(&python, target) {
Ok(()) => Ok(()),
Err(err) if err.kind() == io::ErrorKind::NotFound => {
Err(Error::MissingExecutable(python.clone()))
}
Err(err) => Err(Error::LinkExecutable {
from: python,
to: target.to_path_buf(),
err,
}),
}
} else if cfg!(windows) {
// TODO(zanieb): Install GUI launchers as well
let launcher = windows_python_launcher(&python, false)?;
match File::create(target)?.write_all(launcher.as_ref()) {
Ok(()) => Ok(()),
Err(err) => Err(Error::LinkExecutable {
from: python,
to: target.to_path_buf(),
err,
}),
}
} else {
unimplemented!("Only Windows and Unix systems are supported.")
}
}

match uv_fs::symlink_copy_fallback_file(&python, &python_in_bin) {
Ok(()) => Ok(python_in_bin),
Err(err) if err.kind() == io::ErrorKind::NotFound => {
Err(Error::MissingExecutable(python.clone()))
/// Returns `true` if the path is a link to this installation's binary, e.g., as created by
/// [`ManagedPythonInstallation::create_bin_link`].
pub fn is_bin_link(&self, path: &Path) -> bool {
if cfg!(unix) {
is_same_file(path, self.executable()).unwrap_or_default()
} else if cfg!(windows) {
let Some(launcher) = Launcher::try_from_path(path).unwrap_or_default() else {
return false;
};
if !matches!(launcher.kind, uv_trampoline_builder::LauncherKind::Python) {
return false;
}
Err(err) => Err(Error::LinkExecutable {
from: python,
to: python_in_bin,
err,
}),
launcher.python_path == self.executable()
} else {
unreachable!("Only Windows and Unix are supported")
}
}
}
Expand Down
2 changes: 2 additions & 0 deletions crates/uv-trampoline-builder/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,8 @@ workspace = true

[dependencies]
uv-fs = { workspace = true }

fs-err = {workspace = true }
thiserror = { workspace = true }
zip = { workspace = true }

Expand Down
155 changes: 150 additions & 5 deletions crates/uv-trampoline-builder/src/lib.rs
Original file line number Diff line number Diff line change
@@ -1,6 +1,8 @@
use std::io::{Cursor, Write};
use std::path::Path;
use std::io::{self, Cursor, Read, Seek, Write};
use std::path::{Path, PathBuf};
use std::str::Utf8Error;

use fs_err::File;
use thiserror::Error;
use uv_fs::Simplified;
use zip::write::FileOptions;
Expand Down Expand Up @@ -30,28 +32,157 @@ const LAUNCHER_AARCH64_GUI: &[u8] =
const LAUNCHER_AARCH64_CONSOLE: &[u8] =
include_bytes!("../../uv-trampoline/trampolines/uv-trampoline-aarch64-console.exe");

// See `uv-trampoline::bounce`. These numbers must match.
const PATH_LENGTH_SIZE: usize = size_of::<u32>();
const MAX_PATH_LENGTH: u32 = 32 * 1024;
const MAGIC_NUMBER_SIZE: usize = 4;

#[derive(Debug)]
pub struct Launcher {
pub kind: LauncherKind,
pub python_path: PathBuf,
}

impl Launcher {
/// Read [`Launcher`] metadata from a trampoline executable file.
///
/// Returns `Ok(None)` if the file is not a trampoline executable.
/// Returns `Err` if the file looks like a trampoline executable but is formatted incorrectly.
///
/// Expects the following metadata to be at the end of the file:
///
/// ```text
/// - file path (no greater than 32KB)
/// - file path length (u32)
/// - magic number(4 bytes)
/// ```
///
/// This should only be used on Windows, but should just return `Ok(None)` on other platforms.
///
/// This is an implementation of [`uv-trampoline::bounce::read_trampoline_metadata`] that
/// returns errors instead of panicking. Unlike the utility there, we don't assume that the
/// file we are reading is a trampoline.
#[allow(clippy::cast_possible_wrap)]
pub fn try_from_path(path: &Path) -> Result<Option<Self>, Error> {
let mut file = File::open(path)?;

// Read the magic number
let Some(kind) = LauncherKind::try_from_file(&mut file)? else {
return Ok(None);
};

// Seek to the start of the path length.
let path_length_offset = (MAGIC_NUMBER_SIZE + PATH_LENGTH_SIZE) as i64;
file.seek(io::SeekFrom::End(-path_length_offset))
.map_err(|err| {
Error::InvalidLauncherSeek("path length".to_string(), path_length_offset, err)
})?;

// Read the path length
let mut buffer = [0; PATH_LENGTH_SIZE];
file.read_exact(&mut buffer)
.map_err(|err| Error::InvalidLauncherRead("path length".to_string(), err))?;

let path_length = {
let raw_length = u32::from_le_bytes(buffer);

if raw_length > MAX_PATH_LENGTH {
return Err(Error::InvalidPathLength(raw_length));
}

// SAFETY: Above we guarantee the length is less than 32KB
raw_length as usize
};

// Seek to the start of the path
let path_offset = (MAGIC_NUMBER_SIZE + PATH_LENGTH_SIZE + path_length) as i64;
file.seek(io::SeekFrom::End(-path_offset)).map_err(|err| {
Error::InvalidLauncherSeek("executable path".to_string(), path_offset, err)
})?;

// Read the path
let mut buffer = vec![0u8; path_length];
file.read_exact(&mut buffer)
.map_err(|err| Error::InvalidLauncherRead("executable path".to_string(), err))?;

let path = PathBuf::from(
String::from_utf8(buffer).map_err(|err| Error::InvalidPath(err.utf8_error()))?,
);

Ok(Some(Self {
kind,
python_path: path,
}))
}
}

/// The kind of trampoline launcher to create.
///
/// See [`uv-trampoline::bounce::TrampolineKind`].
enum LauncherKind {
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum LauncherKind {
/// The trampoline should execute itself, it's a zipped Python script.
Script,
/// The trampoline should just execute Python, it's a proxy Python executable.
Python,
}

impl LauncherKind {
const fn magic_number(&self) -> &'static [u8; 4] {
/// Return the magic number for this [`LauncherKind`].
const fn magic_number(self) -> &'static [u8; 4] {
match self {
Self::Script => b"UVSC",
Self::Python => b"UVPY",
}
}

/// Read a [`LauncherKind`] from 4 byte buffer.
///
/// If the buffer does not contain a matching magic number, `None` is returned.
fn try_from_bytes(bytes: [u8; MAGIC_NUMBER_SIZE]) -> Option<Self> {
if &bytes == Self::Script.magic_number() {
return Some(Self::Script);
}
if &bytes == Self::Python.magic_number() {
return Some(Self::Python);
}
None
}

/// Read a [`LauncherKind`] from a file handle, based on the magic number.
///
/// This will mutate the file handle, seeking to the end of the file.
///
/// If the file cannot be read, an [`io::Error`] is returned. If the path is not a launcher,
/// `None` is returned.
#[allow(clippy::cast_possible_wrap)]
pub fn try_from_file(file: &mut File) -> Result<Option<Self>, Error> {
// If the file is less than four bytes, it's not a launcher.
let Ok(_) = file.seek(io::SeekFrom::End(-(MAGIC_NUMBER_SIZE as i64))) else {
return Ok(None);
};

let mut buffer = [0; MAGIC_NUMBER_SIZE];
file.read_exact(&mut buffer)
.map_err(|err| Error::InvalidLauncherRead("magic number".to_string(), err))?;

Ok(Self::try_from_bytes(buffer))
}
}

/// Note: The caller is responsible for adding the path of the wheel we're installing.
#[derive(Error, Debug)]
pub enum Error {
#[error(transparent)]
Io(#[from] io::Error),
#[error("Only paths with a length up to 32KB are supported but found a length of {0} bytes")]
InvalidPathLength(u32),
#[error("Failed to parse executable path")]
InvalidPath(#[source] Utf8Error),
#[error("Failed to seek to {0} at offset {1}")]
InvalidLauncherSeek(String, i64, #[source] io::Error),
#[error("Failed to read launcher {0}")]
InvalidLauncherRead(String, #[source] io::Error),
#[error(
"Unable to create Windows launcher for: {0} (only x86_64, x86, and arm64 are supported)"
)]
Expand Down Expand Up @@ -192,7 +323,7 @@ mod test {

use which::which;

use super::{windows_python_launcher, windows_script_launcher};
use super::{windows_python_launcher, windows_script_launcher, Launcher, LauncherKind};

#[test]
#[cfg(all(windows, target_arch = "x86", feature = "production"))]
Expand Down Expand Up @@ -340,6 +471,13 @@ if __name__ == "__main__":
.stdout(stdout_predicate)
.stderr(stderr_predicate);

let launcher = Launcher::try_from_path(console_bin_path.path())
.expect("We should succeed at reading the launcher")
.expect("The launcher should be valid");

assert!(launcher.kind == LauncherKind::Script);
assert!(launcher.python_path == python_executable_path);

Ok(())
}

Expand Down Expand Up @@ -371,6 +509,13 @@ if __name__ == "__main__":
.success()
.stdout("Hello from Python Launcher\r\n");

let launcher = Launcher::try_from_path(console_bin_path.path())
.expect("We should succeed at reading the launcher")
.expect("The launcher should be valid");

assert!(launcher.kind == LauncherKind::Python);
assert!(launcher.python_path == python_executable_path);

Ok(())
}

Expand Down
Loading

0 comments on commit 6d494bb

Please sign in to comment.