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

difftreelog

source

remowt/cmds/remowt-agent/src/main.rs11.9 KiBsourcehistory
1use std::borrow::Cow;2use std::collections::{BTreeMap, HashMap};3use std::fs::Permissions;4use std::future::pending;5use std::io;6use std::os::unix::fs::PermissionsExt as _;7use std::path::PathBuf;8use std::sync::{Arc, Mutex, OnceLock};910use bifrostlink::declarative::RemoteEndpoints;11use bifrostlink::Rpc;12use bifrostlink_ports::stdio::from_stdio;13use bifrostlink_ports::unix_socket::from_socket;14use clap::Parser;15use remowt_endpoints::{16	forward::Forward, fs::Fs, iroh_tunnel::IrohTunnel, nix_daemon::NixDaemon, pty::Pty,17	subprocess::Subprocess, systemd::Systemd,18};19use remowt_link_shared::iroh_tunnel::TunnelDialer;20use remowt_link_shared::{editor::EditorEndpointsClient, Address, BifConfig};21use remowt_polkit_shared::{emphasize, BackendRequest, Identity, PidDisplay};22use remowt_ui_prompt::bifrost::PromptEndpointsClient;23use remowt_ui_prompt::rofi::RofiPrompter;24use remowt_ui_prompt::{PrependSourcePrompter, Prompter, Source};25use tokio::fs;26use tokio::net::UnixStream;27use tokio::runtime::Builder;28use tokio::task::AbortHandle;29use tracing::{debug, trace};30use zbus::fdo;31use zbus::zvariant::{OwnedValue, Str};32use zbus::{interface, proxy, Connection};33use zbus_polkit::policykit1::Subject;3435use self::helper::{Helper, SocketHelper, SuidHelper};3637pub mod askpass;38pub mod bus;39pub mod editor;40pub mod helper;4142struct CancelTaskOnDrop {43	tasks: Arc<Mutex<HashMap<String, AbortHandle>>>,44	handle: String,45}46impl Drop for CancelTaskOnDrop {47	fn drop(&mut self) {48		debug!("cancel on drop");49		if let Some(task) = self50			.tasks51			.lock()52			.expect("not poisoned")53			.remove(&self.handle)54		{55			task.abort();56		}57	}58}5960struct Agent<H, P> {61	tasks: Arc<Mutex<HashMap<String, AbortHandle>>>,62	helper: H,63	prompter: P,64}65impl<H, P> Agent<H, P> {66	fn new(helper: H, prompter: P) -> Self {67		Agent {68			tasks: Arc::new(Mutex::new(HashMap::new())),69			helper,70			prompter,71		}72	}73}7475#[interface(name = "org.freedesktop.PolicyKit1.AuthenticationAgent")]76impl<H, P> Agent<H, P>77where78	H: Helper + Clone + Send + Sync + 'static,79	P: Prompter + Clone + Send + Sync + 'static,80{81	/// BeginAuthentication method82	#[allow(clippy::too_many_arguments)]83	async fn begin_authentication(84		&self,85		action_id: String,86		message: String,87		_icon_name: String,88		mut details: BTreeMap<String, String>,89		cookie: String,90		identities: Vec<Identity>,91	) -> zbus::fdo::Result<()> {92		use std::fmt::Write;93		debug!("begin auth");94		let _cancel_guard = Arc::new(OnceLock::new());95		let task = {96			let helper = self.helper.clone();97			let prompter = self.prompter.clone();98			let cookie = cookie.clone();99			let _cancel_guard = _cancel_guard.clone();100			tokio::task::spawn(async move {101				let _cancel_guard = _cancel_guard.clone();102				trace!("conversation task");103				let mut description = format!("{message}\n\n<b>Action id:</b> {action_id}",);104				if let Some(subject) = details.remove("polkit.caller-pid") {105					let _ = write!(description, "\n<b>Caller:</b> ");106					if let Ok(pid) = subject.parse::<u32>() {107						let _ = write!(description, "{}", PidDisplay(pid));108					} else {109						let _ = write!(description, "{}", emphasize("invalid pid"));110					}111				}112				if let Some(subject) = details.remove("polkit.subject-pid") {113					let _ = write!(description, "\n<b>Subject:</b> ");114					if let Ok(pid) = subject.parse::<u32>() {115						let _ = write!(description, "{}", PidDisplay(pid));116					} else {117						let _ = write!(description, "{}", emphasize("invalid pid"));118					}119				}120				let mut prompter = PrependSourcePrompter {121					source: vec![Source(Cow::Borrowed("polkit agent"))],122					description: description.clone(),123					prompter,124				};125126				let identity_displays: Vec<String> =127					identities.iter().map(|v| v.to_string()).collect();128				let identity_displays: Vec<&str> =129					identity_displays.iter().map(|v| v.as_str()).collect();130				debug!("choose identity");131				let choosen_identity = match identity_displays.len() {132					0 => {133						return Err(fdo::Error::AuthFailed(134							"no identity to authenticate as".to_owned(),135						))136					}137					1 => 0,138					_ => {139						prompter140							.prompt_enum(141								"Identity",142								"Select identity to use for polkit authorization",143								&identity_displays,144								&[],145							)146							.await?147					}148				};149				debug!("identity chosen");150151				let _ = write!(152					description,153					"\n<b>Identity:</b> {}",154					identities[choosen_identity as usize]155				);156				prompter.description = description;157158				prompter.source.push(Source(Cow::Borrowed("polkit daemon")));159160				helper161					.help_me(162						&cookie,163						prompter,164						identities[choosen_identity as usize].clone(),165					)166					.await167					.map_err(|e| fdo::Error::Failed(e.to_string()))?;168				// let connection = Connection::system().await?;169				// let helper = PolkitHelperProxy::new(&connection).await?;170171				Ok(())172			})173		};174		self.tasks175			.lock()176			.unwrap()177			.insert(cookie.clone(), task.abort_handle());178		debug!("abort handle stored");179		let _ = _cancel_guard.set(CancelTaskOnDrop {180			tasks: self.tasks.clone(),181			handle: cookie.clone(),182		});183184		let _ = task.await;185186		Ok(())187	}188189	/// CancelAuthentication method190	async fn cancel_authentication(&self, cookie: &str) -> zbus::fdo::Result<()> {191		debug!("auth cancelled");192		if let Some(abort) = self.tasks.lock().unwrap().remove(cookie) {193			debug!("abort handle found");194			abort.abort();195		}196		// debug!("Authentication cancled ! {cookie}");197		Ok(())198	}199}200201const OBJ_PATH: &str = "/org/freedesktop/PolicyKit1/AuthenticationAgent";202203#[proxy(204	interface = "lach.PolkitHelper",205	default_service = "lach.polkit.helper1",206	default_path = "/lach/PolkitHelper"207)]208trait PolkitHelper {209	fn init_conversation(&self, request: BackendRequest) -> zbus::Result<()>;210}211212#[derive(Parser)]213enum Opts {214	AskPass {215		prompt: String,216		description: String,217	},218	Editor {219		/// Argument to nvim220		path: String,221	},222	/// Expose a remote TCP/UDP address on the local (User) machine, printing the local port.223	Forward {224		/// Protocol: `tcp` (ssh or iroh) or `udp` (iroh-only).225		proto: ForwardProto,226		/// Remote socket address to expose locally, e.g. `127.0.0.1:8080`.227		addr: String,228	},229	RealAgent {230		#[arg(long)]231		path: Option<PathBuf>,232		/// Expect own address to be AgentPrivileged, skip installing polkit agent233		#[arg(long)]234		privileged: bool,235		#[arg(long)]236		local: bool,237	},238	LocalAgent,239}240241#[derive(Clone, Copy, clap::ValueEnum)]242enum ForwardProto {243	Tcp,244	Udp,245}246247fn main() -> anyhow::Result<()> {248	tracing_subscriber::fmt()249		.with_writer(io::stderr)250		.without_time()251		.init();252	let opts = Opts::parse();253254	let runtime = Builder::new_current_thread().enable_all().build()?;255256	match opts {257		Opts::AskPass {258			prompt,259			description,260		} => runtime.block_on(askpass::ask(&prompt, description)),261		Opts::LocalAgent => runtime.block_on(main_real()),262		Opts::Editor { path } => runtime.block_on(editor::edit(path)),263		Opts::Forward { proto, addr } => {264			runtime.block_on(editor::forward(matches!(proto, ForwardProto::Udp), addr))265		}266		Opts::RealAgent {267			path,268			privileged,269			local,270		} => runtime.block_on(main_real_agent(path, privileged, local)),271	}272}273async fn main_real() -> anyhow::Result<()> {274	let system_conn = Connection::system().await?;275	let helper = SocketHelper {276		fallback: SuidHelper,277	};278	register_auth_agent(&system_conn, Agent::new(helper, RofiPrompter)).await?;279280	let session_conn = Connection::session().await?;281	askpass::serve(&session_conn, RofiPrompter).await?;282283	let _keep_alive = (system_conn, session_conn);284	pending().await285}286async fn main_real_agent(287	path: Option<PathBuf>,288	privileged: bool,289	local: bool,290) -> anyhow::Result<()> {291	let address = if privileged {292		Address::AgentPrivileged293	} else {294		Address::Agent295	};296	let mut rpc = Rpc::<BifConfig>::new(address);297298	let dialer = Arc::new(TunnelDialer::new());299	Fs::new().register_endpoints(&mut rpc);300	Systemd.register_endpoints(&mut rpc);301	Pty::new(dialer.clone()).register_endpoints(&mut rpc);302	Subprocess::new(dialer.clone()).register_endpoints(&mut rpc);303	NixDaemon::new(dialer.clone()).register_endpoints(&mut rpc);304	IrohTunnel::new(dialer.clone()).register_endpoints(&mut rpc);305	Forward::new(dialer.clone()).register_endpoints(&mut rpc);306307	remowt_plugin::host::serve(&mut rpc);308309	let user_prompter = PromptEndpointsClient::wrap(rpc.remote(Address::User));310	let editor_client = EditorEndpointsClient::wrap(rpc.remote(Address::User));311312	let bus = bus::spawn().await?;313	askpass::serve(&bus.conn, user_prompter.clone()).await?;314	editor::serve(&bus.conn, editor_client).await?;315316	let helpers = tempfile::Builder::new().prefix("remowt-path.").tempdir()?;317	let exe = std::env::current_exe()?;318	let askpass_helper = helpers.path().join("remowt-askpass");319	let editor_helper = helpers.path().join("remowt-editor");320	let forward_helper = helpers.path().join("remowt-forward");321	{322		let script = format!(323			"#!/bin/sh\nexec {} ask-pass \"password\" \"$1\"\n",324			sh_quote(&exe.to_string_lossy())325		);326		fs::write(&askpass_helper, script).await?;327		fs::set_permissions(&askpass_helper, Permissions::from_mode(0o755)).await?;328	}329	{330		let script = format!(331			"#!/bin/sh\nexec {} editor \"$1\"\n",332			sh_quote(&exe.to_string_lossy())333		);334		fs::write(&editor_helper, script).await?;335		fs::set_permissions(&editor_helper, Permissions::from_mode(0o755)).await?;336	}337	{338		let script = format!(339			"#!/bin/sh\nexec {} forward \"$@\"\n",340			sh_quote(&exe.to_string_lossy())341		);342		fs::write(&forward_helper, script).await?;343		fs::set_permissions(&forward_helper, Permissions::from_mode(0o755)).await?;344	}345346	// Safety: Hoping tokio own threads won't read any of those...347	unsafe {348		prepend_path(helpers.path());349		std::env::set_var("SUDO_ASKPASS", &askpass_helper);350		std::env::set_var("SSH_ASKPASS", &askpass_helper);351		std::env::set_var("SSH_ASKPASS_REQUIRE", "force");352		std::env::set_var("EDITOR", &editor_helper);353		std::env::set_var("VISUAL", &editor_helper);354		std::env::set_var("DBUS_SESSION_BUS_ADDRESS", &bus.address);355	}356357	let port = match path {358		Some(path) => from_socket(UnixStream::connect(path).await?),359		None => from_stdio(),360	};361	rpc.add_direct(Address::User, port, bifrostlink::Rtt(0));362363	let polkit_conn = if !privileged && !local {364		// The unprivileged agent doubles as a polkit authentication agent so365		// `run0` (e.g. our own elevation) routes its prompt to the User over366		// bifrost instead of failing on a tty-less session.367		let conn = Connection::system().await?;368		let helper = SocketHelper {369			fallback: SuidHelper,370		};371		register_auth_agent(&conn, Agent::new(helper, user_prompter)).await?;372		Some(conn)373	} else {374		None375	};376377	let _keep_alive = (bus, helpers, polkit_conn);378	pending().await379}380381async fn register_auth_agent<H, P>(conn: &Connection, agent: Agent<H, P>) -> anyhow::Result<()>382where383	H: Helper + Clone + Send + Sync + 'static,384	P: Prompter + Clone + Send + Sync + 'static,385{386	let proxy = zbus_polkit::policykit1::AuthorityProxy::new(conn).await?;387	conn.object_server().at(OBJ_PATH, agent).await?;388389	let subject = auth_agent_subject()?;390	proxy391		.register_authentication_agent(&subject, "C", OBJ_PATH)392		.await?;393	debug!(kind = subject.subject_kind, "registered polkit agent");394	Ok(())395}396397fn auth_agent_subject() -> anyhow::Result<Subject> {398	let mut details = HashMap::new();399	if let Ok(session_id) = std::env::var("XDG_SESSION_ID") {400		let val: OwnedValue = Str::from(session_id).into();401		details.insert("session-id".to_string(), val);402		return Ok(Subject {403			subject_kind: "unix-session".to_string(),404			subject_details: details,405		});406	}407408	details.insert("pid".to_string(), OwnedValue::from(std::process::id()));409	Ok(Subject {410		subject_kind: "unix-process".to_string(),411		subject_details: details,412	})413}414415fn sh_quote(s: &str) -> String {416	format!("'{}'", s.replace('\'', "'\\''"))417}418419/// Prepend `dir` to the process `PATH`.420///421/// # SAFETY422///423/// Same as `set_var`424unsafe fn prepend_path(dir: &std::path::Path) {425	let value = match std::env::var_os("PATH") {426		Some(existing) => {427			let mut v = dir.as_os_str().to_owned();428			v.push(":");429			v.push(existing);430			v431		}432		None => dir.as_os_str().to_owned(),433	};434	unsafe {435		std::env::set_var("PATH", value);436	}437}