nitro_plugin 0.31.0

/// Manager for loading and caching WASM efficiently
pub mod loader;

use std::{
	fs::File,
	marker::PhantomData,
	path::{Path, PathBuf},
	process::Stdio,
	sync::Arc,
	time::Instant,
};

use anyhow::{Context, bail};
use nitro_net::download::{self, Client};
use nitro_shared::{
	Side,
	io::{home_dir, update_link},
	nitro_executable::NitroExecutableRegistry,
	no_window,
	output::{Message, MessageContents, MessageLevel, NitroOutput},
	util::{ARCH_STRING, OS_STRING},
};
use tokio::{
	process::Command,
	sync::{Mutex, oneshot},
	task::JoinSet,
};
use wasmtime::{
	Store,
	component::{HasSelf, Linker},
};
use wasmtime_wasi::{
	DirPerms, FilePerms, ResourceTable, WasiCtx, WasiCtxBuilder, WasiCtxView, WasiView,
};
use wasmtime_wasi_http::{WasiHttpCtx, WasiHttpView};

use crate::{
	hook::{
		Hook,
		call::{HookCallArg, HookHandle},
		wasm::loader::WASMLoader,
	},
	host::PluginContext,
	plugin::PluginPersistence,
	plugin_debug_enabled,
};

#[allow(missing_docs)]
mod bindings {
	wasmtime::component::bindgen!({
		path: "src/interface.wit",
		imports: { default: async },
		exports: { default: async }
	});
}

/// Calls a WASM hook handler
pub(crate) async fn call_wasm<H: Hook + Sized>(
	hook: &H,
	arg: HookCallArg<'_, H>,
	o: &mut impl NitroOutput,
) -> anyhow::Result<HookHandle<H>> {
	let _ = hook;

	// Use a full output for synchronous hooks
	let o = if !H::is_asynchronous() || H::get_takes_over() {
		o.get_greater_copy()
	} else {
		o.get_lesser_copy()
	};

	let o = Arc::new(Mutex::new(o));

	let (result_sender, result) = oneshot::channel();

	Ok(HookHandle::wasm(
		WASMHookHandle {
			plugin_id: arg.plugin_id.to_string(),
			o,
			wasm_path: PathBuf::from(arg.cmd),
			arg: serde_json::to_string(&arg.arg)?,
			result_sender: Some(result_sender),
			result,
			custom_config: arg.ctx.custom_config,
			context: arg.ctx.global_context.cloned(),
			persistence: arg.persistence.clone(),
			wasm_loader: arg.wasm_loader,
			data_dir: arg.paths.data_dir.to_string_lossy().to_string(),
			config_dir: arg.paths.config_dir.to_string_lossy().to_string(),
			plugin_dir: arg
				.working_dir
				.unwrap_or(Path::new(""))
				.to_string_lossy()
				.to_string(),
			_phantom: PhantomData,
		},
		arg.plugin_id.to_string(),
		arg.persistence,
	))
}

/// Hook handler internals for a WASM hook
pub(super) struct WASMHookHandle<H: Hook> {
	pub plugin_id: String,
	o: Arc<Mutex<Box<dyn NitroOutput + Sync>>>,
	wasm_path: PathBuf,
	arg: String,
	result_sender: Option<oneshot::Sender<anyhow::Result<H::Result>>>,
	result: oneshot::Receiver<anyhow::Result<H::Result>>,
	custom_config: Option<String>,
	context: Option<Arc<dyn PluginContext>>,
	persistence: Arc<Mutex<PluginPersistence>>,
	wasm_loader: Arc<Mutex<WASMLoader>>,
	data_dir: String,
	config_dir: String,
	plugin_dir: String,
	_phantom: PhantomData<H>,
}

