Win Trampoline: Use Python executable path encoded in binary (#1803)

.github/workflows/ci.yml

@@ -2,7 +2,7 @@ name: CI
 
 on:
   push:
-    branches: [main]
+    branches: [ main ]
   pull_request:
   workflow_dispatch:
 
@@ -32,7 +32,7 @@ jobs:
     name: "cargo clippy"
     strategy:
       matrix:
-        os: [ubuntu-latest, windows-latest]
+        os: [ ubuntu-latest, windows-latest ]
     runs-on: ${{ matrix.os }}
     steps:
       - uses: actions/checkout@v4
@@ -113,7 +113,7 @@ jobs:
           workspaces: "crates/uv-trampoline"
       - name: "Clippy"
         working-directory: crates/uv-trampoline
-        run: cargo clippy --all-features --locked -- -D warnings
+        run: cargo clippy --all-features --locked --target x86_64-pc-windows-msvc -- -D warnings
       - name: "Build"
         working-directory: crates/uv-trampoline
-        run: cargo build --release -Z build-std=core,panic_abort,alloc -Z build-std-features=compiler-builtins-mem --target x86_64-pc-windows-msvc
+        run: cargo build --release --target x86_64-pc-windows-msvc

crates/install-wheel-rs/src/wheel.rs

@@ -33,6 +33,8 @@ use crate::{find_dist_info, Error};
 /// `#!/usr/bin/env python`
 pub const SHEBANG_PYTHON: &str = "#!/usr/bin/env python";
 
+const LAUNCHER_MAGIC_NUMBER: [u8; 4] = [b'U', b'V', b'U', b'V'];
+
 #[cfg(all(windows, target_arch = "x86_64"))]
 const LAUNCHER_X86_64_GUI: &[u8] =
     include_bytes!("../../uv-trampoline/trampolines/uv-trampoline-x86_64-gui.exe");
@@ -281,19 +283,7 @@ fn unpack_wheel_files<R: Read + Seek>(
 }
 
 fn get_shebang(location: &InstallLocation<impl AsRef<Path>>) -> String {
-    let path = location.python().to_string_lossy().to_string();
-    let path = if cfg!(windows) {
-        // https://stackoverflow.com/a/50323079
-        const VERBATIM_PREFIX: &str = r"\\?\";
-        if let Some(stripped) = path.strip_prefix(VERBATIM_PREFIX) {
-            stripped.to_string()
-        } else {
-            path
-        }
-    } else {
-        path
-    };
-    format!("#!{path}")
+    format!("#!{}", location.python().normalized().display())
 }
 
 /// A Windows script is a minimal .exe launcher binary with the python entrypoint script appended as
@@ -305,6 +295,7 @@ fn get_shebang(location: &InstallLocation<impl AsRef<Path>>) -> String {
 pub(crate) fn windows_script_launcher(
     launcher_python_script: &str,
     is_gui: bool,
+    installation: &InstallLocation<impl AsRef<Path>>,
 ) -> Result<Vec<u8>, Error> {
     // This method should only be called on Windows, but we avoid `#[cfg(windows)]` to retain
     // compilation on all platforms.
@@ -352,9 +343,20 @@ pub(crate) fn windows_script_launcher(
         archive.finish().expect(error_msg);
     }
 
+    let python = installation.python();
+    let python_path = python.normalized().to_string_lossy();
+
     let mut launcher: Vec<u8> = Vec::with_capacity(launcher_bin.len() + payload.len());
     launcher.extend_from_slice(launcher_bin);
     launcher.extend_from_slice(&payload);
+    launcher.extend_from_slice(python_path.as_bytes());
+    launcher.extend_from_slice(
+        &u32::try_from(python_path.as_bytes().len())
+            .expect("File Path to be smaller than 4GB")
+            .to_le_bytes(),
+    );
+    launcher.extend_from_slice(&LAUNCHER_MAGIC_NUMBER);
+
     Ok(launcher)
 }
 
@@ -393,7 +395,7 @@ pub(crate) fn write_script_entrypoints(
             write_file_recorded(
                 site_packages,
                 &entrypoint_relative,
-                &windows_script_launcher(&launcher_python_script, is_gui)?,
+                &windows_script_launcher(&launcher_python_script, is_gui, location)?,
                 record,
             )?;
         } else {
@@ -949,7 +951,7 @@ pub fn parse_key_value_file(
 ///
 /// Wheel 1.0: <https://www.python.org/dev/peps/pep-0427/>
 #[allow(clippy::too_many_arguments)]
-#[instrument(skip_all, fields(name = %filename.name))]
+#[instrument(skip_all, fields(name = % filename.name))]
 pub fn install_wheel(
     location: &InstallLocation<LockedDir>,
     reader: impl Read + Seek,

crates/uv-trampoline/.cargo/config.toml

@@ -0,0 +1,3 @@
+[unstable]
+build-std = ["core", "panic_abort", "alloc", "std"]
+build-std-features = ["compiler-builtins-mem"]

crates/uv-trampoline/Cargo.toml

@@ -43,15 +43,7 @@ windows-sys = { version = "0.52.0", features = [
     "Win32_System_WindowsProgramming",
     "Win32_UI_WindowsAndMessaging",
 ] }
-# This provides implementations of memcpy, memset, etc., which the compiler assumes
-# are available. But there's also a hidden copy of this crate inside `core`/`alloc`,
-# and they may or may not conflict depending on how clever the linker is feeling.
-# The issue is that the hidden copy doesn't have the "mem" feature enabled, and we
-# need it. So two options:
-# - Uncomment this, and cross fingers that it doesn't cause conflicts
-# - Use -Zbuild-std=... -Zbuild-std-features=compiler-builtins-mem, which enables
-#   the mem feature on the built-in builtins.
-#compiler_builtins = { version = "*", features = ["mem"]}
+
 ufmt-write = "0.1.0"
 ufmt = "0.2.0"
 

crates/uv-trampoline/README.md

@@ -1,6 +1,7 @@
 # Windows trampolines
 
-This is a fork of [posy trampolines](https://github.com/njsmith/posy/tree/dda22e6f90f5fefa339b869dd2bbe107f5b48448/src/trampolines/windows-trampolines/posy-trampoline).
+This is a fork
+of [posy trampolines](https://github.com/njsmith/posy/tree/dda22e6f90f5fefa339b869dd2bbe107f5b48448/src/trampolines/windows-trampolines/posy-trampoline).
 
 # What is this?
 
@@ -13,27 +14,33 @@ That's what this does: it's a generic "trampoline" that lets us generate custom
 `.exe`s for arbitrary Python scripts, and when invoked it bounces to invoking
 `python <the script>` instead.
 
-
 # How do you use it?
 
 Basically, this looks up `python.exe` (for console programs) or
 `pythonw.exe` (for GUI programs) in the adjacent directory, and invokes
 `python[w].exe path\to\the\<the .exe>`.
 
-The intended use is: take your Python script, name it `__main__.py`, and pack it
-into a `.zip` file. Then concatenate that `.zip` file onto the end of one of our
-prebuilt `.exe`s.
+The intended use is:
+
+* take your Python script, name it `__main__.py`, and pack it
+  into a `.zip` file. Then concatenate that `.zip` file onto the end of one of our
+  prebuilt `.exe`s.
+* After the zip file content, write the path to the Python executable that the script uses to run
+  the Python script as UTF-8 encoded string, followed by the path's length as a 32-bit little-endian
+  integer.
+* At the very end, write the magic number `UVUV` in bytes.
+
+|       `launcher.exe`        |
+|:---------------------------:|
+|  `<zipped python script>`   |
+|   `<path to python.exe>`    |
+| `<len(path to python.exe)>` |
+| `<b'U', b'V', b'U', b'V'>`  |
 
 Then when you run `python` on the `.exe`, it will see the `.zip` trailer at the
 end of the `.exe`, and automagically look inside to find and execute
 `__main__.py`. Easy-peasy.
 
-(TODO: we should probably make the Python-finding logic slightly more flexible
-at some point -- in particular to support more conventional venv-style
-installation where you find `python` by looking in the directory next to the
-trampoline `.exe` -- but this is good enough to get started.)
-
-
 # Why does this exist?
 
 I probably could have used Vinay's C++ implementation from `distlib`, but what's
@@ -47,7 +54,6 @@ Python-finding logic we want. But mostly it was just an interesting challenge.
 This does owe a *lot* to the `distlib` implementation though. The overall logic
 is copied more-or-less directly.
 
-
 # Anything I should know for hacking on this?
 
 In order to minimize binary size, this uses `#![no_std]`, `panic="abort"`, and
@@ -64,7 +70,7 @@ this:
   Though uh, this does mean that literally all of our code is `unsafe`. Sorry!
 
 - `runtime.rs` has the core glue to get panicking, heap allocation, and linking
-  working. 
+  working.
 
 - `diagnostics.rs` uses `ufmt` and some cute Windows tricks to get a convenient
   version of `eprintln!` that works without `std`, and automatically prints to
@@ -85,7 +91,6 @@ Miscellaneous tips:
   `.unwrap_unchecked()` avoids this. Similar for `slice[idx]` vs
   `slice.get_unchecked(idx)`.
 
-
 # How do you build this stupid thing?
 
 Building this can be frustrating, because the low-level compiler/runtime
@@ -99,15 +104,8 @@ might not realize that, and still emit references to the unwinding helper
 `__CxxFrameHandler3`. And then the linker blows up because that symbol doesn't
 exist.
 
-Two approaches that are reasonably likely to work:
-
-- Uncomment `compiler-builtins` in `Cargo.toml`, and build normally: `cargo
-  build --profile release`.
-
-- Leave `compiler-builtins` commented-out, and build like: `cargo build
-  --release -Z build-std=core,panic_abort,alloc -Z
-  build-std-features=compiler-builtins-mem --target x86_64-pc-windows-msvc`
-
+`cargo build --release --target x86_64-pc-windows-msvc`
+or `cargo build --release --target aarch64-pc-windows-msvc`
 
 Hopefully in the future as `#![no_std]` develops, this will get smoother.
 
@@ -125,6 +123,6 @@ rustup target add aarch64-pc-windows-msvc
 ```
 
 ```shell
-cargo +nightly xwin build --release -Z build-std=core,panic_abort,alloc -Z build-std-features=compiler-builtins-mem --target x86_64-pc-windows-msvc
-cargo +nightly xwin build --release -Z build-std=core,panic_abort,alloc -Z build-std-features=compiler-builtins-mem --target aarch64-pc-windows-msvc
+cargo +nightly xwin build --release --target x86_64-pc-windows-msvc
+cargo +nightly xwin build --release --target aarch64-pc-windows-msvc
 ```