git.delta.rocks / fleet / refs/commits / b66112ef1437

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;23use tracing_subscriber::prelude::*;24use tracing_subscriber::EnvFilter;2526#[derive(Parser)]27enum Opts {28	/// Connect to remote host with remowt agent.29	Ssh {30		host: String,31		#[arg(long)]32		escalate: bool,33	},34	/// Connect to local host for testing the connectivity.35	Local {36		#[arg(long)]37		escalate: bool,38	},39}4041fn agents_dir() -> anyhow::Result<PathBuf> {42	std::env::var_os("REMOWT_AGENTS_DIR")43		.map(PathBuf::from)44		.or_else(|| option_env!("REMOWT_AGENTS_DIR").map(PathBuf::from))45		.ok_or_else(|| anyhow!("no remowt-agents bundle"))46}4748#[tokio::main(flavor = "current_thread")]49async fn main() -> anyhow::Result<()> {50	// Log to the journal, not stderr: this is an interactive client that puts51	// the terminal in raw mode, so anything on stderr corrupts the session52	// output. If the journal socket is unavailable, drop logs rather than fall53	// back to stderr.54	if let Ok(journald) = tracing_journald::layer() {55		tracing_subscriber::registry()56			.with(EnvFilter::from_default_env())57			.with(journald.with_syslog_identifier("remowt-ssh".to_owned()))58			.init();59	}60	let opts = Opts::parse();6162	let bundle = AgentBundle::from_dir(agents_dir()?)?;63	let (conn, escalate) = match &opts {64		Opts::Ssh { host, escalate } => (65			Remowt::connect(host, &bundle, "remowt-ssh".to_owned()).await?,66			*escalate,67		),68		Opts::Local { escalate } => (69			Remowt::connect_local(&bundle, "remowt-ssh".to_owned()).await?,70			*escalate,71		),72	};73	let mut rpc = conn.rpc();7475	serve_prompts(76		&mut rpc,77		PrependSourcePrompter {78			prompter: AutoPrompter::new().await,79			source: match opts {80				Opts::Ssh { host, .. } => vec![Source(Cow::Owned(format!("ssh host: {}", host)))],81				Opts::Local { .. } => vec![],82			},83			description: "".to_owned(),84		},85	);86	if conn.ssh().is_some() {87		serve_editor(&mut rpc, SshEditor { conn: conn.clone() });88	}8990	debug!("entering shell");91	run_shell(&conn, escalate).await?;92	debug!("shell ended");9394	Ok(())95}9697async fn run_shell(conn: &Remowt, escalate: bool) -> anyhow::Result<()> {98	let term = match std::env::var("TERM") {99		Ok(v) => v,100		Err(VarError::NotPresent) => "xterm-256color".to_owned(),101		Err(e) => return Err(e.into()),102	};103	let (cols, rows) = term_size().unwrap_or((80, 24));104105	let shell = conn.open_shell(&term, cols, rows, escalate).await?;106	let resizer = shell.resizer();107	let stream = shell.stream;108109	let _raw = RawMode::enable();110111	if let Ok(mut winch) = signal(SignalKind::window_change()) {112		tokio::spawn(async move {113			while winch.recv().await.is_some() {114				if let Some((cols, rows)) = term_size() {115					let _ = resizer.resize(cols, rows).await;116				}117			}118		});119	}120121	let (mut from_remote, mut to_remote) = tokio::io::split(stream);122	let mut stdin = AsyncStdin::new()?;123	let mut stdout = tokio::io::stdout();124125	tokio::select! {126		r = tokio::io::copy(&mut from_remote, &mut stdout) => { r?; }127		_ = tokio::io::copy(&mut stdin, &mut to_remote) => {}128	}129130	Ok(())131}132133struct AsyncStdin {134	fd: AsyncFd<RawFd>,135	original_flags: i32,136}137138impl AsyncStdin {139	fn new() -> io::Result<Self> {140		let raw = libc::STDIN_FILENO;141		// SAFETY: F_GETFL/F_SETFL round-trip on a valid fd.142		let original_flags = unsafe { libc::fcntl(raw, libc::F_GETFL) };143		if original_flags < 0 {144			return Err(io::Error::last_os_error());145		}146		if unsafe { libc::fcntl(raw, libc::F_SETFL, original_flags | libc::O_NONBLOCK) } < 0 {147			return Err(io::Error::last_os_error());148		}149		Ok(Self {150			fd: AsyncFd::new(raw)?,151			original_flags,152		})153	}154}155156impl Drop for AsyncStdin {157	fn drop(&mut self) {158		// SAFETY: restoring the flags we saved on a valid fd.159		unsafe { libc::fcntl(libc::STDIN_FILENO, libc::F_SETFL, self.original_flags) };160	}161}162163impl AsyncRead for AsyncStdin {164	fn poll_read(165		self: Pin<&mut Self>,166		cx: &mut Context<'_>,167		buf: &mut ReadBuf<'_>,168	) -> Poll<io::Result<()>> {169		let this = self.get_mut();170		loop {171			let mut guard = match this.fd.poll_read_ready(cx) {172				Poll::Ready(Ok(g)) => g,173				Poll::Ready(Err(e)) => return Poll::Ready(Err(e)),174				Poll::Pending => return Poll::Pending,175			};176			let unfilled = buf.initialize_unfilled();177			let res = guard.try_io(|inner| {178				let fd = *inner.get_ref();179				// SAFETY: writing into `unfilled`'s own backing storage.180				let n = unsafe { libc::read(fd, unfilled.as_mut_ptr().cast(), unfilled.len()) };181				if n < 0 {182					Err(io::Error::last_os_error())183				} else {184					Ok(n as usize)185				}186			});187			match res {188				Ok(Ok(n)) => {189					buf.advance(n);190					return Poll::Ready(Ok(()));191				}192				Ok(Err(e)) => return Poll::Ready(Err(e)),193				Err(_would_block) => continue,194			}195		}196	}197}198199fn term_size() -> Option<(u16, u16)> {200	let mut ws: libc::winsize = unsafe { std::mem::zeroed() };201	let rc = unsafe { libc::ioctl(libc::STDIN_FILENO, libc::TIOCGWINSZ, &mut ws) };202	if rc != 0 || ws.ws_col == 0 {203		None204	} else {205		Some((ws.ws_col, ws.ws_row))206	}207}208209struct RawMode {210	original: Termios,211}212213impl RawMode {214	fn enable() -> Option<Self> {215		let stdin = std::io::stdin();216		// SAFETY: trivial libc call on a borrowed fd.217		if unsafe { libc::isatty(stdin.as_raw_fd()) } != 1 {218			return None;219		}220		let original = termios::tcgetattr(&stdin).ok()?;221		let mut raw = original.clone();222		termios::cfmakeraw(&mut raw);223		termios::tcsetattr(&stdin, SetArg::TCSANOW, &raw).ok()?;224		Some(Self { original })225	}226}227228impl Drop for RawMode {229	fn drop(&mut self) {230		let _ = termios::tcsetattr(std::io::stdin(), SetArg::TCSANOW, &self.original);231	}232}