git.delta.rocks / remowt / refs/commits / 11447649e1e7

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::auto::AutoPrompter;17use remowt_ui_prompt::bifrost::serve_prompts;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 {28		host: String,29		#[arg(long)]30		escalate: bool,31	},32	/// Connect to local host for testing the connectivity.33	Local {34		#[arg(long)]35		escalate: bool,36	},37}3839fn agents_dir() -> anyhow::Result<PathBuf> {40	std::env::var_os("REMOWT_AGENTS_DIR")41		.map(PathBuf::from)42		.or_else(|| option_env!("REMOWT_AGENTS_DIR").map(PathBuf::from))43		.ok_or_else(|| anyhow!("no remowt-agents bundle"))44}4546#[tokio::main(flavor = "current_thread")]47async fn main() -> anyhow::Result<()> {48	tracing_subscriber::fmt()49		.with_writer(std::io::stderr)50		.without_time()51		.init();52	let opts = Opts::parse();5354	let bundle = AgentBundle::from_dir(agents_dir()?)?;55	let (conn, escalate) = match &opts {56		Opts::Ssh { host, escalate } => (57			Remowt::connect(host, &bundle, "remowt-ssh".to_owned()).await?,58			*escalate,59		),60		Opts::Local { escalate } => (61			Remowt::connect_local(&bundle, "remowt-ssh".to_owned()).await?,62			*escalate,63		),64	};65	let mut rpc = conn.rpc();6667	serve_prompts(68		&mut rpc,69		PrependSourcePrompter {70			prompter: AutoPrompter::new().await,71			source: match opts {72				Opts::Ssh { host, .. } => vec![Source(Cow::Owned(format!("ssh host: {}", host)))],73				Opts::Local { .. } => vec![],74			},75			description: "".to_owned(),76		},77	);78	if let Some(sess) = conn.ssh() {79		serve_editor(&mut rpc, SshEditor { sess });80	}8182	debug!("entering shell");83	run_shell(&conn, escalate).await?;84	debug!("shell ended");8586	Ok(())87}8889async fn run_shell(conn: &Remowt, escalate: bool) -> anyhow::Result<()> {90	let term = match std::env::var("TERM") {91		Ok(v) => v,92		Err(VarError::NotPresent) => "xterm-256color".to_owned(),93		Err(e) => return Err(e.into()),94	};95	let (cols, rows) = term_size().unwrap_or((80, 24));9697	let shell = conn.open_shell(&term, cols, rows, escalate).await?;98	let resizer = shell.resizer();99	let stream = shell.stream;100101	let _raw = RawMode::enable();102103	if let Ok(mut winch) = signal(SignalKind::window_change()) {104		tokio::spawn(async move {105			while winch.recv().await.is_some() {106				if let Some((cols, rows)) = term_size() {107					let _ = resizer.resize(cols, rows).await;108				}109			}110		});111	}112113	let (mut from_remote, mut to_remote) = tokio::io::split(stream);114	let mut stdin = AsyncStdin::new()?;115	let mut stdout = tokio::io::stdout();116117	tokio::select! {118		r = tokio::io::copy(&mut from_remote, &mut stdout) => { r?; }119		_ = tokio::io::copy(&mut stdin, &mut to_remote) => {}120	}121122	Ok(())123}124125struct AsyncStdin {126	fd: AsyncFd<RawFd>,127	original_flags: i32,128}129130impl AsyncStdin {131	fn new() -> io::Result<Self> {132		let raw = libc::STDIN_FILENO;133		// SAFETY: F_GETFL/F_SETFL round-trip on a valid fd.134		let original_flags = unsafe { libc::fcntl(raw, libc::F_GETFL) };135		if original_flags < 0 {136			return Err(io::Error::last_os_error());137		}138		if unsafe { libc::fcntl(raw, libc::F_SETFL, original_flags | libc::O_NONBLOCK) } < 0 {139			return Err(io::Error::last_os_error());140		}141		Ok(Self {142			fd: AsyncFd::new(raw)?,143			original_flags,144		})145	}146}147148impl Drop for AsyncStdin {149	fn drop(&mut self) {150		// SAFETY: restoring the flags we saved on a valid fd.151		unsafe { libc::fcntl(libc::STDIN_FILENO, libc::F_SETFL, self.original_flags) };152	}153}154155impl AsyncRead for AsyncStdin {156	fn poll_read(157		self: Pin<&mut Self>,158		cx: &mut Context<'_>,159		buf: &mut ReadBuf<'_>,160	) -> Poll<io::Result<()>> {161		let this = self.get_mut();162		loop {163			let mut guard = match this.fd.poll_read_ready(cx) {164				Poll::Ready(Ok(g)) => g,165				Poll::Ready(Err(e)) => return Poll::Ready(Err(e)),166				Poll::Pending => return Poll::Pending,167			};168			let unfilled = buf.initialize_unfilled();169			let res = guard.try_io(|inner| {170				let fd = *inner.get_ref();171				// SAFETY: writing into `unfilled`'s own backing storage.172				let n = unsafe { libc::read(fd, unfilled.as_mut_ptr().cast(), unfilled.len()) };173				if n < 0 {174					Err(io::Error::last_os_error())175				} else {176					Ok(n as usize)177				}178			});179			match res {180				Ok(Ok(n)) => {181					buf.advance(n);182					return Poll::Ready(Ok(()));183				}184				Ok(Err(e)) => return Poll::Ready(Err(e)),185				Err(_would_block) => continue,186			}187		}188	}189}190191fn term_size() -> Option<(u16, u16)> {192	let mut ws: libc::winsize = unsafe { std::mem::zeroed() };193	let rc = unsafe { libc::ioctl(libc::STDIN_FILENO, libc::TIOCGWINSZ, &mut ws) };194	if rc != 0 || ws.ws_col == 0 {195		None196	} else {197		Some((ws.ws_col, ws.ws_row))198	}199}200201struct RawMode {202	original: Termios,203}204205impl RawMode {206	fn enable() -> Option<Self> {207		let stdin = std::io::stdin();208		// SAFETY: trivial libc call on a borrowed fd.209		if unsafe { libc::isatty(stdin.as_raw_fd()) } != 1 {210			return None;211		}212		let original = termios::tcgetattr(&stdin).ok()?;213		let mut raw = original.clone();214		termios::cfmakeraw(&mut raw);215		termios::tcsetattr(&stdin, SetArg::TCSANOW, &raw).ok()?;216		Some(Self { original })217	}218}219220impl Drop for RawMode {221	fn drop(&mut self) {222		let _ = termios::tcsetattr(std::io::stdin(), SetArg::TCSANOW, &self.original);223	}224}