git.delta.rocks / fleet / refs/commits / 0ea4b47ffa04

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(80			&mut rpc,81			SshEditor {82				sess,83				conn: conn.clone(),84			},85		);86	}8788	debug!("entering shell");89	run_shell(&conn, escalate).await?;90	debug!("shell ended");9192	Ok(())93}9495async fn run_shell(conn: &Remowt, escalate: bool) -> anyhow::Result<()> {96	let term = match std::env::var("TERM") {97		Ok(v) => v,98		Err(VarError::NotPresent) => "xterm-256color".to_owned(),99		Err(e) => return Err(e.into()),100	};101	let (cols, rows) = term_size().unwrap_or((80, 24));102103	let shell = conn.open_shell(&term, cols, rows, escalate).await?;104	let resizer = shell.resizer();105	let stream = shell.stream;106107	let _raw = RawMode::enable();108109	if let Ok(mut winch) = signal(SignalKind::window_change()) {110		tokio::spawn(async move {111			while winch.recv().await.is_some() {112				if let Some((cols, rows)) = term_size() {113					let _ = resizer.resize(cols, rows).await;114				}115			}116		});117	}118119	let (mut from_remote, mut to_remote) = tokio::io::split(stream);120	let mut stdin = AsyncStdin::new()?;121	let mut stdout = tokio::io::stdout();122123	tokio::select! {124		r = tokio::io::copy(&mut from_remote, &mut stdout) => { r?; }125		_ = tokio::io::copy(&mut stdin, &mut to_remote) => {}126	}127128	Ok(())129}130131struct AsyncStdin {132	fd: AsyncFd<RawFd>,133	original_flags: i32,134}135136impl AsyncStdin {137	fn new() -> io::Result<Self> {138		let raw = libc::STDIN_FILENO;139		// SAFETY: F_GETFL/F_SETFL round-trip on a valid fd.140		let original_flags = unsafe { libc::fcntl(raw, libc::F_GETFL) };141		if original_flags < 0 {142			return Err(io::Error::last_os_error());143		}144		if unsafe { libc::fcntl(raw, libc::F_SETFL, original_flags | libc::O_NONBLOCK) } < 0 {145			return Err(io::Error::last_os_error());146		}147		Ok(Self {148			fd: AsyncFd::new(raw)?,149			original_flags,150		})151	}152}153154impl Drop for AsyncStdin {155	fn drop(&mut self) {156		// SAFETY: restoring the flags we saved on a valid fd.157		unsafe { libc::fcntl(libc::STDIN_FILENO, libc::F_SETFL, self.original_flags) };158	}159}160161impl AsyncRead for AsyncStdin {162	fn poll_read(163		self: Pin<&mut Self>,164		cx: &mut Context<'_>,165		buf: &mut ReadBuf<'_>,166	) -> Poll<io::Result<()>> {167		let this = self.get_mut();168		loop {169			let mut guard = match this.fd.poll_read_ready(cx) {170				Poll::Ready(Ok(g)) => g,171				Poll::Ready(Err(e)) => return Poll::Ready(Err(e)),172				Poll::Pending => return Poll::Pending,173			};174			let unfilled = buf.initialize_unfilled();175			let res = guard.try_io(|inner| {176				let fd = *inner.get_ref();177				// SAFETY: writing into `unfilled`'s own backing storage.178				let n = unsafe { libc::read(fd, unfilled.as_mut_ptr().cast(), unfilled.len()) };179				if n < 0 {180					Err(io::Error::last_os_error())181				} else {182					Ok(n as usize)183				}184			});185			match res {186				Ok(Ok(n)) => {187					buf.advance(n);188					return Poll::Ready(Ok(()));189				}190				Ok(Err(e)) => return Poll::Ready(Err(e)),191				Err(_would_block) => continue,192			}193		}194	}195}196197fn term_size() -> Option<(u16, u16)> {198	let mut ws: libc::winsize = unsafe { std::mem::zeroed() };199	let rc = unsafe { libc::ioctl(libc::STDIN_FILENO, libc::TIOCGWINSZ, &mut ws) };200	if rc != 0 || ws.ws_col == 0 {201		None202	} else {203		Some((ws.ws_col, ws.ws_row))204	}205}206207struct RawMode {208	original: Termios,209}210211impl RawMode {212	fn enable() -> Option<Self> {213		let stdin = std::io::stdin();214		// SAFETY: trivial libc call on a borrowed fd.215		if unsafe { libc::isatty(stdin.as_raw_fd()) } != 1 {216			return None;217		}218		let original = termios::tcgetattr(&stdin).ok()?;219		let mut raw = original.clone();220		termios::cfmakeraw(&mut raw);221		termios::tcsetattr(&stdin, SetArg::TCSANOW, &raw).ok()?;222		Some(Self { original })223	}224}225226impl Drop for RawMode {227	fn drop(&mut self) {228		let _ = termios::tcsetattr(std::io::stdin(), SetArg::TCSANOW, &self.original);229	}230}