impl<H: Hook> WASMHookHandle<H> {
	/// Starts this hook
	pub async fn run(&mut self, o: &mut impl NitroOutput) -> anyhow::Result<()> {
		if !self.result.is_empty() {
			return Ok(());
		}

		let Some(result_sender) = self.result_sender.take() else {
			return Ok(());
		};

		if plugin_debug_enabled() {
			o.display(MessageContents::Simple(format!(
				"Running hook '{}' on plugin '{}'",
				H::get_name_static(),
				self.plugin_id
			)));
		}

		let mut start_time = if std::env::var("NITRO_PLUGIN_PROFILE").is_ok_and(|x| x == "1") {
			Some(Instant::now())
		} else {
			None
		};

		// Initialize the component
		let mut lock = self.wasm_loader.lock().await;
		let component = lock
			.load(self.plugin_id.clone(), &self.wasm_path)
			.await
			.context("Failed to load WASM component")?;
		let engine = lock.engine();
		std::mem::drop(lock);

		if let Some(start_time) = &mut start_time {
			let now = Instant::now();
			println!("Component initialization: {:?}", now - *start_time);
			*start_time = now;
		}

		let mut linker = Linker::new(&engine);

		let mut wasi_ctx = WasiCtxBuilder::new();
		let wasi_ctx = wasi_ctx.inherit_stdio().inherit_env().inherit_network();

		#[cfg(not(target_os = "windows"))]
		let wasi_ctx = wasi_ctx.preopened_dir("/", "/", DirPerms::all(), FilePerms::all())?;
		#[cfg(target_os = "windows")]
		let wasi_ctx = wasi_ctx.preopened_dir("C:\\", "C:\\", DirPerms::all(), FilePerms::all())?;

		let wasi_ctx = wasi_ctx.build();

		let http_ctx = WasiHttpCtx::new();

		let state = State {
			wasi_ctx,
			http_ctx,
			table: ResourceTable::new(),
			custom_config: self.custom_config.clone(),
			context: self.context.clone(),
			persistence: self.persistence.clone(),
			data_dir: self.data_dir.clone(),
			config_dir: self.config_dir.clone(),
			plugin_dir: self.plugin_dir.clone(),
			client: Client::new(),
			o: self.o.clone(),
		};

		let arg = self.arg.clone();
		let plugin_id = self.plugin_id.clone();

		tokio::task::spawn(async move {
			let fun = async move || {
				wasmtime_wasi::p2::add_to_linker_async(&mut linker)
					.context("Failed to add WASI functions to linker")?;
				wasmtime_wasi_http::add_only_http_to_linker_async(&mut linker)
					.context("Failed to add HTTP functions to linker")?;

				bindings::InterfaceWorld::add_to_linker::<_, HasSelf<_>>(&mut linker, |x| x)?;

				if let Some(start_time) = &mut start_time {
					let now = Instant::now();
					println!("Linker initialization: {:?}", now - *start_time);
					*start_time = now;
				}

				let mut store = Store::new(&engine, state);

				let instance =
					bindings::InterfaceWorld::instantiate_async(&mut store, &component, &linker)
						.await
						.context("Failed to construct WASM instance")?;

				if let Some(start_time) = &mut start_time {
					let now = Instant::now();
					println!("Instance initialization: {:?}", now - *start_time);
					*start_time = now;
				}

				let result_code = instance
					.call_run_plugin(
						&mut store,
						H::get_name_static(),
						&arg,
						H::get_version() as u32,
					)
					.await
					.context("Failed to call plugin entrypoint")?;

				if let Some(start_time) = &mut start_time {
					let now = Instant::now();
					println!("Hook runtime: {:?}", now - *start_time);
					*start_time = now;
				}

				let result = if H::get_takes_over() {
					H::Result::default()
				} else {
					let mut result = instance.call_get_result(&mut store).await?;

					if result_code == 1 {
						bail!("Plugin returned an error: {result}");
					}

					unsafe { simd_json::from_str(&mut result) }
						.context("Failed to deserialize hook result")?
				};

				if let Some(start_time) = &mut start_time {
					let now = Instant::now();
					println!("Result handling: {:?}", now - *start_time);
					*start_time = now;
				}

				Ok(result)
			};

			let result = fun()
				.await
				.context(format!("Hook for plugin {plugin_id} failed"));
			let _ = result_sender.send(result);
		});

		Ok(())
	}

	/// Awaits the result of the hook. Hook must have been started or this will run indefinitely.
	pub async fn result(self) -> anyhow::Result<H::Result> {
		self.result.await.context("Channel closed").flatten()
	}

	/// Checks whether the handle has a result
	pub fn has_result(&self) -> bool {
		!self.result.is_empty()
	}
}

/// Host function environment
struct State {
	wasi_ctx: WasiCtx,
	http_ctx: WasiHttpCtx,
	table: ResourceTable,
	custom_config: Option<String>,
	context: Option<Arc<dyn PluginContext>>,
	persistence: Arc<Mutex<PluginPersistence>>,
	data_dir: String,
	config_dir: String,
	plugin_dir: String,
	client: Client,
	o: Arc<Mutex<Box<dyn NitroOutput + Sync>>>,
}

impl WasiView for State {
	fn ctx(&mut self) -> wasmtime_wasi::WasiCtxView<'_> {
		WasiCtxView {
			ctx: &mut self.wasi_ctx,
			table: &mut self.table,
		}
	}
}

