1use std::io;
2use std::path::{Path, PathBuf};
3use std::str::Utf8Error;
4
5use fs_err::File;
6use thiserror::Error;
7
8use uv_fs::Simplified;
9
10#[cfg(all(windows, target_arch = "x86"))]
11const LAUNCHER_I686_GUI: &[u8] = include_bytes!("../trampolines/uv-trampoline-i686-gui.exe");
12
13#[cfg(all(windows, target_arch = "x86"))]
14const LAUNCHER_I686_CONSOLE: &[u8] =
15 include_bytes!("../trampolines/uv-trampoline-i686-console.exe");
16
17#[cfg(all(windows, target_arch = "x86_64"))]
18const LAUNCHER_X86_64_GUI: &[u8] = include_bytes!("../trampolines/uv-trampoline-x86_64-gui.exe");
19
20#[cfg(all(windows, target_arch = "x86_64"))]
21const LAUNCHER_X86_64_CONSOLE: &[u8] =
22 include_bytes!("../trampolines/uv-trampoline-x86_64-console.exe");
23
24#[cfg(all(windows, target_arch = "aarch64"))]
25const LAUNCHER_AARCH64_GUI: &[u8] = include_bytes!("../trampolines/uv-trampoline-aarch64-gui.exe");
26
27#[cfg(all(windows, target_arch = "aarch64"))]
28const LAUNCHER_AARCH64_CONSOLE: &[u8] =
29 include_bytes!("../trampolines/uv-trampoline-aarch64-console.exe");
30
31#[cfg(windows)]
33const RT_RCDATA: u16 = 10;
34
35#[cfg(windows)]
37const RESOURCE_TRAMPOLINE_KIND: windows::core::PCWSTR = windows::core::w!("UV_TRAMPOLINE_KIND");
38#[cfg(windows)]
39const RESOURCE_PYTHON_PATH: windows::core::PCWSTR = windows::core::w!("UV_PYTHON_PATH");
40#[cfg(windows)]
44const RESOURCE_SCRIPT_DATA: windows::core::PCWSTR = windows::core::w!("UV_SCRIPT_DATA");
45
46#[derive(Debug)]
47pub struct Launcher {
48 pub kind: LauncherKind,
49 pub python_path: PathBuf,
50 pub script_data: Option<Vec<u8>>,
51}
52
53impl Launcher {
54 #[cfg(not(windows))]
59 pub fn try_from_path(_path: &Path) -> Result<Option<Self>, Error> {
60 Ok(None)
61 }
62
63 #[cfg(windows)]
68 pub fn try_from_path(path: &Path) -> Result<Option<Self>, Error> {
69 use std::os::windows::ffi::OsStrExt;
70 use windows::Win32::System::LibraryLoader::LOAD_LIBRARY_AS_DATAFILE;
71 use windows::Win32::System::LibraryLoader::LoadLibraryExW;
72
73 let path_str = path
74 .as_os_str()
75 .encode_wide()
76 .chain(std::iter::once(0))
77 .collect::<Vec<_>>();
78
79 #[allow(unsafe_code)]
81 let Some(module) = (unsafe {
82 LoadLibraryExW(
83 windows::core::PCWSTR(path_str.as_ptr()),
84 None,
85 LOAD_LIBRARY_AS_DATAFILE,
86 )
87 .ok()
88 }) else {
89 return Ok(None);
90 };
91
92 let result = (|| {
93 let Some(kind_data) = read_resource(module, RESOURCE_TRAMPOLINE_KIND) else {
94 return Ok(None);
95 };
96 let Some(kind) = LauncherKind::from_resource_value(kind_data[0]) else {
97 return Err(Error::UnprocessableMetadata);
98 };
99
100 let Some(path_data) = read_resource(module, RESOURCE_PYTHON_PATH) else {
101 return Ok(None);
102 };
103 let python_path = PathBuf::from(
104 String::from_utf8(path_data).map_err(|err| Error::InvalidPath(err.utf8_error()))?,
105 );
106
107 let script_data = read_resource(module, RESOURCE_SCRIPT_DATA);
108
109 Ok(Some(Self {
110 kind,
111 python_path,
112 script_data,
113 }))
114 })();
115
116 #[allow(unsafe_code)]
118 unsafe {
119 windows::Win32::Foundation::FreeLibrary(module)
120 .map_err(|err| Error::Io(io::Error::from_raw_os_error(err.code().0)))?;
121 };
122
123 result
124 }
125
126 #[cfg(not(windows))]
131 pub fn write_to_file(self, _file: &mut File, _is_gui: bool) -> Result<(), Error> {
132 Err(Error::NotWindows)
133 }
134
135 #[cfg(windows)]
137 pub fn write_to_file(self, file: &mut File, is_gui: bool) -> Result<(), Error> {
138 use std::io::Write;
139 use uv_fs::Simplified;
140
141 let python_path = self.python_path.simplified_display().to_string();
142
143 let temp_dir = tempfile::TempDir::new()?;
145 let temp_file = temp_dir
146 .path()
147 .join(format!("uv-trampoline-{}.exe", std::process::id()));
148
149 fs_err::write(&temp_file, get_launcher_bin(is_gui)?)?;
151
152 let resources = &[
154 (
155 RESOURCE_TRAMPOLINE_KIND,
156 &[self.kind.to_resource_value()][..],
157 ),
158 (RESOURCE_PYTHON_PATH, python_path.as_bytes()),
159 ];
160 if let Some(script_data) = self.script_data {
161 let mut all_resources = resources.to_vec();
162 all_resources.push((RESOURCE_SCRIPT_DATA, &script_data));
163 write_resources(&temp_file, &all_resources)?;
164 } else {
165 write_resources(&temp_file, resources)?;
166 }
167
168 let launcher = fs_err::read(&temp_file)?;
170 fs_err::remove_file(&temp_file)?;
171
172 file.write_all(&launcher)?;
174
175 Ok(())
176 }
177
178 #[must_use]
179 pub fn with_python_path(self, path: PathBuf) -> Self {
180 Self {
181 kind: self.kind,
182 python_path: path,
183 script_data: self.script_data,
184 }
185 }
186}
187
188#[derive(Debug, Clone, Copy, PartialEq, Eq)]
192pub enum LauncherKind {
193 Script,
195 Python,
197}
198
199impl LauncherKind {
200 #[cfg(windows)]
201 fn to_resource_value(self) -> u8 {
202 match self {
203 Self::Script => 1,
204 Self::Python => 2,
205 }
206 }
207
208 #[cfg(windows)]
209 fn from_resource_value(value: u8) -> Option<Self> {
210 match value {
211 1 => Some(Self::Script),
212 2 => Some(Self::Python),
213 _ => None,
214 }
215 }
216}
217
218#[derive(Error, Debug)]
220pub enum Error {
221 #[error(transparent)]
222 Io(#[from] io::Error),
223 #[error("Failed to parse executable path")]
224 InvalidPath(#[source] Utf8Error),
225 #[error(
226 "Unable to create Windows launcher for: {0} (only x86_64, x86, and arm64 are supported)"
227 )]
228 UnsupportedWindowsArch(&'static str),
229 #[error("Unable to create Windows launcher on non-Windows platform")]
230 NotWindows,
231 #[error("Cannot process launcher metadata from resource")]
232 UnprocessableMetadata,
233 #[cfg(windows)]
234 #[error("Failed to write Windows launcher ZIP payload")]
235 AsyncZip(#[from] async_zip::error::ZipError),
236 #[error("Resources over 2^32 bytes are not supported")]
237 ResourceTooLarge,
238 #[error("Failed to update Windows PE resources: {}", path.user_display())]
239 WriteResources {
240 path: PathBuf,
241 #[source]
242 err: io::Error,
243 },
244}
245
246#[allow(clippy::unnecessary_wraps, unused_variables)]
247#[cfg(windows)]
248fn get_launcher_bin(gui: bool) -> Result<&'static [u8], Error> {
249 Ok(match std::env::consts::ARCH {
250 #[cfg(all(windows, target_arch = "x86"))]
251 "x86" => {
252 if gui {
253 LAUNCHER_I686_GUI
254 } else {
255 LAUNCHER_I686_CONSOLE
256 }
257 }
258 #[cfg(all(windows, target_arch = "x86_64"))]
259 "x86_64" => {
260 if gui {
261 LAUNCHER_X86_64_GUI
262 } else {
263 LAUNCHER_X86_64_CONSOLE
264 }
265 }
266 #[cfg(all(windows, target_arch = "aarch64"))]
267 "aarch64" => {
268 if gui {
269 LAUNCHER_AARCH64_GUI
270 } else {
271 LAUNCHER_AARCH64_CONSOLE
272 }
273 }
274 #[cfg(windows)]
275 arch => {
276 return Err(Error::UnsupportedWindowsArch(arch));
277 }
278 })
279}
280
281#[cfg(windows)]
283fn write_resources(path: &Path, resources: &[(windows::core::PCWSTR, &[u8])]) -> Result<(), Error> {
284 #[allow(unsafe_code)]
286 unsafe {
287 use std::os::windows::ffi::OsStrExt;
288 use windows::Win32::System::LibraryLoader::{
289 BeginUpdateResourceW, EndUpdateResourceW, UpdateResourceW,
290 };
291
292 let map_err = |err: windows::core::Error| Error::WriteResources {
293 path: path.to_path_buf(),
294 err: io::Error::from_raw_os_error(err.code().0),
295 };
296
297 let path_str = path
298 .as_os_str()
299 .encode_wide()
300 .chain(std::iter::once(0))
301 .collect::<Vec<_>>();
302 let handle = BeginUpdateResourceW(windows::core::PCWSTR(path_str.as_ptr()), false)
303 .map_err(map_err)?;
304
305 for (name, data) in resources {
306 UpdateResourceW(
307 handle,
308 windows::core::PCWSTR(RT_RCDATA as *const _),
309 *name,
310 0,
311 Some(data.as_ptr().cast()),
312 u32::try_from(data.len()).map_err(|_| Error::ResourceTooLarge)?,
313 )
314 .map_err(&map_err)?;
315 }
316
317 EndUpdateResourceW(handle, false).map_err(map_err)?;
318 }
319
320 Ok(())
321}
322
323#[cfg(windows)]
325fn read_resource(
326 handle: windows::Win32::Foundation::HMODULE,
327 name: windows::core::PCWSTR,
328) -> Option<Vec<u8>> {
329 #[allow(unsafe_code)]
331 unsafe {
332 use windows::Win32::System::LibraryLoader::{
333 FindResourceW, LoadResource, LockResource, SizeofResource,
334 };
335 let resource = FindResourceW(
337 Some(handle),
338 name,
339 windows::core::PCWSTR(RT_RCDATA as *const _),
340 );
341 if resource.is_invalid() {
342 return None;
343 }
344
345 let size = SizeofResource(Some(handle), resource);
347 if size == 0 {
348 return None;
349 }
350 let data = LoadResource(Some(handle), resource).ok()?;
351 let ptr = LockResource(data) as *const u8;
352 if ptr.is_null() {
353 return None;
354 }
355
356 Some(std::slice::from_raw_parts(ptr, size as usize).to_vec())
358 }
359}
360
361#[cfg(not(windows))]
366pub fn windows_script_launcher(
367 _launcher_python_script: &str,
368 _is_gui: bool,
369 _python_executable: impl AsRef<Path>,
370) -> Result<Vec<u8>, Error> {
371 Err(Error::NotWindows)
372}
373
374#[cfg(windows)]
381pub fn windows_script_launcher(
382 launcher_python_script: &str,
383 is_gui: bool,
384 python_executable: impl AsRef<Path>,
385) -> Result<Vec<u8>, Error> {
386 use async_zip::base::write::ZipFileWriter;
387 use async_zip::{Compression, ZipEntryBuilder};
388 use futures_lite::future::block_on;
389 use futures_lite::io::Cursor;
390
391 use uv_fs::Simplified;
392
393 let launcher_bin: &[u8] = get_launcher_bin(is_gui)?;
394
395 let mut archive = ZipFileWriter::new(Cursor::new(Vec::new()));
399 let entry = ZipEntryBuilder::new("__main__.py".to_string().into(), Compression::Stored);
400 block_on(archive.write_entry_whole(entry, launcher_python_script.as_bytes()))?;
401 let payload = block_on(archive.close())?.into_inner();
402
403 let python = python_executable.as_ref();
404 let python_path = python.simplified_display().to_string();
405
406 let temp_dir = tempfile::TempDir::new()?;
409 let temp_file = temp_dir
410 .path()
411 .join(format!("uv-trampoline-{}.exe", std::process::id()));
412 fs_err::write(&temp_file, launcher_bin)?;
413
414 let resources = &[
416 (
417 RESOURCE_TRAMPOLINE_KIND,
418 &[LauncherKind::Script.to_resource_value()][..],
419 ),
420 (RESOURCE_PYTHON_PATH, python_path.as_bytes()),
421 (RESOURCE_SCRIPT_DATA, &payload),
422 ];
423 write_resources(&temp_file, resources)?;
424
425 let launcher = fs_err::read(&temp_file)?;
430 fs_err::remove_file(temp_file)?;
431
432 Ok(launcher)
433}
434
435#[cfg(not(windows))]
440pub fn windows_python_launcher(
441 _python_executable: impl AsRef<Path>,
442 _is_gui: bool,
443) -> Result<Vec<u8>, Error> {
444 Err(Error::NotWindows)
445}
446
447#[cfg(windows)]
453pub fn windows_python_launcher(
454 python_executable: impl AsRef<Path>,
455 is_gui: bool,
456) -> Result<Vec<u8>, Error> {
457 use uv_fs::Simplified;
458
459 let launcher_bin: &[u8] = get_launcher_bin(is_gui)?;
460
461 let python = python_executable.as_ref();
462 let python_path = python.simplified_display().to_string();
463
464 let temp_dir = tempfile::TempDir::new()?;
466 let temp_file = temp_dir
467 .path()
468 .join(format!("uv-trampoline-{}.exe", std::process::id()));
469 fs_err::write(&temp_file, launcher_bin)?;
470
471 let resources = &[
473 (
474 RESOURCE_TRAMPOLINE_KIND,
475 &[LauncherKind::Python.to_resource_value()][..],
476 ),
477 (RESOURCE_PYTHON_PATH, python_path.as_bytes()),
478 ];
479 write_resources(&temp_file, resources)?;
480
481 let launcher = fs_err::read(&temp_file)?;
483 fs_err::remove_file(temp_file)?;
484
485 Ok(launcher)
486}
487
488#[cfg(all(test, windows))]
489#[expect(clippy::print_stdout)]
490mod test {
491 use std::io::Write;
492 use std::path::Path;
493 use std::path::PathBuf;
494 use std::process::Command;
495
496 use anyhow::Result;
497 use assert_cmd::prelude::OutputAssertExt;
498 use assert_fs::prelude::PathChild;
499 use fs_err::File;
500
501 use which::which;
502
503 use super::{Launcher, LauncherKind, windows_python_launcher, windows_script_launcher};
504
505 #[test]
506 #[cfg(all(windows, target_arch = "x86", feature = "production"))]
507 fn test_launchers_are_small() {
508 assert!(
510 super::LAUNCHER_I686_GUI.len() < 50 * 1024,
511 "GUI launcher: {}",
512 super::LAUNCHER_I686_GUI.len()
513 );
514 assert!(
515 super::LAUNCHER_I686_CONSOLE.len() < 50 * 1024,
516 "CLI launcher: {}",
517 super::LAUNCHER_I686_CONSOLE.len()
518 );
519 }
520
521 #[test]
522 #[cfg(all(windows, target_arch = "x86_64", feature = "production"))]
523 fn test_launchers_are_small() {
524 assert!(
526 super::LAUNCHER_X86_64_GUI.len() < 50 * 1024,
527 "GUI launcher: {}",
528 super::LAUNCHER_X86_64_GUI.len()
529 );
530 assert!(
531 super::LAUNCHER_X86_64_CONSOLE.len() < 50 * 1024,
532 "CLI launcher: {}",
533 super::LAUNCHER_X86_64_CONSOLE.len()
534 );
535 }
536
537 #[test]
538 #[cfg(all(windows, target_arch = "aarch64", feature = "production"))]
539 fn test_launchers_are_small() {
540 assert!(
542 super::LAUNCHER_AARCH64_GUI.len() < 50 * 1024,
543 "GUI launcher: {}",
544 super::LAUNCHER_AARCH64_GUI.len()
545 );
546 assert!(
547 super::LAUNCHER_AARCH64_CONSOLE.len() < 50 * 1024,
548 "CLI launcher: {}",
549 super::LAUNCHER_AARCH64_CONSOLE.len()
550 );
551 }
552
553 fn get_script_launcher(shebang: &str, is_gui: bool) -> String {
555 if is_gui {
556 format!(
557 r#"{shebang}
558# -*- coding: utf-8 -*-
559import re
560import sys
561
562def make_gui() -> None:
563 from tkinter import Tk, ttk
564 root = Tk()
565 root.title("uv Test App")
566 frm = ttk.Frame(root, padding=10)
567 frm.grid()
568 ttk.Label(frm, text="Hello from uv-trampoline-gui.exe").grid(column=0, row=0)
569 root.mainloop()
570
571if __name__ == "__main__":
572 sys.argv[0] = re.sub(r"(-script\.pyw|\.exe)?$", "", sys.argv[0])
573 sys.exit(make_gui())
574"#
575 )
576 } else {
577 format!(
578 r#"{shebang}
579# -*- coding: utf-8 -*-
580import re
581import sys
582
583def main_console() -> None:
584 print("Hello from uv-trampoline-console.exe", file=sys.stdout)
585 print("Hello from uv-trampoline-console.exe", file=sys.stderr)
586 for arg in sys.argv[1:]:
587 print(arg, file=sys.stderr)
588
589if __name__ == "__main__":
590 sys.argv[0] = re.sub(r"(-script\.pyw|\.exe)?$", "", sys.argv[0])
591 sys.exit(main_console())
592"#
593 )
594 }
595 }
596
597 fn format_shebang(executable: impl AsRef<Path>) -> String {
599 let executable = executable.as_ref().display().to_string();
601 format!("#!{executable}")
602 }
603
604 fn create_temp_certificate(temp_dir: &tempfile::TempDir) -> Result<(PathBuf, PathBuf)> {
606 use rcgen::{
607 CertificateParams, DnType, ExtendedKeyUsagePurpose, KeyPair, KeyUsagePurpose, SanType,
608 };
609
610 let mut params = CertificateParams::default();
611 params.key_usages.push(KeyUsagePurpose::DigitalSignature);
612 params
613 .extended_key_usages
614 .push(ExtendedKeyUsagePurpose::CodeSigning);
615 params
616 .distinguished_name
617 .push(DnType::OrganizationName, "Astral Software Inc.");
618 params
619 .distinguished_name
620 .push(DnType::CommonName, "uv-test-signer");
621 params
622 .subject_alt_names
623 .push(SanType::DnsName("uv-test-signer".try_into()?));
624
625 let private_key = KeyPair::generate()?;
626 let public_cert = params.self_signed(&private_key)?;
627
628 let public_cert_path = temp_dir.path().join("uv-trampoline-test.crt");
629 let private_key_path = temp_dir.path().join("uv-trampoline-test.key");
630 fs_err::write(public_cert_path.as_path(), public_cert.pem())?;
631 fs_err::write(private_key_path.as_path(), private_key.serialize_pem())?;
632
633 Ok((public_cert_path, private_key_path))
634 }
635
636 fn sign_authenticode(bin_path: impl AsRef<Path>) {
638 let temp_dir = tempfile::TempDir::new().expect("Failed to create temporary directory");
639 let (public_cert, private_key) =
640 create_temp_certificate(&temp_dir).expect("Failed to create self-signed certificate");
641
642 Command::new("pwsh")
644 .args([
645 "-NoProfile",
646 "-NonInteractive",
647 "-Command",
648 &format!(
649 r"
650 $ErrorActionPreference = 'Stop'
651 Import-Module Microsoft.PowerShell.Security
652 $cert = [System.Security.Cryptography.X509Certificates.X509Certificate2]::CreateFromPemFile('{}', '{}')
653 Set-AuthenticodeSignature -FilePath '{}' -Certificate $cert;
654 ",
655 public_cert.display().to_string().replace('\'', "''"),
656 private_key.display().to_string().replace('\'', "''"),
657 bin_path.as_ref().display().to_string().replace('\'', "''"),
658 ),
659 ])
660 .env_remove("PSModulePath")
661 .assert()
662 .success();
663
664 println!("Signed binary: {}", bin_path.as_ref().display());
665 }
666
667 #[test]
668 fn console_script_launcher() -> Result<()> {
669 let temp_dir = assert_fs::TempDir::new()?;
671 let console_bin_path = temp_dir.child("launcher.console.exe");
672
673 let python_executable_path = which("python")?;
675
676 let launcher_console_script =
678 get_script_launcher(&format_shebang(&python_executable_path), false);
679
680 let console_launcher =
682 windows_script_launcher(&launcher_console_script, false, &python_executable_path)?;
683
684 File::create(console_bin_path.path())?.write_all(console_launcher.as_ref())?;
686
687 println!(
688 "Wrote Console Launcher in {}",
689 console_bin_path.path().display()
690 );
691
692 let stdout_predicate = "Hello from uv-trampoline-console.exe\r\n";
693 let stderr_predicate = "Hello from uv-trampoline-console.exe\r\n";
694
695 #[cfg(windows)]
697 Command::new(console_bin_path.path())
698 .assert()
699 .success()
700 .stdout(stdout_predicate)
701 .stderr(stderr_predicate);
702
703 let args_to_test = vec!["foo", "bar", "foo bar", "foo \"bar\"", "foo 'bar'"];
704 let stderr_predicate = format!("{}{}\r\n", stderr_predicate, args_to_test.join("\r\n"));
705
706 Command::new(console_bin_path.path())
708 .args(args_to_test)
709 .assert()
710 .success()
711 .stdout(stdout_predicate)
712 .stderr(stderr_predicate);
713
714 let launcher = Launcher::try_from_path(console_bin_path.path())
715 .expect("We should succeed at reading the launcher")
716 .expect("The launcher should be valid");
717
718 assert!(launcher.kind == LauncherKind::Script);
719 assert!(launcher.python_path == python_executable_path);
720
721 sign_authenticode(console_bin_path.path());
723
724 let stdout_predicate = "Hello from uv-trampoline-console.exe\r\n";
725 let stderr_predicate = "Hello from uv-trampoline-console.exe\r\n";
726 Command::new(console_bin_path.path())
727 .assert()
728 .success()
729 .stdout(stdout_predicate)
730 .stderr(stderr_predicate);
731
732 Ok(())
733 }
734
735 #[test]
736 fn console_python_launcher() -> Result<()> {
737 let temp_dir = assert_fs::TempDir::new()?;
739 let console_bin_path = temp_dir.child("launcher.console.exe");
740
741 let python_executable_path = which("python")?;
743
744 let console_launcher = windows_python_launcher(&python_executable_path, false)?;
746
747 {
749 File::create(console_bin_path.path())?.write_all(console_launcher.as_ref())?;
750 }
751
752 println!(
753 "Wrote Python Launcher in {}",
754 console_bin_path.path().display()
755 );
756
757 Command::new(console_bin_path.path())
759 .arg("-c")
760 .arg("print('Hello from Python Launcher')")
761 .assert()
762 .success()
763 .stdout("Hello from Python Launcher\r\n");
764
765 let launcher = Launcher::try_from_path(console_bin_path.path())
766 .expect("We should succeed at reading the launcher")
767 .expect("The launcher should be valid");
768
769 assert!(launcher.kind == LauncherKind::Python);
770 assert!(launcher.python_path == python_executable_path);
771
772 sign_authenticode(console_bin_path.path());
774 Command::new(console_bin_path.path())
775 .arg("-c")
776 .arg("print('Hello from Python Launcher')")
777 .assert()
778 .success()
779 .stdout("Hello from Python Launcher\r\n");
780
781 Ok(())
782 }
783
784 #[test]
785 #[ignore = "This test will spawn a GUI and wait until you close the window."]
786 fn gui_launcher() -> Result<()> {
787 let temp_dir = assert_fs::TempDir::new()?;
789 let gui_bin_path = temp_dir.child("launcher.gui.exe");
790
791 let pythonw_executable_path = which("pythonw")?;
793
794 let launcher_gui_script =
796 get_script_launcher(&format_shebang(&pythonw_executable_path), true);
797
798 let gui_launcher =
800 windows_script_launcher(&launcher_gui_script, true, &pythonw_executable_path)?;
801
802 {
804 File::create(gui_bin_path.path())?.write_all(gui_launcher.as_ref())?;
805 }
806
807 println!("Wrote GUI Launcher in {}", gui_bin_path.path().display());
808
809 Command::new(gui_bin_path.path()).assert().success();
812
813 Ok(())
814 }
815}