git.delta.rocks / fleet / refs/commits / 087938500f2e

difftreelog

source

remowt/crates/remowt-client/src/lib.rs19.6 KiBsourcehistory
1use std::collections::HashMap;2use std::env;3use std::path::PathBuf;4use std::sync::atomic::AtomicU64;5use std::sync::{Arc, Mutex};67use anyhow::{anyhow, bail, ensure, Context as _, Result};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::client::{connect, Config, Handle, Handler, Msg, Session};16use russh::keys::agent::client::AgentClient;17use russh::keys::agent::AgentIdentity;18use russh::keys::check_known_hosts;19use russh::keys::ssh_key::PublicKey;20use russh::Channel;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 mut sess = connect(291			Arc::new(Config::default()),292			(hostname.clone(), port),293			SshHandler {294				host: hostname,295				port,296				subs: subs.clone(),297			},298		)299		.await?;300301		let mut agent = AgentClient::connect_env().await?;302		let rsa_hash = sess.best_supported_rsa_hash().await?.flatten();303		let mut authenticated = false;304		for ident in agent.request_identities().await? {305			let AgentIdentity::PublicKey { key, .. } = ident else {306				continue;307			};308			if sess309				.authenticate_publickey_with(user.clone(), key, rsa_hash, &mut agent)310				.await?311				.success()312			{313				authenticated = true;314				break;315			}316		}317		ensure!(authenticated, "ssh authentication failed");318319		let sess = Arc::new(sess);320321		debug!("deploying agent");322		let agent_path = deploy_agent(&sess, bundle).await?;323324		debug!("runtime dir");325		let runtime_dir = remote_runtime_dir(&sess).await?;326327		let rpc = Rpc::<BifConfig>::new(Address::User);328329		let cmd_chan = sess.channel_open_session().await?;330		debug!("starting agent");331		cmd_chan332			.exec(true, format!("{} real-agent", sh_quote(&agent_path)))333			.await?;334335		let child = SshExecChild::from_exec(cmd_chan);336		drain_to_tracing(child.stderr, "agent".to_owned(), true);337		rpc.add_direct(338			Address::Agent,339			child_port(child.stdout, child.stdin),340			Rtt(0),341		);342343		let remowt = Self(Arc::new(RemowtInner {344			transport: Transport::Ssh {345				sess,346				subs,347				runtime_dir,348				agent_path,349			},350			rpc,351			elevated: tokio::sync::OnceCell::new(),352			children: Mutex::new(Vec::new()),353			_runtime_tmp: None,354			user: remowt_user,355			iroh: IrohState::default(),356		}));357		remowt.setup_iroh().await;358		Ok(remowt)359	}360361	/// "Connect" to the local machine's agent, by starting the agent binary locally.362	pub async fn connect_local(bundle: &AgentBundle, user: String) -> Result<Self> {363		let agent_path = bundle.local_binary()?;364		let mut child = tokio::process::Command::new(&agent_path)365			.arg("real-agent")366			.arg("--local")367			.stdin(std::process::Stdio::piped())368			.stdout(std::process::Stdio::piped())369			.kill_on_drop(true)370			.spawn()371			.with_context(|| format!("spawning agent binary {}", agent_path.display()))?;372		let stdin = child.stdin.take().expect("stdin piped");373		let stdout = child.stdout.take().expect("stdout piped");374375		let rpc = Rpc::<BifConfig>::new(Address::User);376		rpc.add_direct(Address::Agent, child_port(stdout, stdin), Rtt(0));377378		let (runtime_dir, runtime_tmp) = local_runtime_dir()?;379380		Ok(Self(Arc::new(RemowtInner {381			transport: Transport::Local {382				agent_path,383				runtime_dir,384			},385			rpc,386			elevated: tokio::sync::OnceCell::new(),387			children: Mutex::new(vec![child]),388			_runtime_tmp: runtime_tmp,389			user,390			iroh: IrohState::default(),391		})))392	}393394	/// Get the handle to the underlying russh session handle.395	pub fn ssh(&self) -> Option<Arc<Handle<SshHandler>>> {396		match &self.0.transport {397			Transport::Ssh { sess, .. } => Some(sess.clone()),398			Transport::Local { .. } => None,399		}400	}401402	pub fn rpc(&self) -> Rpc<BifConfig> {403		self.0.rpc.clone()404	}405406	pub async fn load_plugin(&self, id: u16, name: &str) -> Result<()> {407		let client: PluginEndpointsClient<BifConfig> = self.endpoints();408		client409			.load_plugin(id, name.to_owned())410			.await?411			.map_err(|e| anyhow!("agent failed to load plugin: {e}"))412	}413	pub async fn run0_load_plugin_path(&self, id: u16, path: &str) -> Result<()> {414		self.ensure_escalated().await?;415		let client: PluginEndpointsClient<BifConfig> =416			PluginEndpointsClient::wrap(self.0.rpc.remote(Address::AgentPrivileged));417		client418			.load_plugin_path(id, path.to_owned())419			.await?420			.map_err(|e| anyhow!("privileged agent failed to load plugin: {e}"))421	}422	pub fn plugin_endpoints<R: RemoteEndpoints<BifConfig>>(&self, id: u16) -> R {423		R::wrap(self.0.rpc.remote(Address::Plugin(id)))424	}425426	pub fn endpoints<R: RemoteEndpoints<BifConfig>>(&self) -> R {427		R::wrap(self.0.rpc.remote(Address::Agent))428	}429	pub async fn run0_endpoints<R: RemoteEndpoints<BifConfig>>(&self) -> Result<R> {430		self.ensure_escalated().await?;431		Ok(R::wrap(self.0.rpc.remote(Address::AgentPrivileged)))432	}433434	async fn ensure_escalated(&self) -> Result<()> {435		self.0436			.elevated437			.get_or_try_init(|| async {438				let (agent_path, local) = match &self.0.transport {439					Transport::Ssh { agent_path, .. } => (agent_path.as_str().to_owned(), false),440					Transport::Local { agent_path, .. } => (441						agent_path442							.to_str()443							.ok_or_else(|| anyhow!("local agent path is not utf-8"))?444							.to_owned(),445						true,446					),447				};448449				let (tool, flags) = self.detect_escalation().await?;450				let mut args: Vec<String> = Vec::new();451				args.push("-w".to_owned());452				args.push(tool.to_owned());453				args.extend(flags.iter().copied().map(str::to_owned));454				if tool == "run0" {455					args.push(format!(456						"--unit={}-{}.service",457						self.0.user,458						Uuid::new_v4().simple()459					));460				}461				args.push(agent_path);462				args.push("real-agent".to_owned());463				args.push("--privileged".to_owned());464				if local {465					args.push("--local".to_owned());466				}467468				let child = self469					.spawn(SpawnOptions {470						program: "setsid".to_owned(),471						args,472						stdin: StdioMode::Pipe,473						stdout: StdioMode::Pipe,474						stderr: StderrMode::Inherit,475						..Default::default()476					})477					.await478					.context("spawning privileged agent")?;479480				let stdin = child481					.stdin482					.ok_or_else(|| anyhow!("privileged agent stdin missing"))?;483				let stdout = child484					.stdout485					.ok_or_else(|| anyhow!("privileged agent stdout missing"))?;486487				let port = child_port(stdout, stdin);488				self.0489					.rpc490					.add_direct(Address::AgentPrivileged, port, Rtt(0));491				anyhow::Ok(())492			})493			.await?;494		Ok(())495	}496497	async fn detect_escalation(&self) -> Result<(&'static str, &'static [&'static str])> {498		for (tool, flags) in ESCALATORS {499			let probe = self500				.spawn(SpawnOptions {501					program: (*tool).to_owned(),502					args: vec!["--version".to_owned()],503					stdout: StdioMode::Null,504					stderr: StderrMode::Null,505					..Default::default()506				})507				.await;508			if let Ok(child) = probe {509				let _ = child.wait().await;510				return Ok((tool, flags));511			}512		}513		bail!("no escalation tool found")514	}515516	/// XDG_RUNTIME_DIR on the remote machine.517	pub fn runtime_dir(&self) -> Utf8PathBuf {518		match &self.0.transport {519			Transport::Ssh { runtime_dir, .. } => runtime_dir.clone(),520			Transport::Local { runtime_dir, .. } => runtime_dir.clone(),521		}522	}523524	/// Bind unix listener socket on the remote machine with auto-generated path, returning the path.525	pub async fn bind_runtime_unix(&self, hint: &str) -> Result<(RemowtListener, Utf8PathBuf)> {526		let sock = self527			.runtime_dir()528			.join(format!("remowt-{hint}-{}.sock", Uuid::new_v4()));529		let listener = self.bind_unix(&sock).await?;530		Ok((listener, sock))531	}532533	/// Bind unix listener socket on the remote machine on the specified path.534	pub async fn bind_unix(&self, path: &Utf8Path) -> Result<RemowtListener> {535		match &self.0.transport {536			Transport::Ssh { sess, subs, .. } => {537				let (tx, rx) = oneshot::channel();538				subs.lock().expect("lock").insert(path.to_owned(), tx);539				sess.streamlocal_forward(path.to_owned()).await?;540				Ok(RemowtListener::Ssh(rx))541			}542			Transport::Local { .. } => {543				let _ = std::fs::remove_file(path);544				Ok(RemowtListener::Local(545					UnixListener::bind(path)?,546					path.to_owned(),547				))548			}549		}550	}551552	/// Bind a data tunnel, preferring the iroh fast path when it is up. Escalated tunnels553	/// (the privileged agent, a separate process with no iroh connection of its own) and the554	/// local transport always use the ssh/unix path.555	pub async fn bind_fast_tunnel(556		&self,557		hint: &str,558		escalated: bool,559	) -> Result<(RemowtListener, TunnelAddr)> {560		if !escalated && self.0.iroh.conn.get().is_some() {561			let token = self562				.0563				.iroh564				.next_token565				.fetch_add(1, std::sync::atomic::Ordering::Relaxed);566			let (tx, rx) = oneshot::channel();567			self.0.iroh.subs.lock().expect("lock").insert(token, tx);568			return Ok((RemowtListener::Iroh(rx), TunnelAddr::Iroh { token }));569		}570		let (listener, path) = self.bind_runtime_unix(hint).await?;571		Ok((listener, TunnelAddr::Unix(path)))572	}573574	async fn setup_iroh(&self) {575		if std::env::var_os("REMOWT_NO_IROH").is_some() {576			debug!("REMOWT_NO_IROH set, skipping iroh fast tunnel");577			return;578		}579		if !matches!(self.0.transport, Transport::Ssh { .. }) {580			return;581		}582		if let Err(e) = self.try_setup_iroh().await {583			warn!("iroh fast tunnel unavailable, using ssh: {e}");584		}585	}586587	async fn try_setup_iroh(&self) -> Result<()> {588		use remowt_endpoints::iroh_tunnel::IrohTunnelClient;589		use remowt_link_shared::iroh_tunnel::{build_endpoint, ssh_custom_addr, REMOWT_ALPN};590591		let (listener, sock) = self.bind_runtime_unix("iroh-xport").await?;592		let secret = iroh::SecretKey::generate();593		let client_id = secret.public();594595		let client: IrohTunnelClient<BifConfig> =596			IrohTunnelClient::wrap(self.0.rpc.remote(Address::Agent));597		let (accepted, agent_id) = tokio::join!(listener.accept(), client.setup(client_id, sock));598		let stream = accepted?;599		let agent_id = agent_id?.map_err(|e| anyhow!("agent iroh setup failed: {e}"))?;600601		let ep = build_endpoint(secret, stream, agent_id, false).await?;602		let addr = iroh::EndpointAddr::from_parts(603			agent_id,604			[iroh::TransportAddr::Custom(ssh_custom_addr(agent_id))],605		);606		let conn = ep.connect(addr, REMOWT_ALPN).await?;607		ensure!(conn.remote_id() == agent_id, "iroh peer identity mismatch");608		info!("iroh fast tunnel established");609610		let subs = self.0.iroh.subs.clone();611		let accept_conn = conn.clone();612		tokio::spawn(async move {613			loop {614				match accept_conn.accept_bi().await {615					Ok((send, mut recv)) => {616						let subs = subs.clone();617						tokio::spawn(async move {618							let mut buf = [0u8; 8];619							if tokio::io::AsyncReadExt::read_exact(&mut recv, &mut buf)620								.await621								.is_err()622							{623								return;624							}625							let token = u64::from_be_bytes(buf);626							let tx = subs.lock().expect("lock").remove(&token);627							if let Some(tx) = tx {628								let _ = tx.send(RemowtStream::Iroh(IrohBiStream::new(send, recv)));629							}630						});631					}632					Err(e) => {633						debug!("iroh accept loop ended: {e}");634						break;635					}636				}637			}638		});639640		let log_conn = conn.clone();641		tokio::spawn(async move {642			tokio::time::sleep(std::time::Duration::from_secs(2)).await;643			for p in log_conn.paths().iter() {644				if p.is_selected() {645					info!(rtt = ?p.rtt(), "iroh selected path: {}", p.remote_addr());646				}647			}648		});649650		*self.0.iroh.router.lock().expect("lock") = Some(DatagramRouter::spawn(conn.clone()));651		*self.0.iroh.endpoint.lock().expect("lock") = Some(ep);652		let _ = self.0.iroh.conn.set(conn);653		Ok(())654	}655656	pub fn datagram_router(&self) -> Option<Arc<DatagramRouter>> {657		self.0.iroh.router.lock().expect("lock").clone()658	}659}660661pub(crate) fn drain_to_tracing(662	stream: impl AsyncRead + Unpin + 'static + Send,663	context: String,664	stderr: bool,665) -> JoinHandle<()> {666	tokio::spawn(async move {667		let mut reader = BufReader::new(stream);668		let mut buf = Vec::with_capacity(4096);669		loop {670			buf.clear();671			match reader.read_until(b'\n', &mut buf).await {672				Ok(0) => break,673				Ok(_) => {674					let line = String::from_utf8_lossy(buf.strip_suffix(b"\n").unwrap_or(&buf));675					if stderr {676						warn!(context = %context, "{line}");677					} else {678						info!(context = %context, "{line}");679					}680				}681				Err(e) => {682					warn!(context = %context, "child stdio read failed: {e}");683					break;684				}685			}686		}687	})688}689690fn local_runtime_dir() -> Result<(Utf8PathBuf, Option<TempDir>)> {691	if let Ok(dir) = env::var("XDG_RUNTIME_DIR") {692		if !dir.is_empty() {693			return Ok((Utf8PathBuf::from(dir), None));694		}695	}696	let tmp = tempfile::Builder::new()697		.prefix("remowt.")698		.rand_bytes(12)699		.tempdir()?;700	let dir = Utf8PathBuf::from_path_buf(tmp.path().to_owned())701		.map_err(|p| anyhow!("temp dir {} is not utf-8", p.display()))?;702	Ok((dir, Some(tmp)))703}704705async fn remote_runtime_dir(sess: &Handle<SshHandler>) -> Result<Utf8PathBuf> {706	let dir = run_string_ok(sess, "echo \"$XDG_RUNTIME_DIR\"").await?;707	let dir = dir.trim();708	if dir.is_empty() {709		let tmp = run_string_ok(sess, "mktemp -d remowt.XXXXXXXXXXXX --tmpdir").await?;710		Ok(Utf8PathBuf::from(tmp))711	} else {712		Ok(Utf8PathBuf::from(dir))713	}714}