From 2bf68c1e2a9e5b906a7983b9d0441476c781d347 Mon Sep 17 00:00:00 2001
From: Erica Marigold <hi@devcomp.xyz>
Date: Wed, 22 Nov 2023 17:49:30 +0530
Subject: [PATCH] refactor: polish a few things and clean up code

---
 src/cli/build.rs | 33 +++++++++++++++++++--------------
 src/cli/mod.rs   | 31 ++++++++++++++++++++++++-------
 2 files changed, 43 insertions(+), 21 deletions(-)

diff --git a/src/cli/build.rs b/src/cli/build.rs
index c25f8d6..6bf9871 100644
--- a/src/cli/build.rs
+++ b/src/cli/build.rs
@@ -1,8 +1,12 @@
 use std::{
     env,
     path::{Path, PathBuf},
+    process::ExitCode,
+};
+use tokio::{
+    fs::{self, OpenOptions},
+    io::AsyncWriteExt,
 };
-use tokio::fs;
 
 use anyhow::Result;
 use mlua::Compiler as LuaCompiler;
@@ -10,13 +14,13 @@ use mlua::Compiler as LuaCompiler;
 pub async fn build_standalone<T: AsRef<Path> + Into<PathBuf>>(
     output_path: T,
     code: impl AsRef<[u8]>,
-) -> Result<()> {
+) -> Result<ExitCode> {
     // First, we read the contents of the lune interpreter as our starting point
     let mut patched_bin = fs::read(env::current_exe()?).await?;
     let base_bin_offset = u64::try_from(patched_bin.len())?;
 
     // The signature which separates indicates the presence of bytecode to execute
-    // If a binary contains this signature, that must mean it is a standalone binar
+    // If a binary contains this signature, that must mean it is a standalone binary
     let signature: Vec<u8> = vec![0x4f, 0x3e, 0xf8, 0x41, 0xc3, 0x3a, 0x52, 0x16];
 
     // Compile luau input into bytecode
@@ -26,27 +30,28 @@ pub async fn build_standalone<T: AsRef<Path> + Into<PathBuf>>(
         .set_debug_level(0)
         .compile(code);
 
-    println!("{}", bytecode.len());
-
     patched_bin.append(&mut bytecode.clone());
 
     let mut meta = base_bin_offset.to_ne_bytes().to_vec();
 
-    // bytecode base size files signature
-    // meta.append(&mut [0, 0, 0, 0].to_vec()); // 4 extra padding bytes after 4 byte u64
-    meta.append(&mut (bytecode.len() as u64).to_ne_bytes().to_vec()); // FIXME: len is greater than u8::max
+    // Include metadata in the META chunk, each field is 8 bytes
+    meta.append(&mut (bytecode.len() as u64).to_ne_bytes().to_vec()); // Size of bytecode, used to calculate end offset at runtime
     meta.append(&mut 1_u64.to_ne_bytes().to_vec()); // Number of files, padded with null bytes
-                                                    // meta.append(&mut [0, 0, 0, 0].to_vec()); // 4 extra padding bytes after 4 byte u32
 
     patched_bin.append(&mut meta);
 
     // Append the signature to the base binary
-    for byte in signature {
-        patched_bin.push(byte);
-    }
+    patched_bin.append(&mut signature.clone());
 
     // Write the compiled binary to file
-    fs::write(output_path, patched_bin).await?;
+    OpenOptions::new()
+        .write(true)
+        .create(true)
+        .mode(0o770)
+        .open(output_path)
+        .await?
+        .write_all(&patched_bin)
+        .await?;
 
-    Ok(())
+    Ok(ExitCode::SUCCESS)
 }
diff --git a/src/cli/mod.rs b/src/cli/mod.rs
index caee24c..71cde7e 100644
--- a/src/cli/mod.rs
+++ b/src/cli/mod.rs
@@ -155,17 +155,26 @@ impl Cli {
 
             // Read the current lune binary to memory
             let bin = read_to_vec(env::current_exe()?).await?;
-            // let mut bin = vec![0; tmp_bin.len() % signature.len()];
-
-            // println!("padding: {:?}", bin);
-
-            // bin.append(&mut tmp_bin);
 
             let mut bytecode_offset = 0;
             let mut bytecode_size = 0;
 
-            // println!("{}",);
+            // standalone binary structure (reversed, 8 bytes per field)
+            // [0] => signature
+            // ----------------
+            // -- META Chunk --
+            // [1] => file count
+            // [2] => bytecode size
+            // [3] => bytecode offset
+            // ----------------
+            // -- MISC Chunk --
+            // [4..n] => bytecode (variable size)
+            // ----------------
+            // NOTE: All integers are 8 byte unsigned 64 bit (u64's).
 
+            // The rchunks will have unequally sized sections in the beginning
+            // but that doesn't matter to us because we don't need anything past the
+            // middle chunks where the bytecode is stored
             for (idx, chunk) in bin.rchunks(signature.len()).enumerate() {
                 if idx == 0 && chunk != signature {
                     // We don't have a standalone binary
@@ -181,7 +190,13 @@ impl Cli {
                 }
             }
 
+            // If we were able to retrieve the required metadata, we load
+            // and execute the bytecode
             if bytecode_offset != 0 && bytecode_size != 0 {
+                // FIXME: Passing arguments does not work like it should, because the first
+                // argument provided is treated as the script path. We should probably also not
+                // allow any runner functionality within standalone binaries
+
                 let result = Lune::new()
                     .with_args(self.script_args.clone())
                     .run(
@@ -200,6 +215,8 @@ impl Cli {
                 });
             }
 
+            // If not in a standalone context and we don't have any arguments
+            // display the interactive REPL interface
             return repl::show_interface().await;
         }
         // Figure out if we should read from stdin or from a file,
@@ -232,7 +249,7 @@ impl Cli {
 
             return Ok(
                 match build_standalone(output_path, strip_shebang(script_contents.clone())).await {
-                    Ok(()) => ExitCode::SUCCESS,
+                    Ok(exitcode) => exitcode,
                     Err(err) => {
                         eprintln!("{err}");
                         ExitCode::FAILURE