From fd89c670b91e31dd41758448291750e045fd4594 Mon Sep 17 00:00:00 2001
From: madomado <madonuko@outlook.com>
Date: Tue, 3 Oct 2023 20:29:57 +0800
Subject: [PATCH] feat: script dependency

This adds a new feature where with the `needs:` attribute, other scripts
can be specified as dependencies of this script.

Let's just hope that the code works :3
---
 Cargo.lock     |  14 ++---
 src/builder.rs | 143 ++++++++++++++++++++++++++++++-------------------
 src/config.rs  |   4 +-
 3 files changed, 97 insertions(+), 64 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 0da6e4f..717966f 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -854,9 +854,9 @@ dependencies = [
 
 [[package]]
 name = "memchr"
-version = "2.6.3"
+version = "2.6.4"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "8f232d6ef707e1956a43342693d2a31e72989554d58299d7a88738cc95b0d35c"
+checksum = "f665ee40bc4a3c5590afb1e9677db74a508659dfd71e126420da8274909a0167"
 
 [[package]]
 name = "merge-struct"
@@ -1162,13 +1162,13 @@ dependencies = [
 
 [[package]]
 name = "regex"
-version = "1.9.5"
+version = "1.9.6"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "697061221ea1b4a94a624f67d0ae2bfe4e22b8a17b6a192afb11046542cc8c47"
+checksum = "ebee201405406dbf528b8b672104ae6d6d63e6d118cb10e4d51abbc7b58044ff"
 dependencies = [
  "aho-corasick",
  "memchr",
- "regex-automata 0.3.8",
+ "regex-automata 0.3.9",
  "regex-syntax 0.7.5",
 ]
 
@@ -1183,9 +1183,9 @@ dependencies = [
 
 [[package]]
 name = "regex-automata"
-version = "0.3.8"
+version = "0.3.9"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "c2f401f4955220693b56f8ec66ee9c78abffd8d1c4f23dc41a23839eb88f0795"
+checksum = "59b23e92ee4318893fa3fe3e6fb365258efbfe6ac6ab30f090cdcbb7aa37efa9"
 dependencies = [
  "aho-corasick",
  "memchr",
diff --git a/src/builder.rs b/src/builder.rs
index f83bcf6..b73bfbd 100644
--- a/src/builder.rs
+++ b/src/builder.rs
@@ -1,7 +1,7 @@
-use color_eyre::Result;
+use color_eyre::{eyre::eyre, Result};
 use serde_derive::{Deserialize, Serialize};
 use std::{
-	collections::BTreeMap,
+	collections::{BTreeMap, HashMap},
 	fs,
 	io::{Seek, Write},
 	path::PathBuf,
@@ -67,7 +67,7 @@ impl RootBuilder for DnfRootBuilder {
 	fn build(&self, chroot: PathBuf, manifest: &Manifest) -> Result<()> {
 		info!("Running Pre-install scripts");
 
-		run_scripts(manifest.scripts.pre.clone(), &chroot, false)?;
+		run_all_scripts(manifest.scripts.pre.clone(), &chroot, false)?;
 
 		if let Some(disk) = manifest.clone().disk {
 			let f = disk.fstab(&chroot)?;
@@ -132,65 +132,96 @@ impl RootBuilder for DnfRootBuilder {
 
 		info!("Running post-install scripts");
 
-		run_scripts(manifest.scripts.post.clone(), &chroot, true)?;
+		run_all_scripts(manifest.scripts.post.clone(), &chroot, true)?;
 
 		Ok(())
 	}
 }
 
-pub fn run_scripts(scripts: Vec<Script>, chroot: &PathBuf, in_chroot: bool) -> Result<()> {
-	for script in scripts {
-		if let Some(mut data) = script.load() {
-			info!(
-				"Running script: {script}",
-				script = script.name.as_ref().unwrap_or(&"<Untitled>".to_string())
-			);
-			info!("Script ID: {id}", id = script.id.as_ref().unwrap_or(&"<NULL>".to_string()));
-
-			let script_name =
-				format!("script-{}", script.id.as_ref().unwrap_or(&"untitled".to_string()));
-			// check if data has shebang
-			if !data.starts_with("#!") {
-				// if not, add one
-				warn!("Script does not have shebang, #!/bin/sh will be added. It is recommended to add a shebang to your script.");
-				data.insert_str(0, "#!/bin/sh\n");
-			}
-			// write data to chroot
-			let fpath = if in_chroot {
-				chroot.join("tmp").join(&script_name)
-			} else {
-				PathBuf::from(format!("katsu-work/{}", &script_name))
+pub fn run_script(script: Script, chroot: &PathBuf, in_chroot: bool) -> Result<()> {
+	let id = script.id.as_ref().map_or("<NULL>", |s| &**s);
+	let Some(mut data) = script.load() else { return Err(eyre!("Cannot load script `{}`", script.id.as_ref().map_or("<NULL>", |s| &**s))); };
+	let name = script.name.as_ref().map_or("<Untitled>", |s| &**s);
+	info!(id, "Running script: {name}");
+
+	let name = format!("script-{}", script.id.as_ref().map_or("untitled", |s| &**s));
+	// check if data has shebang
+	if !data.starts_with("#!") {
+		warn!("Script does not have shebang, #!/bin/sh will be added. It is recommended to add a shebang to your script.");
+		data.insert_str(0, "#!/bin/sh\n");
+	}
+	// write data to chroot
+	let fpath = if in_chroot {
+		chroot.join("tmp").join(&name)
+	} else {
+		PathBuf::from(format!("katsu-work/{name}"))
+	};
+	let mut file = fs::File::create(fpath)?;
+	file.write_all(data.as_bytes())?;
+	file.flush()?;
+	drop(file);
+
+	// now add execute bit
+	if in_chroot {
+		util::run_with_chroot(&chroot, || -> Result<()> {
+			cmd_lib::run_cmd!(
+				chmod +x ${chroot}/tmp/${name};
+				unshare -R ${chroot} /tmp/${name} 2>&1;
+				rm -f ${chroot}/tmp/${name};
+			)?;
+			Ok(())
+		})?;
+	} else {
+		// export envar
+		std::env::set_var("CHROOT", chroot);
+		cmd_lib::run_cmd!(
+			chmod +x katsu-work/${name};
+			/usr/bin/env CHROOT=${chroot} katsu-work/${name} 2>&1;
+			rm -f katsu-work/${name};
+		)?;
+	}
+
+	info!(
+		"===== Script {script} finished =====",
+		script = script.name.as_ref().unwrap_or(&"<Untitled>".to_string())
+	);
+	Ok(())
+}
+
+pub fn run_all_scripts(scripts: Vec<Script>, chroot: &PathBuf, in_chroot: bool) -> Result<()> {
+	let mut scrs: HashMap<String, (Script, bool)> = HashMap::new();
+	scripts.into_iter().for_each(|s| {
+		scrs.insert(s.id.clone().unwrap_or("<?>".into()), (s, false));
+	});
+	run_scripts(scrs, chroot, in_chroot)
+}
+
+pub fn run_scripts(
+	mut scripts: HashMap<String, (Script, bool)>, chroot: &PathBuf, in_chroot: bool,
+) -> Result<()> {
+	for idx in scripts.clone().keys() {
+		// FIXME: if someone dares to optimize things with unsafe, go for it
+		// we can't use get_mut here because we need to do scripts.get_mut() later
+		let Some((scr, done)) = scripts.get(idx) else { unreachable!() };
+		if *done {
+			continue;
+		}
+		let id = scr.id.clone().unwrap_or("<NULL>".into());
+		let mut needs = HashMap::new();
+		for need in scr.needs.clone() {
+			let Some((s, done)) = scripts.get_mut(&need) else {
+				return Err(eyre!("Script `{need}` required by `{id}` not found"));
 			};
-			let mut file = fs::File::create(fpath)?;
-			file.write_all(data.as_bytes())?;
-			file.flush()?;
-			drop(file);
-
-			// now add execute bit
-			if in_chroot {
-				util::run_with_chroot(&chroot, || -> color_eyre::Result<()> {
-					cmd_lib::run_cmd!(
-						chmod +x ${chroot}/tmp/${script_name};
-						unshare -R ${chroot} /tmp/${script_name} 2>&1;
-						rm -f ${chroot}/tmp/${script_name};
-					)?;
-					Ok(())
-				})?;
-			} else {
-				// export envar
-				std::env::set_var("CHROOT", chroot);
-				cmd_lib::run_cmd!(
-					chmod +x katsu-work/${script_name};
-					/usr/bin/env CHROOT=${chroot} katsu-work/${script_name} 2>&1;
-					rm -f katsu-work/${script_name};
-				)?;
+			if *done {
+				continue;
 			}
-
-			info!(
-				"===== Script {script} finished =====",
-				script = script.name.as_ref().unwrap_or(&"<Untitled>".to_string())
-			);
+			needs.insert(need, (std::mem::take(s), false));
+			*done = true;
 		}
+		run_scripts(needs, chroot, in_chroot)?;
+		let Some((scr, done)) = scripts.get_mut(idx) else { unreachable!() };
+		run_script(std::mem::take(scr), chroot, in_chroot)?;
+		*done = true;
 	}
 	Ok(())
 }
@@ -215,11 +246,11 @@ impl ImageBuilder for DiskImageBuilder {
 
 		if manifest.clone().disk.is_none() {
 			// error out
-			return Err(color_eyre::eyre::eyre!("Disk layout not specified"));
+			return Err(eyre!("Disk layout not specified"));
 		} else {
 			info!("Disk layout specified");
 			if manifest.clone().disk.unwrap().size.is_none() {
-				return Err(color_eyre::eyre::eyre!("Disk size not specified"));
+				return Err(eyre!("Disk size not specified"));
 			}
 		}
 
diff --git a/src/config.rs b/src/config.rs
index 04ba457..be77f52 100644
--- a/src/config.rs
+++ b/src/config.rs
@@ -115,13 +115,15 @@ pub struct ScriptsManifest {
 	pub post: Vec<Script>,
 }
 
-#[derive(Deserialize, Debug, Clone, Serialize, PartialEq, Eq)]
+#[derive(Deserialize, Debug, Clone, Serialize, PartialEq, Eq, Default)]
 // load script from file, or inline if there's one specified
 pub struct Script {
 	pub id: Option<String>,
 	pub name: Option<String>,
 	pub file: Option<PathBuf>,
 	pub inline: Option<String>,
+	#[serde(default)]
+	pub needs: Vec<String>,
 }
 
 impl Script {