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