git.delta.rocks / fleet / refs/commits / b66112ef1437

1use std::collections::HashMap;2use std::env;3use std::path::PathBuf;4use std::sync::atomic::AtomicU64;5use std::sync::{Arc, Mutex};67use anyhow::{Context as _, Result, anyhow, bail, ensure};8use bifrostlink::declarative::RemoteEndpoints;9use bifrostlink::{Remote, Rpc, Rtt};10use camino::{Utf8Path, Utf8PathBuf};11use remowt_link_shared::iroh_tunnel::{DatagramRouter, IrohBiStream, TunnelAddr};12use remowt_link_shared::plugin::PluginEndpointsClient;13use remowt_link_shared::port::child_port;14use remowt_link_shared::{Address, BifConfig};15use russh::Channel;16use russh::client::{Config, Handle, Handler, Msg, Session, connect};17use russh::keys::agent::AgentIdentity;18use russh::keys::agent::client::AgentClient;19use russh::keys::check_known_hosts;20use russh::keys::ssh_key::PublicKey;21use tempfile::TempDir;22use tokio::io::AsyncRead;23use tokio::net::UnixListener;24use tokio::sync::oneshot;25use tokio::task::JoinHandle;26use tokio::{27	fs,28	io::{AsyncBufReadExt as _, AsyncReadExt as _, AsyncWriteExt as _, BufReader},29};30use tracing::{debug, info, warn};31use uuid::Uuid;3233pub mod editor;34mod forwarded;35mod shell;36mod ssh_exec;37mod subprocess;3839use self::ssh_exec::SshExecChild;40pub use self::subprocess::{RemowtChild, SpawnOptions, StderrMode, StdioMode};41pub use forwarded::{RemowtListener, RemowtStream};42pub use shell::{RemowtShell, RemowtShellResizer};4344type Subs = Arc<Mutex<HashMap<Utf8PathBuf, oneshot::Sender<Channel<Msg>>>>>;4546fn sh_quote(s: impl AsRef<str>) -> String {47	format!("'{}'", s.as_ref().replace('\'', "'\\''"))48}4950const ESCALATORS: [(&str, &[&str]); 2] = [("run0", &["--background=", "--pipe"]), ("sudo", &[])];5152pub struct AgentBundle {53	dir: PathBuf,54	hashes: HashMap<String, String>,55}5657impl AgentBundle {58	pub fn from_dir(dir: impl Into<PathBuf>) -> Result<Self> {59		let dir = dir.into();60		let hashes_path = dir.join("hashes");61		let raw = std::fs::read_to_string(&hashes_path)62			.with_context(|| format!("reading agent hashes at {}", hashes_path.display()))?;63		let mut hashes = HashMap::new();64		for line in raw.lines() {65			let line = line.trim();66			if line.is_empty() {67				continue;68			}69			let (arch, hash) = line70				.split_once(char::is_whitespace)71				.ok_or_else(|| anyhow!("malformed hashes line: {line:?}"))?;72			hashes.insert(arch.to_owned(), hash.trim().to_owned());73		}74		ensure!(75			!hashes.is_empty(),76			"agent bundle {} has no hashes",77			dir.display()78		);79		Ok(Self { dir, hashes })80	}8182	fn binary(&self, arch: &str) -> PathBuf {83		self.dir.join(format!("remowt-agent-{arch}"))84	}8586	fn local_binary(&self) -> Result<PathBuf> {87		let arch = env::consts::ARCH;88		let path = self.binary(arch);89		ensure!(90			path.is_file(),91			"no local remowt-agent build for arch {arch} in bundle {}",92			self.dir.display()93		);94		Ok(path)95	}96}9798async fn run(sess: &Handle<SshHandler>, cmd: &str) -> Result<(Option<u32>, Vec<u8>)> {99	let ch = sess.channel_open_session().await?;100	ch.exec(true, cmd).await?;101102	let mut child = SshExecChild::from_exec(ch);103	drop(child.stdin);104	drain_to_tracing(child.stderr, cmd.to_owned(), true);105106	let mut out = Vec::new();107	child.stdout.read_to_end(&mut out).await?;108	let code = child.exit.await.ok().flatten();109	Ok((code, out))110}111112async fn run_string_ok(sess: &Handle<SshHandler>, cmd: &str) -> Result<String> {113	let (code, mut out) = run(sess, cmd).await?;114	ensure!(115		code == Some(0),116		"remote command failed (exit {code:?}): {cmd}"117	);118	if !out.is_empty() {119		ensure!(120			out.ends_with(b"\n"),121			"remote command was not newline-terminated: {cmd}: {out:?}"122		);123		out.pop();124	}125	String::from_utf8(out).context("expected utf8 output for command")126}127128async fn deploy_agent(sess: &Handle<SshHandler>, bundle: &AgentBundle) -> Result<Utf8PathBuf> {129	debug!("uname -a");130	let arch = run_string_ok(sess, "uname -m").await?;131	let hash = bundle132		.hashes133		.get(&arch)134		.ok_or_else(|| anyhow!("no remowt-agent build for remote arch {arch:?}"))?;135136	debug!("get dir");137	let cache = run_string_ok(sess, "echo \"$XDG_CACHE_HOME\"").await?;138	let dir = if cache.is_empty() {139		let home = run_string_ok(sess, "echo \"$HOME\"").await?;140		ensure!(141			!home.is_empty(),142			"remote $HOME and $XDG_CACHE_HOME both empty"143		);144		Utf8PathBuf::from(home).join(".cache/remowt")145	} else {146		Utf8PathBuf::from(cache).join("remowt")147	};148	let path = dir.join(hash);149150	debug!("presence");151	let (present, _) = run(sess, &format!("test -x {}", sh_quote(&path))).await?;152	if present != Some(0) {153		let bin = bundle.binary(&arch);154		debug!("read");155		let bytes = fs::read(&bin)156			.await157			.with_context(|| format!("reading agent binary {}", bin.display()))?;158		debug!("upload");159		upload_agent(sess, &dir, &path, bytes).await?;160	}161	Ok(path)162}163164async fn upload_agent(165	sess: &Handle<SshHandler>,166	dir: &Utf8Path,167	path: &Utf8Path,168	bytes: Vec<u8>,169) -> Result<()> {170	debug!("mkdirp");171	run_string_ok(sess, &format!("mkdir -p {}", sh_quote(dir))).await?;172173	let tmp = dir.join(format!("tmp.{}", Uuid::new_v4()));174	let ch = sess.channel_open_session().await?;175	debug!("cat");176	ch.exec(true, format!("cat > {}", sh_quote(&tmp))).await?;177178	let mut child = SshExecChild::from_exec(ch);179	child180		.stdin181		.write_all(&bytes)182		.await183		.context("sending agent binary")?;184	child185		.stdin186		.shutdown()187		.await188		.context("sending agent binary")?;189	let code = child.wait().await;190	ensure!(code == Some(0), "agent upload failed (exit {code:?})");191192	debug!("chmod");193	run_string_ok(sess, &format!("chmod 0755 {}", sh_quote(&tmp))).await?;194	run_string_ok(195		sess,196		&format!("mv -f {} {}", sh_quote(&tmp), sh_quote(path)),197	)198	.await?;199	Ok(())200}201202pub struct SshHandler {203	host: String,204	port: u16,205	subs: Subs,206}207impl Handler for SshHandler {208	type Error = russh::Error;209	async fn check_server_key(210		&mut self,211		server_public_key: &PublicKey,212	) -> Result<bool, Self::Error> {213		Ok(check_known_hosts(&self.host, self.port, server_public_key)?)214	}215	async fn server_channel_open_forwarded_streamlocal(216		&mut self,217		channel: Channel<Msg>,218		socket_path: &str,219		_session: &mut Session,220	) -> Result<(), Self::Error> {221		let Some(ch) = self222			.subs223			.lock()224			.expect("lock")225			.remove(&Utf8PathBuf::from(socket_path))226		else {227			return Err(russh::Error::WrongChannel);228		};229		let _ = ch.send(channel);230		Ok(())231	}232}233234enum Transport {235	Ssh {236		sess: Arc<Handle<SshHandler>>,237		subs: Subs,238		runtime_dir: Utf8PathBuf,239		agent_path: Utf8PathBuf,240	},241	Local {242		agent_path: PathBuf,243		runtime_dir: Utf8PathBuf,244	},245}246247struct RemowtInner {248	transport: Transport,249	rpc: Rpc<BifConfig>,250	elevated: tokio::sync::OnceCell<()>,251	#[allow(dead_code)]252	children: Mutex<Vec<tokio::process::Child>>,253	_runtime_tmp: Option<TempDir>,254	user: String,255	iroh: IrohState,256}257258#[derive(Default)]259struct IrohState {260	conn: tokio::sync::OnceCell<iroh::endpoint::Connection>,261	#[allow(dead_code)]262	endpoint: Mutex<Option<iroh::Endpoint>>,263	subs: Arc<Mutex<HashMap<u64, oneshot::Sender<RemowtStream>>>>,264	next_token: AtomicU64,265	router: Mutex<Option<Arc<DatagramRouter>>>,266}267268#[derive(Clone)]269pub struct Remowt(Arc<RemowtInner>);270271pub type RemowtRemote = Remote<BifConfig>;272273impl Remowt {274	/// Connect to the remote host over ssh, detect the architecture and deploy the required275	/// agent binary.276	pub async fn connect(host: &str, bundle: &AgentBundle, remowt_user: String) -> Result<Self> {277		let conf = russh_config::parse_home(host)?;278		let port = conf.host_config.port.or(conf.port).unwrap_or(22);279		let hostname = conf280			.host_config281			.hostname282			.clone()283			.unwrap_or_else(|| conf.host_name.clone());284		let user = conf285			.user286			.clone()287			.unwrap_or_else(|| env::var("USER").unwrap_or_else(|_| "root".to_owned()));288289		let subs: Subs = Arc::new(Mutex::new(HashMap::new()));290		let config = Config {291			nodelay: true,292			..Config::default()293		};294		let mut sess = connect(295			Arc::new(config),296			(hostname.clone(), port),297			SshHandler {298				host: hostname,299				port,300				subs: subs.clone(),301			},302		)303		.await?;304305		let mut agent = AgentClient::connect_env().await?;306		let rsa_hash = sess.best_supported_rsa_hash().await?.flatten();307		let mut authenticated = false;308		for ident in agent.request_identities().await? {309			let AgentIdentity::PublicKey { key, .. } = ident else {310				continue;311			};312			if sess313				.authenticate_publickey_with(user.clone(), key, rsa_hash, &mut agent)314				.await?315				.success()316			{317				authenticated = true;318				break;319			}320		}321		ensure!(authenticated, "ssh authentication failed");322323		let sess = Arc::new(sess);324325		debug!("deploying agent");326		let agent_path = deploy_agent(&sess, bundle).await?;327328		debug!("runtime dir");329		let runtime_dir = remote_runtime_dir(&sess).await?;330331		let rpc = Rpc::<BifConfig>::new(Address::User);332333		let cmd_chan = sess.channel_open_session().await?;334		debug!("starting agent");335		cmd_chan336			.exec(true, format!("{} real-agent", sh_quote(&agent_path)))337			.await?;338339		let child = SshExecChild::from_exec(cmd_chan);340		drain_to_tracing(child.stderr, "agent".to_owned(), true);341		rpc.add_direct(342			Address::Agent,343			child_port(child.stdout, child.stdin),344			Rtt(0),345		);346347		let remowt = Self(Arc::new(RemowtInner {348			transport: Transport::Ssh {349				sess,350				subs,351				runtime_dir,352				agent_path,353			},354			rpc,355			elevated: tokio::sync::OnceCell::new(),356			children: Mutex::new(Vec::new()),357			_runtime_tmp: None,358			user: remowt_user,359			iroh: IrohState::default(),360		}));361		remowt.setup_iroh().await;362		Ok(remowt)363	}364365	/// "Connect" to the local machine's agent, by starting the agent binary locally.366	pub async fn connect_local(bundle: &AgentBundle, user: String) -> Result<Self> {367		let agent_path = bundle.local_binary()?;368		let mut child = tokio::process::Command::new(&agent_path)369			.arg("real-agent")370			.arg("--local")371			.stdin(std::process::Stdio::piped())372			.stdout(std::process::Stdio::piped())373			.kill_on_drop(true)374			.spawn()375			.with_context(|| format!("spawning agent binary {}", agent_path.display()))?;376		let stdin = child.stdin.take().expect("stdin piped");377		let stdout = child.stdout.take().expect("stdout piped");378379		let rpc = Rpc::<BifConfig>::new(Address::User);380		rpc.add_direct(Address::Agent, child_port(stdout, stdin), Rtt(0));381382		let (runtime_dir, runtime_tmp) = local_runtime_dir()?;383384		Ok(Self(Arc::new(RemowtInner {385			transport: Transport::Local {386				agent_path,387				runtime_dir,388			},389			rpc,390			elevated: tokio::sync::OnceCell::new(),391			children: Mutex::new(vec![child]),392			_runtime_tmp: runtime_tmp,393			user,394			iroh: IrohState::default(),395		})))396	}397398	/// Get the handle to the underlying russh session handle.399	pub fn ssh(&self) -> Option<Arc<Handle<SshHandler>>> {400		match &self.0.transport {401			Transport::Ssh { sess, .. } => Some(sess.clone()),402			Transport::Local { .. } => None,403		}404	}405406	pub fn rpc(&self) -> Rpc<BifConfig> {407		self.0.rpc.clone()408	}409410	pub async fn load_plugin(&self, id: u16, name: &str) -> Result<()> {411		let client: PluginEndpointsClient<BifConfig> = self.endpoints();412		client413			.load_plugin(id, name.to_owned())414			.await?415			.map_err(|e| anyhow!("agent failed to load plugin: {e}"))416	}417	pub async fn run0_load_plugin_path(&self, id: u16, path: &str) -> Result<()> {418		self.ensure_escalated().await?;419		let client: PluginEndpointsClient<BifConfig> =420			PluginEndpointsClient::wrap(self.0.rpc.remote(Address::AgentPrivileged));421		client422			.load_plugin_path(id, path.to_owned())423			.await?424			.map_err(|e| anyhow!("privileged agent failed to load plugin: {e}"))425	}426	pub fn plugin_endpoints<R: RemoteEndpoints<BifConfig>>(&self, id: u16) -> R {427		R::wrap(self.0.rpc.remote(Address::Plugin(id)))428	}429430	pub fn endpoints<R: RemoteEndpoints<BifConfig>>(&self) -> R {431		R::wrap(self.0.rpc.remote(Address::Agent))432	}433	pub async fn run0_endpoints<R: RemoteEndpoints<BifConfig>>(&self) -> Result<R> {434		self.ensure_escalated().await?;435		Ok(R::wrap(self.0.rpc.remote(Address::AgentPrivileged)))436	}437438	async fn ensure_escalated(&self) -> Result<()> {439		self.0440			.elevated441			.get_or_try_init(|| async {442				let (agent_path, local) = match &self.0.transport {443					Transport::Ssh { agent_path, .. } => (agent_path.as_str().to_owned(), false),444					Transport::Local { agent_path, .. } => (445						agent_path446							.to_str()447							.ok_or_else(|| anyhow!("local agent path is not utf-8"))?448							.to_owned(),449						true,450					),451				};452453				let (tool, flags) = self.detect_escalation().await?;454				let mut args: Vec<String> = Vec::new();455				args.push("-w".to_owned());456				args.push(tool.to_owned());457				args.extend(flags.iter().copied().map(str::to_owned));458				if tool == "run0" {459					args.push(format!(460						"--unit={}-{}.service",461						self.0.user,462						Uuid::new_v4().simple()463					));464				}465				args.push(agent_path);466				args.push("real-agent".to_owned());467				args.push("--privileged".to_owned());468				if local {469					args.push("--local".to_owned());470				}471472				let child = self473					.spawn(SpawnOptions {474						program: "setsid".to_owned(),475						args,476						stdin: StdioMode::Pipe,477						stdout: StdioMode::Pipe,478						stderr: StderrMode::Inherit,479						..Default::default()480					})481					.await482					.context("spawning privileged agent")?;483484				let stdin = child485					.stdin486					.ok_or_else(|| anyhow!("privileged agent stdin missing"))?;487				let stdout = child488					.stdout489					.ok_or_else(|| anyhow!("privileged agent stdout missing"))?;490491				let port = child_port(stdout, stdin);492				self.0493					.rpc494					.add_direct(Address::AgentPrivileged, port, Rtt(0));495				anyhow::Ok(())496			})497			.await?;498		Ok(())499	}500501	async fn detect_escalation(&self) -> Result<(&'static str, &'static [&'static str])> {502		for (tool, flags) in ESCALATORS {503			let probe = self504				.spawn(SpawnOptions {505					program: (*tool).to_owned(),506					args: vec!["--version".to_owned()],507					stdout: StdioMode::Null,508					stderr: StderrMode::Null,509					..Default::default()510				})511				.await;512			if let Ok(child) = probe {513				let _ = child.wait().await;514				return Ok((tool, flags));515			}516		}517		bail!("no escalation tool found")518	}519520	/// XDG_RUNTIME_DIR on the remote machine.521	pub fn runtime_dir(&self) -> Utf8PathBuf {522		match &self.0.transport {523			Transport::Ssh { runtime_dir, .. } => runtime_dir.clone(),524			Transport::Local { runtime_dir, .. } => runtime_dir.clone(),525		}526	}527528	/// Bind unix listener socket on the remote machine with auto-generated path, returning the path.529	pub async fn bind_runtime_unix(&self, hint: &str) -> Result<(RemowtListener, Utf8PathBuf)> {530		let sock = self531			.runtime_dir()532			.join(format!("remowt-{hint}-{}.sock", Uuid::new_v4()));533		let listener = self.bind_unix(&sock).await?;534		Ok((listener, sock))535	}536537	/// Bind unix listener socket on the remote machine on the specified path.538	pub async fn bind_unix(&self, path: &Utf8Path) -> Result<RemowtListener> {539		match &self.0.transport {540			Transport::Ssh { sess, subs, .. } => {541				let (tx, rx) = oneshot::channel();542				subs.lock().expect("lock").insert(path.to_owned(), tx);543				sess.streamlocal_forward(path.to_owned()).await?;544				Ok(RemowtListener::Ssh(rx))545			}546			Transport::Local { .. } => {547				let _ = std::fs::remove_file(path);548				Ok(RemowtListener::Local(549					UnixListener::bind(path)?,550					path.to_owned(),551				))552			}553		}554	}555556	/// Bind a data tunnel, preferring the iroh fast path when it is up.557	pub async fn bind_fast_tunnel(558		&self,559		hint: &str,560		escalated: bool,561	) -> Result<(RemowtListener, TunnelAddr)> {562		if !escalated && self.0.iroh.conn.get().is_some() {563			let token = self564				.0565				.iroh566				.next_token567				.fetch_add(1, std::sync::atomic::Ordering::Relaxed);568			let (tx, rx) = oneshot::channel();569			self.0.iroh.subs.lock().expect("lock").insert(token, tx);570			return Ok((RemowtListener::Iroh(rx), TunnelAddr::Iroh { token }));571		}572		let (listener, path) = self.bind_runtime_unix(hint).await?;573		Ok((listener, TunnelAddr::Unix(path)))574	}575576	async fn setup_iroh(&self) {577		if std::env::var_os("REMOWT_NO_IROH").is_some() {578			debug!("REMOWT_NO_IROH set, skipping iroh fast tunnel");579			return;580		}581		if !matches!(self.0.transport, Transport::Ssh { .. }) {582			return;583		}584		if let Err(e) = self.try_setup_iroh().await {585			warn!("iroh fast tunnel unavailable, using ssh: {e}");586		}587	}588589	async fn try_setup_iroh(&self) -> Result<()> {590		use remowt_endpoints::iroh_tunnel::IrohTunnelClient;591		use remowt_link_shared::iroh_tunnel::{REMOWT_ALPN, build_endpoint, ssh_custom_addr};592593		let (listener, sock) = self.bind_runtime_unix("iroh-xport").await?;594		let secret = iroh::SecretKey::generate();595		let client_id = secret.public();596597		let client: IrohTunnelClient<BifConfig> =598			IrohTunnelClient::wrap(self.0.rpc.remote(Address::Agent));599		let (accepted, agent_id) = tokio::join!(listener.accept(), client.setup(client_id, sock));600		let stream = accepted?;601		let agent_id = agent_id?.map_err(|e| anyhow!("agent iroh setup failed: {e}"))?;602603		let ep = build_endpoint(secret, stream, agent_id, false).await?;604		let addr = iroh::EndpointAddr::from_parts(605			agent_id,606			[iroh::TransportAddr::Custom(ssh_custom_addr(agent_id))],607		);608		let conn = ep.connect(addr, REMOWT_ALPN).await?;609		ensure!(conn.remote_id() == agent_id, "iroh peer identity mismatch");610		debug!("iroh fast tunnel established");611612		let subs = self.0.iroh.subs.clone();613		let accept_conn = conn.clone();614		tokio::spawn(async move {615			loop {616				match accept_conn.accept_bi().await {617					Ok((send, mut recv)) => {618						let subs = subs.clone();619						tokio::spawn(async move {620							let mut buf = [0u8; 8];621							if tokio::io::AsyncReadExt::read_exact(&mut recv, &mut buf)622								.await623								.is_err()624							{625								return;626							}627							let token = u64::from_be_bytes(buf);628							let tx = subs.lock().expect("lock").remove(&token);629							if let Some(tx) = tx {630								let _ = tx.send(RemowtStream::Iroh(IrohBiStream::new(send, recv)));631							}632						});633					}634					Err(e) => {635						debug!("iroh accept loop ended: {e}");636						break;637					}638				}639			}640		});641642		*self.0.iroh.router.lock().expect("lock") = Some(DatagramRouter::spawn(conn.clone()));643		*self.0.iroh.endpoint.lock().expect("lock") = Some(ep);644		let _ = self.0.iroh.conn.set(conn);645		Ok(())646	}647648	pub fn datagram_router(&self) -> Option<Arc<DatagramRouter>> {649		self.0.iroh.router.lock().expect("lock").clone()650	}651}652653pub(crate) fn drain_to_tracing(654	stream: impl AsyncRead + Unpin + 'static + Send,655	context: String,656	stderr: bool,657) -> JoinHandle<()> {658	tokio::spawn(async move {659		let mut reader = BufReader::new(stream);660		let mut buf = Vec::with_capacity(4096);661		loop {662			buf.clear();663			match reader.read_until(b'\n', &mut buf).await {664				Ok(0) => break,665				Ok(_) => {666					let line = String::from_utf8_lossy(buf.strip_suffix(b"\n").unwrap_or(&buf));667					if stderr {668						warn!(context = %context, "{line}");669					} else {670						info!(context = %context, "{line}");671					}672				}673				Err(e) => {674					warn!(context = %context, "child stdio read failed: {e}");675					break;676				}677			}678		}679	})680}681682fn local_runtime_dir() -> Result<(Utf8PathBuf, Option<TempDir>)> {683	if let Ok(dir) = env::var("XDG_RUNTIME_DIR")684		&& !dir.is_empty()685	{686		return Ok((Utf8PathBuf::from(dir), None));687	}688	let tmp = tempfile::Builder::new()689		.prefix("remowt.")690		.rand_bytes(12)691		.tempdir()?;692	let dir = Utf8PathBuf::from_path_buf(tmp.path().to_owned())693		.map_err(|p| anyhow!("temp dir {} is not utf-8", p.display()))?;694	Ok((dir, Some(tmp)))695}696697async fn remote_runtime_dir(sess: &Handle<SshHandler>) -> Result<Utf8PathBuf> {698	let dir = run_string_ok(sess, "echo \"$XDG_RUNTIME_DIR\"").await?;699	let dir = dir.trim();700	if dir.is_empty() {701		let tmp = run_string_ok(sess, "mktemp -d remowt.XXXXXXXXXXXX --tmpdir").await?;702		Ok(Utf8PathBuf::from(tmp))703	} else {704		Ok(Utf8PathBuf::from(dir))705	}706}