git.delta.rocks / remowt / refs/commits / 7e5126608cef

1use std::borrow::Cow;2use std::env::VarError;3use std::io;4use std::os::fd::{AsRawFd, RawFd};5use std::path::PathBuf;6use std::pin::Pin;7use std::task::{Context, Poll};89use anyhow::anyhow;10use clap::Parser;11use nix::libc;12use nix::sys::termios::{self, SetArg, Termios};13use remowt_client::editor::SshEditor;14use remowt_client::{AgentBundle, Remowt};15use remowt_link_shared::editor::serve_editor;16use remowt_ui_prompt::bifrost::serve_prompts;17use remowt_ui_prompt::rofi::RofiPrompter;18use remowt_ui_prompt::{PrependSourcePrompter, Source};19use tokio::io::unix::AsyncFd;20use tokio::io::{AsyncRead, ReadBuf};21use tokio::signal::unix::{signal, SignalKind};22use tracing::debug;2324#[derive(Parser)]25enum Opts {26	/// Connect to remote host with remowt agent.27	Ssh { host: String },28	/// Connect to local host for testing the connectivity.29	Local,30}3132fn agents_dir() -> anyhow::Result<PathBuf> {33	std::env::var_os("REMOWT_AGENTS_DIR")34		.map(PathBuf::from)35		.or_else(|| option_env!("REMOWT_AGENTS_DIR").map(PathBuf::from))36		.ok_or_else(|| anyhow!("no remowt-agents bundle"))37}3839#[tokio::main(flavor = "current_thread")]40async fn main() -> anyhow::Result<()> {41	tracing_subscriber::fmt()42		.with_writer(std::io::stderr)43		.without_time()44		.init();45	let opts = Opts::parse();4647	let bundle = AgentBundle::from_dir(agents_dir()?)?;48	let conn = match &opts {49		Opts::Ssh { host } => Remowt::connect(host, &bundle).await?,50		Opts::Local => Remowt::connect_local(&bundle).await?,51	};52	let mut rpc = conn.rpc();5354	serve_prompts(55		&mut rpc,56		PrependSourcePrompter {57			prompter: RofiPrompter,58			source: match opts {59				Opts::Ssh { host } => vec![Source(Cow::Owned(format!("ssh host: {}", host)))],60				Opts::Local => vec![],61			},62			description: "".to_owned(),63		},64	);65	if let Some(sess) = conn.ssh() {66		serve_editor(&mut rpc, SshEditor { sess });67	}6869	debug!("entering shell");70	run_shell(&conn).await?;71	debug!("shell ended");7273	Ok(())74}7576async fn run_shell(conn: &Remowt) -> anyhow::Result<()> {77	let term = match std::env::var("TERM") {78		Ok(v) => v,79		Err(VarError::NotPresent) => "xterm-256color".to_owned(),80		Err(e) => return Err(e.into()),81	};82	let (cols, rows) = term_size().unwrap_or((80, 24));8384	let shell = conn.open_shell(&term, cols, rows).await?;85	let resizer = shell.resizer();86	let stream = shell.stream;8788	let _raw = RawMode::enable();8990	if let Ok(mut winch) = signal(SignalKind::window_change()) {91		tokio::spawn(async move {92			while winch.recv().await.is_some() {93				if let Some((cols, rows)) = term_size() {94					let _ = resizer.resize(cols, rows).await;95				}96			}97		});98	}99100	let (mut from_remote, mut to_remote) = tokio::io::split(stream);101	let mut stdin = AsyncStdin::new()?;102	let mut stdout = tokio::io::stdout();103104	tokio::select! {105		r = tokio::io::copy(&mut from_remote, &mut stdout) => { r?; }106		_ = tokio::io::copy(&mut stdin, &mut to_remote) => {}107	}108109	Ok(())110}111112struct AsyncStdin {113	fd: AsyncFd<RawFd>,114	original_flags: i32,115}116117impl AsyncStdin {118	fn new() -> io::Result<Self> {119		let raw = libc::STDIN_FILENO;120		// SAFETY: F_GETFL/F_SETFL round-trip on a valid fd.121		let original_flags = unsafe { libc::fcntl(raw, libc::F_GETFL) };122		if original_flags < 0 {123			return Err(io::Error::last_os_error());124		}125		if unsafe { libc::fcntl(raw, libc::F_SETFL, original_flags | libc::O_NONBLOCK) } < 0 {126			return Err(io::Error::last_os_error());127		}128		Ok(Self {129			fd: AsyncFd::new(raw)?,130			original_flags,131		})132	}133}134135impl Drop for AsyncStdin {136	fn drop(&mut self) {137		// SAFETY: restoring the flags we saved on a valid fd.138		unsafe { libc::fcntl(libc::STDIN_FILENO, libc::F_SETFL, self.original_flags) };139	}140}141142impl AsyncRead for AsyncStdin {143	fn poll_read(144		self: Pin<&mut Self>,145		cx: &mut Context<'_>,146		buf: &mut ReadBuf<'_>,147	) -> Poll<io::Result<()>> {148		let this = self.get_mut();149		loop {150			let mut guard = match this.fd.poll_read_ready(cx) {151				Poll::Ready(Ok(g)) => g,152				Poll::Ready(Err(e)) => return Poll::Ready(Err(e)),153				Poll::Pending => return Poll::Pending,154			};155			let unfilled = buf.initialize_unfilled();156			let res = guard.try_io(|inner| {157				let fd = *inner.get_ref();158				// SAFETY: writing into `unfilled`'s own backing storage.159				let n = unsafe { libc::read(fd, unfilled.as_mut_ptr().cast(), unfilled.len()) };160				if n < 0 {161					Err(io::Error::last_os_error())162				} else {163					Ok(n as usize)164				}165			});166			match res {167				Ok(Ok(n)) => {168					buf.advance(n);169					return Poll::Ready(Ok(()));170				}171				Ok(Err(e)) => return Poll::Ready(Err(e)),172				Err(_would_block) => continue,173			}174		}175	}176}177178fn term_size() -> Option<(u16, u16)> {179	let mut ws: libc::winsize = unsafe { std::mem::zeroed() };180	let rc = unsafe { libc::ioctl(libc::STDIN_FILENO, libc::TIOCGWINSZ, &mut ws) };181	if rc != 0 || ws.ws_col == 0 {182		None183	} else {184		Some((ws.ws_col, ws.ws_row))185	}186}187188struct RawMode {189	original: Termios,190}191192impl RawMode {193	fn enable() -> Option<Self> {194		let stdin = std::io::stdin();195		// SAFETY: trivial libc call on a borrowed fd.196		if unsafe { libc::isatty(stdin.as_raw_fd()) } != 1 {197			return None;198		}199		let original = termios::tcgetattr(&stdin).ok()?;200		let mut raw = original.clone();201		termios::cfmakeraw(&mut raw);202		termios::tcsetattr(&stdin, SetArg::TCSANOW, &raw).ok()?;203		Some(Self { original })204	}205}206207impl Drop for RawMode {208	fn drop(&mut self) {209		let _ = termios::tcsetattr(std::io::stdin(), SetArg::TCSANOW, &self.original);210	}211}