git.delta.rocks / remowt / refs/commits / 6d9cf16dada2

difftreelog

source

cmds/remowt-agent/src/main.rs10.8 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	fs::Fs, nix_daemon::NixDaemon, pty::Pty, subprocess::Subprocess, systemd::Systemd,17};18use remowt_link_shared::{editor::EditorEndpointsClient, Address, BifConfig};19use remowt_polkit_shared::{emphasize, BackendRequest, Identity, PidDisplay};20use remowt_ui_prompt::bifrost::PromptEndpointsClient;21use remowt_ui_prompt::rofi::RofiPrompter;22use remowt_ui_prompt::{PrependSourcePrompter, Prompter, Source};23use tokio::fs;24use tokio::net::UnixStream;25use tokio::runtime::Builder;26use tokio::task::AbortHandle;27use tracing::{debug, trace};28use zbus::fdo;29use zbus::zvariant::{OwnedValue, Str};30use zbus::{interface, proxy, Connection};31use zbus_polkit::policykit1::Subject;3233use self::helper::{Helper, SocketHelper, SuidHelper};3435pub mod askpass;36pub mod bus;37pub mod editor;38pub mod helper;3940struct CancelTaskOnDrop {41	tasks: Arc<Mutex<HashMap<String, AbortHandle>>>,42	handle: String,43}44impl Drop for CancelTaskOnDrop {45	fn drop(&mut self) {46		debug!("cancel on drop");47		if let Some(task) = self48			.tasks49			.lock()50			.expect("not poisoned")51			.remove(&self.handle)52		{53			task.abort();54		}55	}56}5758struct Agent<H, P> {59	tasks: Arc<Mutex<HashMap<String, AbortHandle>>>,60	helper: H,61	prompter: P,62}63impl<H, P> Agent<H, P> {64	fn new(helper: H, prompter: P) -> Self {65		Agent {66			tasks: Arc::new(Mutex::new(HashMap::new())),67			helper,68			prompter,69		}70	}71}7273#[interface(name = "org.freedesktop.PolicyKit1.AuthenticationAgent")]74impl<H, P> Agent<H, P>75where76	H: Helper + Clone + Send + Sync + 'static,77	P: Prompter + Clone + Send + Sync + 'static,78{79	/// BeginAuthentication method80	#[allow(clippy::too_many_arguments)]81	async fn begin_authentication(82		&self,83		action_id: String,84		message: String,85		_icon_name: String,86		mut details: BTreeMap<String, String>,87		cookie: String,88		identities: Vec<Identity>,89	) -> zbus::fdo::Result<()> {90		use std::fmt::Write;91		debug!("begin auth");92		let _cancel_guard = Arc::new(OnceLock::new());93		let task = {94			let helper = self.helper.clone();95			let prompter = self.prompter.clone();96			let cookie = cookie.clone();97			let _cancel_guard = _cancel_guard.clone();98			tokio::task::spawn(async move {99				let _cancel_guard = _cancel_guard.clone();100				trace!("conversation task");101				let mut description = format!("{message}\n\n<b>Action id:</b> {action_id}",);102				if let Some(subject) = details.remove("polkit.caller-pid") {103					let _ = write!(description, "\n<b>Caller:</b> ");104					if let Ok(pid) = subject.parse::<u32>() {105						let _ = write!(description, "{}", PidDisplay(pid));106					} else {107						let _ = write!(description, "{}", emphasize("invalid pid"));108					}109				}110				if let Some(subject) = details.remove("polkit.subject-pid") {111					let _ = write!(description, "\n<b>Subject:</b> ");112					if let Ok(pid) = subject.parse::<u32>() {113						let _ = write!(description, "{}", PidDisplay(pid));114					} else {115						let _ = write!(description, "{}", emphasize("invalid pid"));116					}117				}118				let mut prompter = PrependSourcePrompter {119					source: vec![Source(Cow::Borrowed("polkit agent"))],120					description: description.clone(),121					prompter,122				};123124				let identity_displays: Vec<String> =125					identities.iter().map(|v| v.to_string()).collect();126				let identity_displays: Vec<&str> =127					identity_displays.iter().map(|v| v.as_str()).collect();128				debug!("choose identity");129				let choosen_identity = match identity_displays.len() {130					0 => {131						return Err(fdo::Error::AuthFailed(132							"no identity to authenticate as".to_owned(),133						))134					}135					1 => 0,136					_ => {137						prompter138							.prompt_enum(139								"Identity",140								"Select identity to use for polkit authorization",141								&identity_displays,142								&[],143							)144							.await?145					}146				};147				debug!("identity chosen");148149				let _ = write!(150					description,151					"\n<b>Identity:</b> {}",152					identities[choosen_identity as usize]153				);154				prompter.description = description;155156				prompter.source.push(Source(Cow::Borrowed("polkit daemon")));157158				helper159					.help_me(160						&cookie,161						prompter,162						identities[choosen_identity as usize].clone(),163					)164					.await165					.map_err(|e| fdo::Error::Failed(e.to_string()))?;166				// let connection = Connection::system().await?;167				// let helper = PolkitHelperProxy::new(&connection).await?;168169				Ok(())170			})171		};172		self.tasks173			.lock()174			.unwrap()175			.insert(cookie.clone(), task.abort_handle());176		debug!("abort handle stored");177		let _ = _cancel_guard.set(CancelTaskOnDrop {178			tasks: self.tasks.clone(),179			handle: cookie.clone(),180		});181182		let _ = task.await;183184		Ok(())185	}186187	/// CancelAuthentication method188	async fn cancel_authentication(&self, cookie: &str) -> zbus::fdo::Result<()> {189		debug!("auth cancelled");190		if let Some(abort) = self.tasks.lock().unwrap().remove(cookie) {191			debug!("abort handle found");192			abort.abort();193		}194		// debug!("Authentication cancled ! {cookie}");195		Ok(())196	}197}198199const OBJ_PATH: &str = "/org/freedesktop/PolicyKit1/AuthenticationAgent";200201#[proxy(202	interface = "lach.PolkitHelper",203	default_service = "lach.polkit.helper1",204	default_path = "/lach/PolkitHelper"205)]206trait PolkitHelper {207	fn init_conversation(&self, request: BackendRequest) -> zbus::Result<()>;208}209210#[derive(Parser)]211enum Opts {212	AskPass {213		prompt: String,214		description: String,215	},216	Editor {217		/// Argument to nvim218		path: String,219	},220	RealAgent {221		#[arg(long)]222		path: Option<PathBuf>,223		/// Expect own address to be AgentPrivileged, skip installing polkit agent224		#[arg(long)]225		privileged: bool,226		#[arg(long)]227		local: bool,228	},229	LocalAgent,230}231232fn main() -> anyhow::Result<()> {233	tracing_subscriber::fmt()234		.with_writer(io::stderr)235		.without_time()236		.init();237	let opts = Opts::parse();238239	let runtime = Builder::new_current_thread().enable_all().build()?;240241	match opts {242		Opts::AskPass {243			prompt,244			description,245		} => runtime.block_on(askpass::ask(&prompt, description)),246		Opts::LocalAgent => runtime.block_on(main_real()),247		Opts::Editor { path } => runtime.block_on(editor::edit(path)),248		Opts::RealAgent {249			path,250			privileged,251			local,252		} => runtime.block_on(main_real_agent(path, privileged, local)),253	}254}255async fn main_real() -> anyhow::Result<()> {256	let system_conn = Connection::system().await?;257	let helper = SocketHelper {258		fallback: SuidHelper,259	};260	register_auth_agent(&system_conn, Agent::new(helper, RofiPrompter)).await?;261262	let session_conn = Connection::session().await?;263	askpass::serve(&session_conn, RofiPrompter).await?;264265	let _keep_alive = (system_conn, session_conn);266	pending().await267}268async fn main_real_agent(269	path: Option<PathBuf>,270	privileged: bool,271	local: bool,272) -> anyhow::Result<()> {273	let address = if privileged {274		Address::AgentPrivileged275	} else {276		Address::Agent277	};278	let mut rpc = Rpc::<BifConfig>::new(address);279280	Fs::new().register_endpoints(&mut rpc);281	Systemd.register_endpoints(&mut rpc);282	Pty::new().register_endpoints(&mut rpc);283	Subprocess::new().register_endpoints(&mut rpc);284	NixDaemon.register_endpoints(&mut rpc);285286	remowt_plugin::host::serve(&mut rpc);287288	let user_prompter = PromptEndpointsClient::wrap(rpc.remote(Address::User));289	let editor_client = EditorEndpointsClient::wrap(rpc.remote(Address::User));290291	let bus = bus::spawn().await?;292	askpass::serve(&bus.conn, user_prompter.clone()).await?;293	editor::serve(&bus.conn, editor_client).await?;294295	let helpers = tempfile::Builder::new().prefix("remowt-path.").tempdir()?;296	let exe = std::env::current_exe()?;297	let askpass_helper = helpers.path().join("remowt-askpass");298	let editor_helper = helpers.path().join("remowt-editor");299	{300		let script = format!(301			"#!/bin/sh\nexec {} ask-pass \"password\" \"$1\"\n",302			sh_quote(&exe.to_string_lossy())303		);304		fs::write(&askpass_helper, script).await?;305		fs::set_permissions(&askpass_helper, Permissions::from_mode(0o755)).await?;306	}307	{308		let script = format!(309			"#!/bin/sh\nexec {} editor \"$1\"\n",310			sh_quote(&exe.to_string_lossy())311		);312		fs::write(&editor_helper, script).await?;313		fs::set_permissions(&editor_helper, Permissions::from_mode(0o755)).await?;314	}315316	// Safety: Hoping tokio own threads won't read any of those...317	unsafe {318		prepend_path(helpers.path());319		std::env::set_var("SUDO_ASKPASS", &askpass_helper);320		std::env::set_var("SSH_ASKPASS", &askpass_helper);321		std::env::set_var("SSH_ASKPASS_REQUIRE", "force");322		std::env::set_var("EDITOR", &editor_helper);323		std::env::set_var("VISUAL", &editor_helper);324		std::env::set_var("DBUS_SESSION_BUS_ADDRESS", &bus.address);325	}326327	let port = match path {328		Some(path) => from_socket(UnixStream::connect(path).await?),329		None => from_stdio(),330	};331	rpc.add_direct(Address::User, port, bifrostlink::Rtt(0));332333	let polkit_conn = if !privileged && !local {334		// The unprivileged agent doubles as a polkit authentication agent so335		// `run0` (e.g. our own elevation) routes its prompt to the User over336		// bifrost instead of failing on a tty-less session.337		let conn = Connection::system().await?;338		let helper = SocketHelper {339			fallback: SuidHelper,340		};341		register_auth_agent(&conn, Agent::new(helper, user_prompter)).await?;342		Some(conn)343	} else {344		None345	};346347	let _keep_alive = (bus, helpers, polkit_conn);348	pending().await349}350351async fn register_auth_agent<H, P>(conn: &Connection, agent: Agent<H, P>) -> anyhow::Result<()>352where353	H: Helper + Clone + Send + Sync + 'static,354	P: Prompter + Clone + Send + Sync + 'static,355{356	let proxy = zbus_polkit::policykit1::AuthorityProxy::new(conn).await?;357	conn.object_server().at(OBJ_PATH, agent).await?;358359	let subject = auth_agent_subject()?;360	proxy361		.register_authentication_agent(&subject, "C", OBJ_PATH)362		.await?;363	debug!(kind = subject.subject_kind, "registered polkit agent");364	Ok(())365}366367fn auth_agent_subject() -> anyhow::Result<Subject> {368	let mut details = HashMap::new();369	if let Ok(session_id) = std::env::var("XDG_SESSION_ID") {370		let val: OwnedValue = Str::from(session_id).into();371		details.insert("session-id".to_string(), val);372		return Ok(Subject {373			subject_kind: "unix-session".to_string(),374			subject_details: details,375		});376	}377378	details.insert("pid".to_string(), OwnedValue::from(std::process::id()));379	Ok(Subject {380		subject_kind: "unix-process".to_string(),381		subject_details: details,382	})383}384385fn sh_quote(s: &str) -> String {386	format!("'{}'", s.replace('\'', "'\\''"))387}388389/// Prepend `dir` to the process `PATH`.390///391/// # SAFETY392///393/// Same as `set_var`394unsafe fn prepend_path(dir: &std::path::Path) {395	let value = match std::env::var_os("PATH") {396		Some(existing) => {397			let mut v = dir.as_os_str().to_owned();398			v.push(":");399			v.push(existing);400			v401		}402		None => dir.as_os_str().to_owned(),403	};404	unsafe {405		std::env::set_var("PATH", value);406	}407}