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 tokio::io::unix::AsyncFd;17use tokio::io::{AsyncRead, ReadBuf};18use tokio::signal::unix::{signal, SignalKind};19use tracing::info;20use remowt_ui_prompt::bifrost::serve_prompts;21use remowt_ui_prompt::rofi::RofiPrompter;22use remowt_ui_prompt::{PrependSourcePrompter, Source};2324#[derive(Parser)]25struct Opts {26 host: String,27}2829fn agents_dir() -> anyhow::Result<PathBuf> {30 std::env::var_os("REMOWT_AGENTS_DIR")31 .map(PathBuf::from)32 .or_else(|| option_env!("REMOWT_AGENTS_DIR").map(PathBuf::from))33 .ok_or_else(|| anyhow!("no remowt-agents bundle"))34}3536#[tokio::main(flavor = "current_thread")]37async fn main() -> anyhow::Result<()> {38 tracing_subscriber::fmt::init();39 let opts = Opts::parse();4041 let bundle = AgentBundle::from_dir(agents_dir()?)?;42 let conn = Remowt::connect(&opts.host, &bundle).await?;43 let mut rpc = conn.rpc();4445 serve_prompts(46 &mut rpc,47 PrependSourcePrompter {48 prompter: RofiPrompter,49 source: vec![Source(Cow::Owned(format!("ssh host: {}", opts.host)))],50 description: "".to_owned(),51 },52 );53 if let Some(sess) = conn.ssh() {54 serve_editor(&mut rpc, SshEditor { sess });55 }5657 info!("entering shell");58 run_shell(&conn).await?;59 info!("shell ended");6061 Ok(())62}6364async fn run_shell(conn: &Remowt) -> anyhow::Result<()> {65 let term = match std::env::var("TERM") {66 Ok(v) => v,67 Err(VarError::NotPresent) => "xterm-256color".to_owned(),68 Err(e) => return Err(e.into()),69 };70 let (cols, rows) = term_size().unwrap_or((80, 24));7172 let shell = conn.open_shell(&term, cols, rows).await?;73 let resizer = shell.resizer();74 let stream = shell.stream;7576 let _raw = RawMode::enable();7778 if let Ok(mut winch) = signal(SignalKind::window_change()) {79 tokio::spawn(async move {80 while winch.recv().await.is_some() {81 if let Some((cols, rows)) = term_size() {82 let _ = resizer.resize(cols, rows).await;83 }84 }85 });86 }8788 let (mut from_remote, mut to_remote) = tokio::io::split(stream);89 let mut stdin = AsyncStdin::new()?;90 let mut stdout = tokio::io::stdout();9192 tokio::select! {93 r = tokio::io::copy(&mut from_remote, &mut stdout) => { r?; }94 _ = tokio::io::copy(&mut stdin, &mut to_remote) => {}95 }9697 Ok(())98}99100struct AsyncStdin {101 fd: AsyncFd<RawFd>,102 original_flags: i32,103}104105impl AsyncStdin {106 fn new() -> io::Result<Self> {107 let raw = libc::STDIN_FILENO;108 109 let original_flags = unsafe { libc::fcntl(raw, libc::F_GETFL) };110 if original_flags < 0 {111 return Err(io::Error::last_os_error());112 }113 if unsafe { libc::fcntl(raw, libc::F_SETFL, original_flags | libc::O_NONBLOCK) } < 0 {114 return Err(io::Error::last_os_error());115 }116 Ok(Self {117 fd: AsyncFd::new(raw)?,118 original_flags,119 })120 }121}122123impl Drop for AsyncStdin {124 fn drop(&mut self) {125 126 unsafe { libc::fcntl(libc::STDIN_FILENO, libc::F_SETFL, self.original_flags) };127 }128}129130impl AsyncRead for AsyncStdin {131 fn poll_read(132 self: Pin<&mut Self>,133 cx: &mut Context<'_>,134 buf: &mut ReadBuf<'_>,135 ) -> Poll<io::Result<()>> {136 let this = self.get_mut();137 loop {138 let mut guard = match this.fd.poll_read_ready(cx) {139 Poll::Ready(Ok(g)) => g,140 Poll::Ready(Err(e)) => return Poll::Ready(Err(e)),141 Poll::Pending => return Poll::Pending,142 };143 let unfilled = buf.initialize_unfilled();144 let res = guard.try_io(|inner| {145 let fd = *inner.get_ref();146 147 let n = unsafe { libc::read(fd, unfilled.as_mut_ptr().cast(), unfilled.len()) };148 if n < 0 {149 Err(io::Error::last_os_error())150 } else {151 Ok(n as usize)152 }153 });154 match res {155 Ok(Ok(n)) => {156 buf.advance(n);157 return Poll::Ready(Ok(()));158 }159 Ok(Err(e)) => return Poll::Ready(Err(e)),160 Err(_would_block) => continue,161 }162 }163 }164}165166fn term_size() -> Option<(u16, u16)> {167 let mut ws: libc::winsize = unsafe { std::mem::zeroed() };168 let rc = unsafe { libc::ioctl(libc::STDIN_FILENO, libc::TIOCGWINSZ, &mut ws) };169 if rc != 0 || ws.ws_col == 0 {170 None171 } else {172 Some((ws.ws_col, ws.ws_row))173 }174}175176struct RawMode {177 original: Termios,178}179180impl RawMode {181 fn enable() -> Option<Self> {182 let stdin = std::io::stdin();183 184 if unsafe { libc::isatty(stdin.as_raw_fd()) } != 1 {185 return None;186 }187 let original = termios::tcgetattr(&stdin).ok()?;188 let mut raw = original.clone();189 termios::cfmakeraw(&mut raw);190 termios::tcsetattr(&stdin, SetArg::TCSANOW, &raw).ok()?;191 Some(Self { original })192 }193}194195impl Drop for RawMode {196 fn drop(&mut self) {197 let _ = termios::tcsetattr(std::io::stdin(), SetArg::TCSANOW, &self.original);198 }199}