Skip to main content

cargo_zng/res/built_in/sfx_res/
sfx_main.rs

1// <windows-subsystem>
2#![windows_subsystem = "windows"]
3// </windows-subsystem>
4
5// source code for the generated self-extracting executable compiled by ./sfx.rs
6
7// [(name, compression, parts)]
8// data split in parts for performance and to avoid hard limits
9static DATA: &[(&str, Compression, &[&[u8]])] = &[];
10
11static ARGS: &[&str] = &[];
12static ENV: &[(&str, &str)] = &[];
13
14#[derive(Clone, Copy)]
15pub enum Compression {
16    None,
17    Zstd,
18    ZstdBcj(BcjFilter),
19}
20
21#[derive(Clone, Copy)]
22pub enum BcjFilter {
23    X86,
24    Arm,
25    Arm64,
26    ArmThumb,
27    Ppc,
28    Sparc,
29    Ia64,
30    Riscv,
31}
32
33use std::{
34    env,
35    fmt::Write as _,
36    fs,
37    io::{self, Read as _},
38    path::PathBuf,
39};
40
41use lzma_rust2::filter::bcj::BcjReader;
42
43pub fn main() {
44    // <windows-subsystem>
45    #[cfg(windows)]
46    attach_console();
47    // </windows-subsystem>
48
49    if let Ok(name) = std::env::var("SFX_GET_DATA") {
50        return serve_data(&name);
51    }
52    run();
53}
54
55macro_rules! err_exit {
56    ($($msg:tt)*) => {
57        {
58            eprintln!($($msg)*);
59            std::process::exit(-1);
60        }
61    };
62}
63trait UnwrapOrExit<T> {
64    fn unwrap_or_exit(self, ctx: &str) -> T;
65}
66impl<T, E: std::error::Error> UnwrapOrExit<T> for Result<T, E> {
67    fn unwrap_or_exit(self, ctx: &str) -> T {
68        match self {
69            Ok(r) => r,
70            Err(e) => err_exit!("{ctx} error, {e}"),
71        }
72    }
73}
74
75fn read_data(name: &str) -> Box<dyn io::Read> {
76    for (key, compression, parts) in DATA {
77        if *key == name {
78            if parts.is_empty() {
79                return Box::new(&[0u8; 0][..]);
80            }
81            let mut data: Box<dyn io::Read> = Box::new(parts[0]);
82            for &part in parts[1..].iter() {
83                data = Box::new(data.chain(part));
84            }
85            match *compression {
86                Compression::None => {}
87                Compression::Zstd => {
88                    let d = zstd::stream::read::Decoder::new(data).unwrap_or_exit("Decoder::new");
89                    data = Box::new(d);
90                }
91                Compression::ZstdBcj(bcj) => {
92                    let d = zstd::stream::read::Decoder::new(data).unwrap_or_exit("Decoder::new");
93                    let d = match bcj {
94                        BcjFilter::X86 => BcjReader::new_x86(d, 0),
95                        BcjFilter::Arm => BcjReader::new_arm(d, 0),
96                        BcjFilter::Arm64 => BcjReader::new_arm64(d, 0),
97                        BcjFilter::ArmThumb => BcjReader::new_arm_thumb(d, 0),
98                        BcjFilter::Ppc => BcjReader::new_ppc(d, 0),
99                        BcjFilter::Sparc => BcjReader::new_sparc(d, 0),
100                        BcjFilter::Ia64 => BcjReader::new_ia64(d, 0),
101                        BcjFilter::Riscv => BcjReader::new_riscv(d, 0),
102                    };
103                    data = Box::new(d);
104                }
105            }
106            return data;
107        }
108    }
109    err_exit!("{name:?} not found")
110}
111
112fn serve_data(name: &str) {
113    let mut data = read_data(name);
114    io::copy(&mut data, &mut std::io::stdout()).unwrap_or_exit("serve_data/copy");
115}
116
117fn run() {
118    let run_file = tmp_run_file().unwrap_or_exit("tmp_run_file");
119    {
120        let mut data = read_data(":run");
121        let mut run_file = fs::File::create_new(&run_file).unwrap_or_exit("File::create_new");
122        io::copy(&mut data, &mut run_file).unwrap_or_exit("run/copy");
123    }
124    #[cfg(unix)]
125    {
126        use std::os::unix::fs::PermissionsExt as _;
127        let perms = fs::metadata(std::env::current_exe().unwrap())
128            .unwrap_or_exit("get-metadata")
129            .permissions();
130        fs::set_permissions(&run_file, perms).unwrap_or_exit("set-metadata");
131    }
132
133    let mut sfx_args = String::new();
134    let mut sep = "";
135    for arg in env::args() {
136        write!(&mut sfx_args, "{sep}{arg}").unwrap_or_exit("");
137        sep = "\n";
138    }
139
140    let mut run = std::process::Command::new(&run_file);
141    for arg in ARGS {
142        run.arg(arg);
143    }
144    for (key, value) in ENV {
145        run.env(key, value);
146    }
147    let s = run.env("SFX_ARGS", sfx_args).status().unwrap_or_exit("run");
148
149    if let Err(e) = fs::remove_file(run_file)
150        && !matches!(e.kind(), io::ErrorKind::NotFound)
151    {
152        err_exit!("run/remove error, {e:?}");
153    }
154    if !s.success() {
155        std::process::exit(s.code().unwrap_or(-1))
156    }
157}
158fn tmp_run_file() -> io::Result<PathBuf> {
159    let tmp = env::temp_dir().join("zng-sfx");
160    if let Err(e) = fs::create_dir(&tmp)
161        && !matches!(e.kind(), io::ErrorKind::AlreadyExists)
162    {
163        return Err(e);
164    }
165    for i in 0..1000 {
166        let tmp = tmp.join(format!("run-{i}.exe"));
167        if let Err(e) = fs::remove_file(&tmp)
168            && matches!(e.kind(), io::ErrorKind::NotFound)
169        {
170            return Ok(tmp);
171        }
172    }
173    Err(io::Error::new(io::ErrorKind::QuotaExceeded, "too many tmp exe"))
174}
175
176// <windows-subsystem>
177#[cfg(windows)]
178pub fn attach_console() {
179    #[link(name = "kernel32")]
180    unsafe extern "system" {
181        fn GetConsoleWindow() -> isize;
182        fn AttachConsole(process_id: u32) -> i32;
183    }
184    unsafe {
185        // If no console is attached, attempt to attach to parent
186        if GetConsoleWindow() == 0 {
187            let _ = AttachConsole(0xFFFFFFFF);
188        }
189    }
190}
191// </windows-subsystem>