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 27 Ssh {28 host: String,29 #[arg(long)]30 escalate: bool,31 },32 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 } => (Remowt::connect(host, &bundle).await?, *escalate),57 Opts::Local { escalate } => (Remowt::connect_local(&bundle).await?, *escalate),58 };59 let mut rpc = conn.rpc();6061 serve_prompts(62 &mut rpc,63 PrependSourcePrompter {64 prompter: AutoPrompter::new().await,65 source: match opts {66 Opts::Ssh { host, .. } => vec![Source(Cow::Owned(format!("ssh host: {}", host)))],67 Opts::Local { .. } => vec![],68 },69 description: "".to_owned(),70 },71 );72 if let Some(sess) = conn.ssh() {73 serve_editor(&mut rpc, SshEditor { sess });74 }7576 debug!("entering shell");77 run_shell(&conn, escalate).await?;78 debug!("shell ended");7980 Ok(())81}8283async fn run_shell(conn: &Remowt, escalate: bool) -> anyhow::Result<()> {84 let term = match std::env::var("TERM") {85 Ok(v) => v,86 Err(VarError::NotPresent) => "xterm-256color".to_owned(),87 Err(e) => return Err(e.into()),88 };89 let (cols, rows) = term_size().unwrap_or((80, 24));9091 let shell = conn.open_shell(&term, cols, rows, escalate).await?;92 let resizer = shell.resizer();93 let stream = shell.stream;9495 let _raw = RawMode::enable();9697 if let Ok(mut winch) = signal(SignalKind::window_change()) {98 tokio::spawn(async move {99 while winch.recv().await.is_some() {100 if let Some((cols, rows)) = term_size() {101 let _ = resizer.resize(cols, rows).await;102 }103 }104 });105 }106107 let (mut from_remote, mut to_remote) = tokio::io::split(stream);108 let mut stdin = AsyncStdin::new()?;109 let mut stdout = tokio::io::stdout();110111 tokio::select! {112 r = tokio::io::copy(&mut from_remote, &mut stdout) => { r?; }113 _ = tokio::io::copy(&mut stdin, &mut to_remote) => {}114 }115116 Ok(())117}118119struct AsyncStdin {120 fd: AsyncFd<RawFd>,121 original_flags: i32,122}123124impl AsyncStdin {125 fn new() -> io::Result<Self> {126 let raw = libc::STDIN_FILENO;127 128 let original_flags = unsafe { libc::fcntl(raw, libc::F_GETFL) };129 if original_flags < 0 {130 return Err(io::Error::last_os_error());131 }132 if unsafe { libc::fcntl(raw, libc::F_SETFL, original_flags | libc::O_NONBLOCK) } < 0 {133 return Err(io::Error::last_os_error());134 }135 Ok(Self {136 fd: AsyncFd::new(raw)?,137 original_flags,138 })139 }140}141142impl Drop for AsyncStdin {143 fn drop(&mut self) {144 145 unsafe { libc::fcntl(libc::STDIN_FILENO, libc::F_SETFL, self.original_flags) };146 }147}148149impl AsyncRead for AsyncStdin {150 fn poll_read(151 self: Pin<&mut Self>,152 cx: &mut Context<'_>,153 buf: &mut ReadBuf<'_>,154 ) -> Poll<io::Result<()>> {155 let this = self.get_mut();156 loop {157 let mut guard = match this.fd.poll_read_ready(cx) {158 Poll::Ready(Ok(g)) => g,159 Poll::Ready(Err(e)) => return Poll::Ready(Err(e)),160 Poll::Pending => return Poll::Pending,161 };162 let unfilled = buf.initialize_unfilled();163 let res = guard.try_io(|inner| {164 let fd = *inner.get_ref();165 166 let n = unsafe { libc::read(fd, unfilled.as_mut_ptr().cast(), unfilled.len()) };167 if n < 0 {168 Err(io::Error::last_os_error())169 } else {170 Ok(n as usize)171 }172 });173 match res {174 Ok(Ok(n)) => {175 buf.advance(n);176 return Poll::Ready(Ok(()));177 }178 Ok(Err(e)) => return Poll::Ready(Err(e)),179 Err(_would_block) => continue,180 }181 }182 }183}184185fn term_size() -> Option<(u16, u16)> {186 let mut ws: libc::winsize = unsafe { std::mem::zeroed() };187 let rc = unsafe { libc::ioctl(libc::STDIN_FILENO, libc::TIOCGWINSZ, &mut ws) };188 if rc != 0 || ws.ws_col == 0 {189 None190 } else {191 Some((ws.ws_col, ws.ws_row))192 }193}194195struct RawMode {196 original: Termios,197}198199impl RawMode {200 fn enable() -> Option<Self> {201 let stdin = std::io::stdin();202 203 if unsafe { libc::isatty(stdin.as_raw_fd()) } != 1 {204 return None;205 }206 let original = termios::tcgetattr(&stdin).ok()?;207 let mut raw = original.clone();208 termios::cfmakeraw(&mut raw);209 termios::tcsetattr(&stdin, SetArg::TCSANOW, &raw).ok()?;210 Some(Self { original })211 }212}213214impl Drop for RawMode {215 fn drop(&mut self) {216 let _ = termios::tcsetattr(std::io::stdin(), SetArg::TCSANOW, &self.original);217 }218}