remowt-plugin 0.1.0

use std::ffi::OsStr;
use std::io;
use std::process::Stdio;
use std::sync::Mutex;

use bifrostlink::{Port, Rpc, Rtt, WeakRpc};
use bytes::{Bytes, BytesMut};
use tokio::io::{AsyncReadExt as _, AsyncWriteExt as _};
use tokio::process::{Child, ChildStdin, ChildStdout, Command};

use remowt_link_shared::plugin::{Error, PluginEndpoints, PluginHost};
use remowt_link_shared::{Address, BifConfig};

pub fn serve(rpc: &mut Rpc<BifConfig>) {
	let host = Host {
		rpc: rpc.clone().downgrade(),
		children: Mutex::new(Vec::new()),
	};
	PluginEndpoints(host).register_endpoints(rpc);
}

struct Host {
	rpc: WeakRpc<BifConfig>,
	children: Mutex<Vec<Child>>,
}

impl Host {
	fn spawn(&self, id: u16, path: impl AsRef<OsStr>) -> Result<(), Error> {
		let rpc = self.rpc.clone().upgrade().ok_or(Error::Gone)?;

		let mut child = Command::new(path)
			.arg(id.to_string())
			.stdin(Stdio::piped())
			.stdout(Stdio::piped())
			.kill_on_drop(true)
			.spawn()
			.map_err(|e| Error::Spawn(e.to_string()))?;
		let stdin = child.stdin.take().expect("stdin piped");
		let stdout = child.stdout.take().expect("stdout piped");

		rpc.add_direct(Address::Plugin(id), child_port(stdout, stdin), Rtt(0));
		self.children.lock().expect("not poisoned").push(child);
		Ok(())
	}
}

impl PluginHost for Host {
	async fn load_plugin(&self, id: u16, name: String) -> Result<(), Error> {
		// TODO: Right now loads plugin next to the binary...
		// But with our CA addressed schema, the plugins should be located in content-addressed subdir...
		// Maybe it should just be scrapped in favor of load_plugin_path.
		if name.is_empty() || name == "." || name == ".." || name.contains(['/', '\0']) {
			return Err(Error::BadName);
		}
		let exe = std::env::current_exe().map_err(|e| Error::Spawn(e.to_string()))?;
		let dir = exe
			.parent()
			.ok_or_else(|| Error::Spawn("primary agent has no parent directory".to_owned()))?;
		self.spawn(id, dir.join(&name))
	}

	async fn load_plugin_path(&self, id: u16, path: String) -> Result<(), Error> {
		if path.is_empty() || path.contains('\0') {
			return Err(Error::BadName);
		}
		self.spawn(id, path)
	}
}

fn child_port(mut stdout: ChildStdout, mut stdin: ChildStdin) -> Port {
	Port::new(|mut rx, tx| async move {
		let reader = async move {
			loop {
				let len = match stdout.read_u32().await {
					Ok(len) => len,
					Err(e) => {
						tracing::error!("plugin stdout read failed: {e}");
						break;
					}
				};
				let mut buf = BytesMut::zeroed(len as usize);
				if let Err(e) = stdout.read_exact(&mut buf).await {
					tracing::error!("plugin stdout read failed: {e}");
					break;
				}
				if tx.send(buf.freeze()).is_err() {
					break;
				}
			}
		};
		let writer = async move {
			while let Some(msg) = rx.recv().await {
				if let Err(e) = write_frame(&mut stdin, msg).await {
					tracing::error!("plugin stdin write failed: {e}");
					break;
				}
			}
		};
		tokio::join!(reader, writer);
	})
}

async fn write_frame(stdin: &mut ChildStdin, msg: Bytes) -> io::Result<()> {
	let len = u32::try_from(msg.len())
		.map_err(|_| io::Error::new(io::ErrorKind::InvalidInput, "message larger than 4GB"))?;
	stdin.write_u32(len).await?;
	stdin.write_all(&msg).await?;
	stdin.flush().await?;
	Ok(())
}