git.delta.rocks / remowt / refs/commits / 960ee4bca68c

difftreelog

source

cmds/remowt-agent/src/main.rs10.1 KiBsourcehistory
1use std::borrow::Cow;2use std::collections::{BTreeMap, HashMap};3use std::fs::Permissions;4use std::future::pending;5use std::os::unix::fs::PermissionsExt as _;6use std::path::PathBuf;7use std::sync::{Arc, Mutex, OnceLock};89use bifrostlink::declarative::RemoteEndpoints;10use bifrostlink::Rpc;11use bifrostlink_ports::stdio::from_stdio;12use bifrostlink_ports::unix_socket::from_socket;13use clap::Parser;14use polkit_shared::{emphasize, BackendRequest, Identity, PidDisplay};15use remowt_link_shared::editor::EditorEndpointsClient;16use remowt_link_shared::{Address, BifConfig, Fs, Pty, Systemd};17use tokio::fs;18use tokio::net::UnixStream;19use tokio::runtime::Builder;20use tokio::task::AbortHandle;21use tracing::{info, trace};22use ui_prompt::bifrost::PromptEndpointsClient;23use ui_prompt::{PrependSourcePrompter, Prompter, Source};24use zbus::fdo;25use zbus::zvariant::{OwnedValue, Str};26use zbus::{interface, proxy, Connection};27use zbus_polkit::policykit1::Subject;2829use self::helper::{Helper, SocketHelper, SuidHelper};3031pub mod askpass;32pub mod bus;33pub mod editor;34pub mod helper;3536struct CancelTaskOnDrop {37	tasks: Arc<Mutex<HashMap<String, AbortHandle>>>,38	handle: String,39}40impl Drop for CancelTaskOnDrop {41	fn drop(&mut self) {42		info!("cancel on drop");43		if let Some(task) = self44			.tasks45			.lock()46			.expect("not poisoned")47			.remove(&self.handle)48		{49			task.abort();50		}51	}52}5354struct Agent<H, P> {55	tasks: Arc<Mutex<HashMap<String, AbortHandle>>>,56	helper: H,57	prompter: P,58}59impl<H, P> Agent<H, P> {60	fn new(helper: H, prompter: P) -> Self {61		Agent {62			tasks: Arc::new(Mutex::new(HashMap::new())),63			helper,64			prompter,65		}66	}67}6869#[interface(name = "org.freedesktop.PolicyKit1.AuthenticationAgent")]70impl<H, P> Agent<H, P>71where72	H: Helper + Clone + Send + Sync + 'static,73	P: Prompter + Clone + Send + Sync + 'static,74{75	/// BeginAuthentication method76	#[allow(clippy::too_many_arguments)]77	async fn begin_authentication(78		&self,79		action_id: String,80		message: String,81		_icon_name: String,82		mut details: BTreeMap<String, String>,83		cookie: String,84		identities: Vec<Identity>,85	) -> zbus::fdo::Result<()> {86		use std::fmt::Write;87		info!("begin auth");88		let _cancel_guard = Arc::new(OnceLock::new());89		let task = {90			let helper = self.helper.clone();91			let prompter = self.prompter.clone();92			let cookie = cookie.clone();93			let _cancel_guard = _cancel_guard.clone();94			tokio::task::spawn(async move {95				let _cancel_guard = _cancel_guard.clone();96				trace!("conversation task");97				let mut description = format!("{message}\n\n<b>Action id:</b> {action_id}",);98				if let Some(subject) = details.remove("polkit.caller-pid") {99					let _ = write!(description, "\n<b>Caller:</b> ");100					if let Ok(pid) = subject.parse::<u32>() {101						let _ = write!(description, "{}", PidDisplay(pid));102					} else {103						let _ = write!(description, "{}", emphasize("invalid pid"));104					}105				}106				if let Some(subject) = details.remove("polkit.subject-pid") {107					let _ = write!(description, "\n<b>Subject:</b> ");108					if let Ok(pid) = subject.parse::<u32>() {109						let _ = write!(description, "{}", PidDisplay(pid));110					} else {111						let _ = write!(description, "{}", emphasize("invalid pid"));112					}113				}114				let mut prompter = PrependSourcePrompter {115					source: vec![Source(Cow::Borrowed("polkit agent"))],116					description: description.clone(),117					prompter,118				};119120				let identity_displays: Vec<String> =121					identities.iter().map(|v| v.to_string()).collect();122				let identity_displays: Vec<&str> =123					identity_displays.iter().map(|v| v.as_str()).collect();124				info!("choose identity");125				let choosen_identity = match identity_displays.len() {126					0 => {127						return Err(fdo::Error::AuthFailed(128							"no identity to authenticate as".to_owned(),129						))130					}131					1 => 0,132					_ => {133						prompter134							.prompt_enum(135								"Identity",136								"Select identity to use for polkit authorization",137								&identity_displays,138								&[],139							)140							.await?141					}142				};143				info!("identity chosen");144145				let _ = write!(146					description,147					"\n<b>Identity:</b> {}",148					identities[choosen_identity as usize]149				);150				prompter.description = description;151152				prompter.source.push(Source(Cow::Borrowed("polkit daemon")));153154				helper155					.help_me(156						&cookie,157						prompter,158						identities[choosen_identity as usize].clone(),159					)160					.await161					.map_err(|e| fdo::Error::Failed(e.to_string()))?;162				// let connection = Connection::system().await?;163				// let helper = PolkitHelperProxy::new(&connection).await?;164165				Ok(())166			})167		};168		self.tasks169			.lock()170			.unwrap()171			.insert(cookie.clone(), task.abort_handle());172		info!("abort handle stored");173		let _ = _cancel_guard.set(CancelTaskOnDrop {174			tasks: self.tasks.clone(),175			handle: cookie.clone(),176		});177178		let _ = task.await;179180		Ok(())181	}182183	/// CancelAuthentication method184	async fn cancel_authentication(&self, cookie: &str) -> zbus::fdo::Result<()> {185		info!("auth cancelled");186		if let Some(abort) = self.tasks.lock().unwrap().remove(cookie) {187			info!("abort handle found");188			abort.abort();189		}190		// debug!("Authentication cancled ! {cookie}");191		Ok(())192	}193}194195const OBJ_PATH: &str = "/org/freedesktop/PolicyKit1/AuthenticationAgent";196197#[proxy(198	interface = "lach.PolkitHelper",199	default_service = "lach.polkit.helper1",200	default_path = "/lach/PolkitHelper"201)]202trait PolkitHelper {203	fn init_conversation(&self, request: BackendRequest) -> zbus::Result<()>;204}205206#[derive(Parser)]207enum Opts {208	AskPass {209		prompt: String,210		description: String,211	},212	Editor {213		/// Argument to nvim214		path: String,215	},216	RealAgent {217		#[arg(long)]218		path: Option<PathBuf>,219		/// Expect own address to be AgentPrivileged, skip installing polkit agent220		#[arg(long)]221		privileged: bool,222	},223}224225fn main() -> anyhow::Result<()> {226	// Log to stderr: `privileged-agent` uses stdout as the bifrost transport,227	// so anything written there would corrupt the stream.228	tracing_subscriber::fmt()229		.with_writer(std::io::stderr)230		.init();231	let opts = Opts::parse();232233	let runtime = Builder::new_current_thread().enable_all().build()?;234235	match opts {236		Opts::AskPass {237			prompt,238			description,239		} => runtime.block_on(askpass::ask(&prompt, description)),240		Opts::Editor { path } => runtime.block_on(editor::edit(path)),241		Opts::RealAgent { path, privileged } => runtime.block_on(main_real_agent(path, privileged)),242	}243}244async fn main_real_agent(path: Option<PathBuf>, privileged: bool) -> anyhow::Result<()> {245	let address = if privileged {246		Address::AgentPrivileged247	} else {248		Address::Agent249	};250	let mut rpc = Rpc::<BifConfig>::new(address);251252	Fs::new().register_endpoints(&mut rpc);253	Systemd.register_endpoints(&mut rpc);254	Pty::new().register_endpoints(&mut rpc);255256	let user_prompter = PromptEndpointsClient::wrap(rpc.remote(Address::User));257	let editor_client = EditorEndpointsClient::wrap(rpc.remote(Address::User));258259	let bus = bus::spawn().await?;260	askpass::serve(&bus.conn, user_prompter.clone()).await?;261	editor::serve(&bus.conn, editor_client).await?;262263	let helpers = tempfile::Builder::new().prefix("remowt-path.").tempdir()?;264	let exe = std::env::current_exe()?;265	let askpass_helper = helpers.path().join("remowt-askpass");266	let editor_helper = helpers.path().join("remowt-editor");267	{268		let script = format!(269			"#!/bin/sh\nexec {} ask-pass \"password\" \"$1\"\n",270			sh_quote(&exe.to_string_lossy())271		);272		fs::write(&askpass_helper, script).await?;273		fs::set_permissions(&askpass_helper, Permissions::from_mode(0o755)).await?;274	}275	{276		let script = format!(277			"#!/bin/sh\nexec {} editor \"$1\"\n",278			sh_quote(&exe.to_string_lossy())279		);280		fs::write(&editor_helper, script).await?;281		fs::set_permissions(&editor_helper, Permissions::from_mode(0o755)).await?;282	}283284	// Safety: Hoping tokio own threads won't read any of those...285	unsafe {286		prepend_path(helpers.path());287		std::env::set_var("SUDO_ASKPASS", &askpass_helper);288		std::env::set_var("SSH_ASKPASS", &askpass_helper);289		std::env::set_var("SSH_ASKPASS_REQUIRE", "force");290		std::env::set_var("EDITOR", &editor_helper);291		std::env::set_var("VISUAL", &editor_helper);292		std::env::set_var("DBUS_SESSION_BUS_ADDRESS", &bus.address);293	}294295	let port = match path {296		Some(path) => from_socket(UnixStream::connect(path).await?),297		None => from_stdio(),298	};299	rpc.add_direct(Address::User, port, bifrostlink::Rtt(0));300301	let polkit_conn = if !privileged {302		// The unprivileged agent doubles as a polkit authentication agent so303		// `run0` (e.g. our own elevation) routes its prompt to the User over304		// bifrost instead of failing on a tty-less session.305		let conn = Connection::system().await?;306		let helper = SocketHelper {307			fallback: SuidHelper,308		};309		register_auth_agent(&conn, Agent::new(helper, user_prompter)).await?;310		Some(conn)311	} else {312		None313	};314315	let _keep_alive = (bus, helpers, polkit_conn);316	pending().await317}318319async fn register_auth_agent<H, P>(conn: &Connection, agent: Agent<H, P>) -> anyhow::Result<()>320where321	H: Helper + Clone + Send + Sync + 'static,322	P: Prompter + Clone + Send + Sync + 'static,323{324	let proxy = zbus_polkit::policykit1::AuthorityProxy::new(conn).await?;325	conn.object_server().at(OBJ_PATH, agent).await?;326327	let subject = auth_agent_subject()?;328	proxy329		.register_authentication_agent(&subject, "C", OBJ_PATH)330		.await?;331	info!(kind = subject.subject_kind, "registered polkit agent");332	Ok(())333}334335fn auth_agent_subject() -> anyhow::Result<Subject> {336	let mut details = HashMap::new();337	if let Ok(session_id) = std::env::var("XDG_SESSION_ID") {338		let val: OwnedValue = Str::from(session_id).into();339		details.insert("session-id".to_string(), val);340		return Ok(Subject {341			subject_kind: "unix-session".to_string(),342			subject_details: details,343		});344	}345346	details.insert("pid".to_string(), OwnedValue::from(std::process::id()));347	Ok(Subject {348		subject_kind: "unix-process".to_string(),349		subject_details: details,350	})351}352353fn sh_quote(s: &str) -> String {354	format!("'{}'", s.replace('\'', "'\\''"))355}356357/// Prepend `dir` to the process `PATH`.358///359/// # SAFETY360///361/// Same as `set_var`362unsafe fn prepend_path(dir: &std::path::Path) {363	let value = match std::env::var_os("PATH") {364		Some(existing) => {365			let mut v = dir.as_os_str().to_owned();366			v.push(":");367			v.push(existing);368			v369		}370		None => dir.as_os_str().to_owned(),371	};372	unsafe {373		std::env::set_var("PATH", value);374	}375}