impl WasiHttpView for State {
	fn ctx(&mut self) -> &mut WasiHttpCtx {
		&mut self.http_ctx
	}

	fn table(&mut self) -> &mut ResourceTable {
		&mut self.table
	}
}

impl bindings::InterfaceWorldImports for State {
	async fn get_custom_config(&mut self) -> Option<String> {
		self.custom_config.clone()
	}

	async fn get_persistent_state(&mut self) -> String {
		serde_json::to_string(&self.persistence.lock().await.state)
			.unwrap_or_else(|_| "null".into())
	}

	async fn set_persistent_state(&mut self, state: String) {
		if let Ok(state) = serde_json::from_str(&state) {
			self.persistence.lock().await.state = state;
		}
	}

	async fn get_data_dir(&mut self) -> String {
		self.data_dir.clone()
	}

	async fn get_config_dir(&mut self) -> String {
		self.config_dir.clone()
	}

	async fn get_plugin_dir(&mut self) -> String {
		self.plugin_dir.clone()
	}

	async fn get_current_dir(&mut self) -> String {
		std::env::current_dir()
			.unwrap_or_default()
			.to_string_lossy()
			.to_string()
	}

	async fn get_home_dir(&mut self) -> String {
		home_dir()
			.map(|x| x.to_string_lossy().to_string())
			.unwrap_or_else(|_| "/home/none".into())
	}

	async fn get_os_string(&mut self) -> String {
		OS_STRING.to_string()
	}

	async fn get_arch_string(&mut self) -> String {
		ARCH_STRING.to_string()
	}

	async fn get_pointer_width(&mut self) -> u32 {
		#[cfg(target_pointer_width = "32")]
		return 32;
		#[cfg(target_pointer_width = "64")]
		return 64;
	}

	async fn update_hardlink(&mut self, src: String, tgt: String) -> Result<(), String> {
		let result = if !PathBuf::from(&tgt).exists() {
			tokio::fs::hard_link(tgt, src).await
		} else {
			Ok(())
		};
		match result {
			Ok(..) => Ok(()),
			Err(e) => Err(format!("{e:?}")),
		}
	}

	async fn update_link(&mut self, src: String, tgt: String) -> Result<(), String> {
		let result = update_link(Path::new(&tgt), Path::new(&src));
		match result {
			Ok(..) => Ok(()),
			Err(e) => Err(format!("{e:?}")),
		}
	}

	async fn download_bytes(&mut self, url: String) -> Result<Vec<u8>, String> {
		let result = download::bytes(url, &self.client).await;
		match result {
			Ok(result) => Ok(result.to_vec()),
			Err(e) => Err(format!("{e:?}")),
		}
	}

	async fn download_text(&mut self, url: String) -> Result<String, String> {
		let result = download::text(url, &self.client).await;
		match result {
			Ok(result) => Ok(result),
			Err(e) => Err(format!("{e:?}")),
		}
	}

	async fn download_file(&mut self, url: String, path: String) -> Result<(), String> {
		let result = download::file(url, path, &self.client).await;
		match result {
			Ok(..) => Ok(()),
			Err(e) => Err(format!("{e:?}")),
		}
	}

	async fn download_files(
		&mut self,
		urls: Vec<String>,
		paths: Vec<String>,
		skip_existing: bool,
	) -> Result<(), String> {
		let mut tasks = JoinSet::new();
		for (url, path) in urls.into_iter().zip(paths) {
			let path = PathBuf::from(path);
			if skip_existing && path.exists() {
				continue;
			}

			let client = self.client.clone();
			tasks.spawn(async move { download::file(url, path, &client).await });
		}

		let mut final_result = Ok(());
		while let Some(result) = tasks.join_next().await {
			match result {
				Ok(result) => {
					if final_result.is_ok() {
						final_result = result.map_err(|e| format!("{e:?}"));
					}
				}
				Err(e) => final_result = Err(e.to_string()),
			}
		}

		final_result
	}

	async fn run_command(
		&mut self,
		cmd: String,
		args: Vec<String>,
		working_dir: Option<String>,
		stdout_file: Option<String>,
		suppress_command_window: bool,
		silent: bool,
		wait: bool,
	) -> Result<(i32, u32), String> {
		let mut command = Command::new(cmd);
		command.args(args);
		if let Some(working_dir) = working_dir {
			command.current_dir(working_dir);
		}
		if suppress_command_window {
			no_window!(command);
		}

		if silent {
			command.stdin(Stdio::null());
			command.stdout(Stdio::null());
			command.stderr(Stdio::null());
		}

		if let Some(stdout_file) = stdout_file {
			let file = File::create(stdout_file).map_err(|e| format!("{e:?}"))?;
			command.stdout(Stdio::from(file));
		}

		let mut child = command.spawn().map_err(|e| format!("{e:?}"))?;
		let pid = child.id().unwrap();

		if wait {
			let status = child.wait().await.map_err(|e| format!("{e:?}"))?;
			Ok((status.code().unwrap_or_default(), pid))
		} else {
			Ok((0, pid))
		}
	}

