New trampoline

crates/rattler-bin/src/commands/create.rs

@@ -26,7 +26,6 @@ use rattler_solve::{
     resolvo, SolverImpl, SolverTask,
 };
 use reqwest::Client;
-use std::sync::Arc;
 use std::{
     borrow::Cow,
     env,
@@ -36,6 +35,7 @@ use std::{
     str::FromStr,
     time::Duration,
 };
+use std::{fs, sync::Arc};
 use tokio::task::JoinHandle;
 
 #[derive(Debug, clap::Parser)]
@@ -60,12 +60,23 @@ pub struct Opt {
 
     #[clap(long)]
     timeout: Option<u64>,
+
+    #[clap(long)]
+    target_prefix: Option<PathBuf>,
 }
 
 pub async fn create(opt: Opt) -> anyhow::Result<()> {
     let channel_config = ChannelConfig::default();
-    let target_prefix = env::current_dir()?.join(".prefix");
-
+    let target_prefix = opt.target_prefix.unwrap_or_else(|| {
+        env::current_dir()
+            .expect("could not find current dir")
+            .join(".prefix")
+    });
+    // First create the target prefix if it doesn't exist yet.
+    fs::create_dir_all(&target_prefix).expect("could not create target prefix");
+    // Canonicalize the target prefix so we can use it in the future.
+    let target_prefix =
+        fs::canonicalize(target_prefix).context("could not canonicalize target prefix")?;
     // Determine the platform we're going to install for
     let install_platform = if let Some(platform) = opt.platform {
         Platform::from_str(&platform)?

crates/rattler/Cargo.toml

@@ -48,6 +48,8 @@ tokio-stream = { workspace = true, features = ["sync"] }
 tracing = { workspace = true }
 url = { workspace = true, features = ["serde"] }
 uuid = { workspace = true, features = ["v4", "fast-rng"] }
+zip.workspace = true
+dunce.workspace = true
 
 [dev-dependencies]
 assert_matches = { workspace = true }

crates/rattler/src/install/entry_point.rs

@@ -7,18 +7,25 @@ use rattler_conda_types::{
 };
 use rattler_digest::HashingWriter;
 use rattler_digest::Sha256;
-use std::{fs::File, io, io::Write, path::Path};
+use std::{
+    fs::{self, File},
+    io::{self, Cursor, Write},
+    path::{Path, PathBuf},
+};
+use zip::{write::FileOptions, ZipWriter};
 
 /// Get the bytes of the windows launcher executable.
 pub fn get_windows_launcher(platform: &Platform) -> &'static [u8] {
     match platform {
         Platform::Win32 => unimplemented!("32 bit windows is not supported for entry points"),
-        Platform::Win64 => include_bytes!("../../resources/launcher64.exe"),
-        Platform::WinArm64 => unimplemented!("arm64 windows is not supported for entry points"),
+        Platform::Win64 => include_bytes!("../../resources/uv-trampoline-x86_64-console.exe"),
+        Platform::WinArm64 => include_bytes!("../../resources/uv-trampoline-aarch64-console.exe"),
         _ => panic!("unsupported platform"),
     }
 }
 
+const LAUNCHER_MAGIC_NUMBER: [u8; 4] = [b'U', b'V', b'U', b'V'];
+
 /// Creates an "entry point" on disk for a Python entrypoint. Entrypoints are executable files that
 /// directly call a certain Python function.
 ///
@@ -39,21 +46,8 @@ pub fn create_windows_python_entry_point(
     entry_point: &EntryPoint,
     python_info: &PythonInfo,
     target_platform: &Platform,
-) -> Result<[PathsEntry; 2], std::io::Error> {
-    // Construct the path to where we will be creating the python entry point script.
-    let relative_path_script_py = python_info
-        .bin_dir
-        .join(format!("{}-script.py", &entry_point.command));
-
-    // Write the contents of the launcher script to disk
-    let script_path = target_dir.join(&relative_path_script_py);
-    std::fs::create_dir_all(
-        script_path
-            .parent()
-            .expect("since we joined with target_dir there must be a parent"),
-    )?;
+) -> Result<[PathsEntry; 1], std::io::Error> {
     let script_contents = python_entry_point_template(target_prefix, entry_point, python_info);
-    let (hash, size) = write_and_hash(&script_path, script_contents)?;
 
     // Construct a path to where we will create the python launcher executable.
     let relative_path_script_exe = python_info
@@ -62,34 +56,50 @@ pub fn create_windows_python_entry_point(
 
     // Include the bytes of the launcher directly in the binary so we can write it to disk.
     let launcher_bytes = get_windows_launcher(target_platform);
-    std::fs::write(target_dir.join(&relative_path_script_exe), launcher_bytes)?;
 
-    let fixed_launcher_digest = rattler_digest::parse_digest_from_hex::<rattler_digest::Sha256>(
-        "28b001bb9a72ae7a24242bfab248d767a1ac5dec981c672a3944f7a072375e9a",
-    )
-    .unwrap();
-
-    Ok([
-        PathsEntry {
-            relative_path: relative_path_script_py,
-            // todo: clobbering of entry points not handled yet
-            original_path: None,
-            path_type: PathType::WindowsPythonEntryPointScript,
-            no_link: false,
-            sha256: Some(hash),
-            sha256_in_prefix: None,
-            size_in_bytes: Some(size as _),
-        },
-        PathsEntry {
-            relative_path: relative_path_script_exe,
-            original_path: None,
-            path_type: PathType::WindowsPythonEntryPointExe,
-            no_link: false,
-            sha256: Some(fixed_launcher_digest),
-            sha256_in_prefix: None,
-            size_in_bytes: Some(launcher_bytes.len() as u64),
-        },
-    ])
+    let mut payload: Vec<u8> = Vec::new();
+    {
+        // We're using the zip writer, but with stored compression
+        // https://github.com/njsmith/posy/blob/04927e657ca97a5e35bb2252d168125de9a3a025/src/trampolines/mod.rs#L75-L82
+        // https://github.com/pypa/distlib/blob/8ed03aab48add854f377ce392efffb79bb4d6091/PC/launcher.c#L259-L271
+        let stored = FileOptions::default().compression_method(zip::CompressionMethod::Stored);
+        let mut archive = ZipWriter::new(Cursor::new(&mut payload));
+        let error_msg = "Writing to Vec<u8> should never fail";
+        archive.start_file("__main__.py", stored).expect(error_msg);
+        archive
+            .write_all(script_contents.as_bytes())
+            .expect(error_msg);
+        archive.finish().expect(error_msg);
+    }
+
+    let python = PathBuf::from(target_prefix).join(&python_info.path);
+    let python_path = dunce::simplified(&python).display().to_string();
+    println!("Python path: {}", python_path);
+
+    let mut launcher: Vec<u8> = Vec::with_capacity(launcher_bytes.len() + payload.len());
+    launcher.extend_from_slice(launcher_bytes);
+    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);
+
+    let target_location = target_dir.join(&relative_path_script_exe);
+    fs::create_dir_all(target_location.parent().unwrap())?;
+    let (sha256, size) = write_and_hash(&target_location, launcher)?;
+
+    Ok([PathsEntry {
+        relative_path: relative_path_script_exe,
+        original_path: None,
+        path_type: PathType::WindowsPythonEntryPointExe,
+        no_link: false,
+        sha256: Some(sha256),
+        sha256_in_prefix: None,
+        size_in_bytes: Some(size as u64),
+    }])
 }
 
 /// Creates an "entry point" on disk for a Python entrypoint. Entrypoints are executable files that

crates/rattler/src/install/mod.rs

@@ -386,9 +386,8 @@ pub async fn link_package(
                         &python_info,
                         &platform,
                     ) {
-                        Ok([a, b]) => {
+                        Ok([a]) => {
                             let _ = tx.blocking_send(Ok((number_of_paths_entries, a)));
-                            let _ = tx.blocking_send(Ok((number_of_paths_entries + 1, b)));
                         }
                         Err(e) => {
                             let _ = tx.blocking_send(Err(
@@ -397,7 +396,7 @@ pub async fn link_package(
                         }
                     }
                 });
-                number_of_paths_entries += 2;
+                number_of_paths_entries += 1;
             } else {
                 driver.spawn_throttled_and_forget(move || {
                     // Return immediately if the receiver was closed. This can happen if a previous step