Use trampolines for Python executables on Windows

zanieb · zanieb · commit d9fdcd202a7b · 2024-10-29T15:37:41.000-05:00
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
@@ -297,7 +297,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 }}
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/crates/uv-python/Cargo.toml b/crates/uv-python/Cargo.toml
@@ -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 }
diff --git a/crates/uv-python/src/managed.rs b/crates/uv-python/src/managed.rs
@@ -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::{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::{
@@ -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)]
@@ -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())]
@@ -485,16 +489,52 @@ impl ManagedPythonInstallation {
             err,
         })?;
 
-        match uv_fs::symlink_copy_fallback_file(&python, target) {
-            Ok(()) => Ok(()),
-            Err(err) if err.kind() == io::ErrorKind::NotFound => {
-                Err(Error::MissingExecutable(python.clone()))
+        if cfg!(unix) {
+            match uv_fs::symlink_copy_fallback_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) 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 {
+            unimplemented!("Only Windows and Unix systems are supported.")
+        }
+    }
+
+    /// 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: target.to_path_buf(),
-                err,
-            }),
+            launcher.python_path == self.executable()
+        } else {
+            unreachable!("Only Windows and Unix are supported")
         }
     }
 }
diff --git a/crates/uv-trampoline-builder/Cargo.toml b/crates/uv-trampoline-builder/Cargo.toml
@@ -23,6 +23,8 @@ workspace = true
 
 [dependencies]
 uv-fs = { workspace = true }
+
+fs-err = {workspace = true }
 thiserror = { workspace = true }
 zip = { workspace = true }
 
diff --git a/crates/uv-trampoline-builder/src/lib.rs b/crates/uv-trampoline-builder/src/lib.rs
@@ -1,6 +1,7 @@
-use std::io::{Cursor, Write};
-use std::path::Path;
+use std::io::{self, Cursor, Read, Seek, Write};
+use std::path::{Path, PathBuf};
 
+use fs_err::File;
 use thiserror::Error;
 use uv_fs::Simplified;
 use zip::write::FileOptions;
@@ -30,28 +31,160 @@ 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 Ok(_) = file.seek(io::SeekFrom::End(
+            -((MAGIC_NUMBER_SIZE + PATH_LENGTH_SIZE) as i64),
+        )) else {
+            return Err(Error::InvalidLauncher(
+                "Unable to seek to the start of the path length".to_string(),
+            ));
+        };
+
+        // 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::InvalidLauncher(format!(
+                    "Only paths with a length up to 32KBs are supported but the Python executable path has a length of {raw_length}"
+                )));
+            }
+
+            // SAFETY: Above we guarantee the length is less than 32KB
+            raw_length as usize
+        };
+
+        // Seek to the start of the path
+        let Ok(_) = file.seek(io::SeekFrom::End(
+            -((MAGIC_NUMBER_SIZE + PATH_LENGTH_SIZE + path_length) as i64),
+        )) else {
+            return Err(Error::InvalidLauncher(
+                "Unable to seek to the start of the path".to_string(),
+            ));
+        };
+
+        // 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(|_| {
+            Error::InvalidLauncher("Python executable path was not valid UTF-8".to_string())
+        })?);
+
+        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("Invalid launcher: {0}")]
+    InvalidLauncher(String),
+    #[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)"
     )]
@@ -192,7 +325,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"))]
@@ -340,6 +473,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(())
     }
 
@@ -371,6 +511,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(())
     }
 
diff --git a/crates/uv/src/commands/python/uninstall.rs b/crates/uv/src/commands/python/uninstall.rs
@@ -7,7 +7,6 @@ use futures::StreamExt;
 use itertools::Itertools;
 use owo_colors::OwoColorize;
 
-use same_file::is_same_file;
 use tracing::{debug, warn};
 use uv_fs::Simplified;
 use uv_python::downloads::PythonDownloadRequest;
@@ -149,9 +148,9 @@ async fn do_uninstall(
         })
         // Only include Python executables that match the installations
         .filter(|path| {
-            matching_installations.iter().any(|installation| {
-                is_same_file(path, installation.executable()).unwrap_or_default()
-            })
+            matching_installations
+                .iter()
+                .any(|installation| installation.is_bin_link(path.as_path()))
         })
         .collect::<BTreeSet<_>>();
 
@@ -218,6 +217,7 @@ async fn do_uninstall(
             .sorted_unstable_by(|a, b| a.key.cmp(&b.key).then_with(|| a.kind.cmp(&b.kind)))
         {
             match event.kind {
+                // TODO(zanieb): Track removed executables and report them all here
                 ChangeEventKind::Removed => {
                     writeln!(
                         printer.stderr(),
diff --git a/crates/uv/tests/it/common/mod.rs b/crates/uv/tests/it/common/mod.rs
@@ -663,7 +663,15 @@ impl TestContext {
             .arg("python")
             .arg("install")
             .env(EnvVars::UV_PYTHON_INSTALL_DIR, managed)
-            .env(EnvVars::UV_PYTHON_BIN_DIR, bin)
+            .env(EnvVars::UV_PYTHON_BIN_DIR, bin.as_os_str())
+            .env(
+                EnvVars::PATH,
+                std::env::join_paths(
+                    std::iter::once(bin)
+                        .chain(std::env::split_paths(&env::var("PATH").unwrap_or_default())),
+                )
+                .unwrap(),
+            )
             .current_dir(&self.temp_dir);
         command
     }
diff --git a/crates/uv/tests/it/python_install.rs b/crates/uv/tests/it/python_install.rs