	async fn get_instances(&mut self) -> Option<Vec<(String, String)>> {
		let Some(context) = &self.context else {
			return None;
		};
		let instances = context.get_instances();

		Some(
			instances
				.iter()
				.filter_map(|(k, v)| {
					if let Ok(config) = serde_json::to_string(v) {
						Some((k.clone(), config))
					} else {
						None
					}
				})
				.collect(),
		)
	}

	async fn get_templates(&mut self) -> Option<Vec<(String, String)>> {
		let Some(context) = &self.context else {
			return None;
		};
		let templates = context.get_templates();

		Some(
			templates
				.iter()
				.filter_map(|(k, v)| {
					if let Ok(config) = serde_json::to_string(v) {
						Some((k.clone(), config))
					} else {
						None
					}
				})
				.collect(),
		)
	}

	async fn get_instance_dir(&mut self, instance: String) -> Result<Option<String>, String> {
		let Some(context) = &self.context else {
			return Err("Missing context".into());
		};
		let instances = context.get_instances();
		let Some(config) = instances.get(&instance) else {
			return Err("Instance does not exist".into());
		};

		let inst_dir = if let Some(inst_dir) = &config.dir {
			if inst_dir == "none" {
				return Ok(None);
			} else {
				inst_dir.clone()
			}
		} else {
			let base_dir = Path::new(&self.data_dir).join("instances").join(instance);
			if config.side == Some(Side::Client) {
				base_dir.join(".minecraft").to_string_lossy().to_string()
			} else {
				base_dir.to_string_lossy().to_string()
			}
		};

		Ok(Some(inst_dir))
	}

	async fn create_instance(&mut self, id: String, config: String) -> Result<(), String> {
		if let Some(context) = &self.context {
			let Ok(config) = serde_json::from_str(&config) else {
				return Err("Failed to deserialize config".into());
			};
			context
				.create_instance(id, config)
				.await
				.map_err(|e| e.to_string())?;
			Ok(())
		} else {
			Err("Context missing".into())
		}
	}

	async fn create_template(&mut self, id: String, config: String) -> Result<(), String> {
		if let Some(context) = &self.context {
			let Ok(config) = serde_json::from_str(&config) else {
				return Err("Failed to deserialize config".into());
			};
			context
				.create_template(id, config)
				.await
				.map_err(|e| e.to_string())?;
			Ok(())
		} else {
			Err("Context missing".into())
		}
	}

	async fn launch_instance(
		&mut self,
		instance: String,
		account: Option<String>,
	) -> Result<(), String> {
		let executable_registry = fmt_err(NitroExecutableRegistry::open(
			&PathBuf::from(&self.data_dir).join("internal"),
		))?;

		let mut command = fmt_err(
			executable_registry
				.launch_instance(&instance, account.as_deref(), None)
				.context("No executable available"),
		)?;

		fmt_err(command.spawn().context("Failed to launch instance"))?;

		Ok(())
	}

	async fn output_display_text(&mut self, text: String, level: u8) {
		let level = match level {
			0 => MessageLevel::Important,
			1 => MessageLevel::Debug,
			2 => MessageLevel::Trace,
			_ => return,
		};

		self.o.lock().await.display_text(text, level);
	}

	async fn output_display_message(&mut self, message: String, level: u8) {
		let Ok(message) = serde_json::from_str::<MessageContents>(&message) else {
			return;
		};

		let level = match level {
			0 => MessageLevel::Important,
			1 => MessageLevel::Debug,
			2 => MessageLevel::Trace,
			_ => return,
		};

		self.o.lock().await.display_message(Message {
			contents: message,
			level,
		});
	}

	async fn output_start_process(&mut self) {
		self.o.lock().await.start_process();
	}

	async fn output_end_process(&mut self) {
		self.o.lock().await.end_process();
	}

	async fn output_start_section(&mut self) {
		self.o.lock().await.start_section();
	}

	async fn output_end_section(&mut self) {
		self.o.lock().await.end_section();
	}
}

fn fmt_err<T>(x: anyhow::Result<T>) -> Result<T, String> {
	x.map_err(|e| e.to_string())
}