zoi_lua/api/
system.rs

1use std::fmt::Write as _;
2use std::io::{BufRead, BufReader, Write as _};
3use std::process::{Child, ChildStdin, ChildStdout, Command, Stdio};
4use std::sync::{Arc, Mutex};
5use std::thread;
6
7use mlua::{self, Lua, Table, Value};
8
9/// Represents an active shell process used for executing commands from Lua.
10struct InnerSession {
11    /// The handle to the child shell process.
12    child: Child,
13    /// The standard input stream for the child process.
14    stdin: ChildStdin,
15    /// The buffered standard output stream for the child process.
16    stdout: BufReader<ChildStdout>,
17    /// The shared buffer for capturing standard error.
18    stderr_buffer: Arc<Mutex<String>>,
19    /// A unique string used to identify the end of a command's output.
20    sentinel: String
21}
22
23impl Drop for InnerSession {
24    fn drop(&mut self) {
25        let _ = self.child.kill();
26        let _ = self.child.wait();
27    }
28}
29
30/// A wrapper around `InnerSession` providing thread-safe access.
31struct ShellSession {
32    /// The inner session state, protected by a mutex.
33    inner: Mutex<InnerSession>
34}
35
36impl ShellSession {
37    /// Creates a new `ShellSession` within the specified build directory.
38    fn new(lua: &Lua, build_dir: &str) -> Result<Self, anyhow::Error> {
39        let sentinel =
40            format!("---ZOI_CMD_COMPLETE_{}---", uuid::Uuid::new_v4());
41
42        let mut child = if cfg!(target_os = "windows") {
43            Command::new("pwsh")
44                .arg("-NoProfile")
45                .arg("-NonInteractive")
46                .arg("-Command")
47                .arg("-")
48                .current_dir(build_dir)
49                .stdin(Stdio::piped())
50                .stdout(Stdio::piped())
51                .stderr(Stdio::piped())
52                .spawn()?
53        } else {
54            Command::new("bash")
55                .arg("--noprofile")
56                .arg("--norc")
57                .current_dir(build_dir)
58                .stdin(Stdio::piped())
59                .stdout(Stdio::piped())
60                .stderr(Stdio::piped())
61                .spawn()?
62        };
63
64        let mut stdin = child.stdin.take().expect("Failed to open stdin");
65        let stdout =
66            BufReader::new(child.stdout.take().expect("Failed to open stdout"));
67        let stderr = child.stderr.take().expect("Failed to open stderr");
68
69        let stderr_buffer = Arc::new(Mutex::new(String::new()));
70        let buffer_clone = Arc::clone(&stderr_buffer);
71
72        // Spawn a thread to consume stderr continuously
73        thread::spawn(move || {
74            let mut reader = BufReader::new(stderr);
75            let mut line = String::new();
76            while reader.read_line(&mut line).unwrap_or(0) > 0 {
77                if let Ok(mut buf) = buffer_clone.lock() {
78                    buf.push_str(&line);
79                }
80                line.clear();
81            }
82        });
83
84        // Inject Zoi environment variables
85        let mut env_cmds = String::new();
86        let globals = lua.globals();
87
88        let vars = [
89            ("BUILD_TYPE", "BUILD_TYPE"),
90            ("SUBPKG", "SUBPKG"),
91            ("BUILD_DIR", "BUILD_DIR"),
92            ("STAGING_DIR", "STAGING_DIR")
93        ];
94
95        for (lua_name, env_name) in vars {
96            if let Ok(val) = globals.get::<String>(lua_name) {
97                if cfg!(target_os = "windows") {
98                    let _ = writeln!(
99                        env_cmds,
100                        "$env:{env_name} = '{}'",
101                        val.replace('\'', "''")
102                    );
103                } else {
104                    let _ = writeln!(env_cmds, "export {env_name}={val:?}");
105                }
106            }
107        }
108
109        // Handle SYSTEM and ZOI tables
110        let tables = [("SYSTEM", "SYSTEM_"), ("ZOI", "ZOI_")];
111        for (table_name, prefix) in tables {
112            if let Ok(table) = globals.get::<Table>(table_name) {
113                for (k, v) in table.pairs::<String, Value>().flatten() {
114                    let val_str = match v {
115                        Value::String(s) => s
116                            .to_str()
117                            .map_err(|e| anyhow::anyhow!(e.to_string()))?
118                            .to_string(),
119                        Value::Integer(i) => i.to_string(),
120                        Value::Number(n) => n.to_string(),
121                        Value::Boolean(b) => b.to_string(),
122                        _ => continue
123                    };
124                    let k_upper = k.to_uppercase();
125                    if cfg!(target_os = "windows") {
126                        let _ = writeln!(
127                            env_cmds,
128                            "$env:{prefix}{k_upper} = '{}'",
129                            val_str.replace('\'', "''")
130                        );
131                    } else {
132                        let _ = writeln!(
133                            env_cmds,
134                            "export {prefix}{k_upper}={val_str:?}"
135                        );
136                    }
137                }
138            }
139        }
140
141        stdin.write_all(env_cmds.as_bytes())?;
142        stdin.flush()?;
143
144        Ok(Self {
145            inner: Mutex::new(InnerSession {
146                child,
147                stdin,
148                stdout,
149                stderr_buffer,
150                sentinel
151            })
152        })
153    }
154}
155
156/// Exposes system command and patching utilities to the Lua environment.
157///
158/// # Errors
159///
160/// Returns an error if the `cmd` function cannot be registered in the Lua
161/// globals.
162///
163/// # Panics
164///
165/// Panics if the internal shell session mutex is poisoned.
166pub fn add_cmd_util(lua: &Lua, quiet: bool) -> Result<(), mlua::Error> {
167    let cmd_fn = lua.create_function(move |lua, command: String| {
168        let build_dir: String = lua.globals().get("BUILD_DIR")?;
169
170        let session_is_dead =
171            if let Some(session) = lua.app_data_ref::<ShellSession>() {
172                let mut inner = session.inner.lock().expect("lock poisoned");
173                inner.child.try_wait().map_or(true, |s| s.is_some())
174            } else {
175                true
176            };
177
178        if session_is_dead {
179            let session = ShellSession::new(lua, &build_dir)
180                .map_err(|e| mlua::Error::RuntimeError(e.to_string()))?;
181            lua.set_app_data(session);
182        }
183
184        let session = lua
185            .app_data_ref::<ShellSession>()
186            .expect("ShellSession missing from app_data");
187        let mut inner = session.inner.lock().expect("lock poisoned");
188
189        if !quiet {
190            println!("Executing: {command}");
191        }
192
193        // Clear stderr buffer before running command
194        if let Ok(mut buf) = inner.stderr_buffer.lock() {
195            buf.clear();
196        }
197
198        let sentinel = inner.sentinel.clone();
199
200        if cfg!(target_os = "windows") {
201            let cmd_text = format!(
202                "$ErrorActionPreference = 'Continue'; & {{ {command} }}; \
203                 \"{sentinel} $LASTEXITCODE\"\n"
204            );
205            inner
206                .stdin
207                .write_all(cmd_text.as_bytes())
208                .map_err(|e| mlua::Error::RuntimeError(e.to_string()))?;
209            inner
210                .stdin
211                .flush()
212                .map_err(|e| mlua::Error::RuntimeError(e.to_string()))?;
213        } else {
214            let cmd_text = format!(
215                "{{ {command} ; }} ; printf \"\\n%s %d\\n\" {sentinel:?} $?\n"
216            );
217            inner
218                .stdin
219                .write_all(cmd_text.as_bytes())
220                .map_err(|e| mlua::Error::RuntimeError(e.to_string()))?;
221            inner
222                .stdin
223                .flush()
224                .map_err(|e| mlua::Error::RuntimeError(e.to_string()))?;
225        }
226
227        let mut stdout_accum = String::new();
228        let mut line = String::new();
229        let exit_code;
230
231        loop {
232            line.clear();
233            let n = inner
234                .stdout
235                .read_line(&mut line)
236                .map_err(|e| mlua::Error::RuntimeError(e.to_string()))?;
237            if n == 0 {
238                return Err(mlua::Error::RuntimeError(
239                    "Shell session ended unexpectedly".to_string()
240                ));
241            }
242
243            if let Some(idx) = line.find(&sentinel) {
244                let out_part = &line[..idx];
245                stdout_accum.push_str(out_part);
246
247                let rest = line[idx..]
248                    .strip_prefix(&sentinel)
249                    .expect("sentinel missing")
250                    .trim();
251                exit_code = rest.parse::<i32>().unwrap_or(0);
252                break;
253            }
254            stdout_accum.push_str(&line);
255        }
256
257        // Retrieve captured stderr
258        let stderr = inner
259            .stderr_buffer
260            .lock()
261            .map(|b| b.clone())
262            .unwrap_or_default();
263
264        if exit_code != 0 && !quiet {
265            eprintln!("[cmd] {stderr}");
266        }
267
268        Ok((stdout_accum.trim_end().to_string(), stderr, exit_code))
269    })?;
270    lua.globals().set("cmd", cmd_fn)?;
271    Ok(())
272}
273
274/// Adds the `zpatch` function to the Lua environment for applying patches.
275///
276/// # Errors
277///
278/// Returns an error if the `zpatch` function cannot be registered in the Lua
279/// globals.
280pub fn add_zpatch(lua: &Lua, quiet: bool) -> Result<(), mlua::Error> {
281    let zpatch_fn = lua.create_function(
282        move |lua, (patch_file, strip): (String, Option<u32>)| {
283            let build_dir: String = lua.globals().get("BUILD_DIR")?;
284            let strip_level = strip.unwrap_or(1);
285
286            if !quiet {
287                println!("Applying patch: {patch_file}");
288            }
289
290            let output = std::process::Command::new("patch")
291                .arg(format!("-p{strip_level}"))
292                .arg("-i")
293                .arg(&patch_file)
294                .current_dir(&build_dir)
295                .output();
296
297            match output {
298                Ok(out) => {
299                    if !out.status.success() {
300                        let stderr =
301                            String::from_utf8_lossy(&out.stderr).to_string();
302                        return Err(mlua::Error::RuntimeError(format!(
303                            "patch failed: {stderr}"
304                        )));
305                    }
306                    if !quiet {
307                        println!("Successfully applied patch {patch_file}");
308                    }
309                    Ok(())
310                }
311                Err(e) => Err(mlua::Error::RuntimeError(format!(
312                    "Failed to execute patch command: {e}"
313                )))
314            }
315        }
316    )?;
317    lua.globals().set("zpatch", zpatch_fn)?;
318    Ok(())
319}
zoi_lua/api/system.rs

zoi_lua/api/
system.rs
