use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tokio::sync::oneshot; use serde::{Serialize, Deserialize}; use tauri::{AppHandle, Manager}; use crate::errors::*; use crate::clientinfo::{self, Client}; use crate::credentials::{ AwsBaseCredential, AwsSessionCredential, }; use crate::ipc::{Approval, RequestNotification}; use crate::state::AppState; use crate::shortcuts::{self, ShortcutAction}; #[cfg(windows)] mod server_win; #[cfg(windows)] pub use server_win::Server; #[cfg(windows)] use server_win::Stream; #[cfg(unix)] mod server_unix; #[cfg(unix)] pub use server_unix::Server; #[cfg(unix)] use server_unix::Stream; pub mod ssh_agent; pub use ssh_agent::Agent; #[derive(Serialize, Deserialize)] pub enum Request { GetAwsCredentials{ base: bool, }, InvokeShortcut(ShortcutAction), } #[derive(Debug, Serialize, Deserialize)] pub enum Response { AwsBase(AwsBaseCredential), AwsSession(AwsSessionCredential), Empty, } struct CloseWaiter<'s> { stream: &'s mut Stream, } impl<'s> CloseWaiter<'s> { async fn wait_for_close(&mut self) -> std::io::Result<()> { let mut buf = [0u8; 8]; loop { match self.stream.read(&mut buf).await { Ok(0) => break Ok(()), Ok(_) => (), Err(e) => break Err(e), } } } } async fn handle(mut stream: Stream, app_handle: AppHandle, client_pid: u32) -> Result<(), HandlerError> { // read from stream until delimiter is reached let mut buf: Vec = Vec::with_capacity(1024); // requests are small, 1KiB is more than enough let mut n = 0; loop { n += stream.read_buf(&mut buf).await?; if let Some(&b'\n') = buf.last() { break; } else if n >= 1024 { return Err(HandlerError::RequestTooLarge); } } let client = clientinfo::get_client(client_pid, true)?; let waiter = CloseWaiter { stream: &mut stream }; let req: Request = serde_json::from_slice(&buf)?; let res = match req { Request::GetAwsCredentials{ base } => get_aws_credentials( base, client, app_handle, waiter ).await, Request::InvokeShortcut(action) => invoke_shortcut(action).await, }; // doesn't make sense to send the error to the client if the client has already left if let Err(HandlerError::Abandoned) = res { return Err(HandlerError::Abandoned); } let res = serde_json::to_vec(&res).unwrap(); stream.write_all(&res).await?; Ok(()) } async fn invoke_shortcut(action: ShortcutAction) -> Result { shortcuts::exec_shortcut(action); Ok(Response::Empty) } async fn get_aws_credentials( base: bool, client: Client, app_handle: AppHandle, mut waiter: CloseWaiter<'_>, ) -> Result { let state = app_handle.state::(); let rehide_ms = { let config = state.config.read().await; config.rehide_ms }; let lease = state.acquire_visibility_lease(rehide_ms).await .map_err(|_e| HandlerError::NoMainWindow)?; // automate this conversion eventually? let (chan_send, chan_recv) = oneshot::channel(); let request_id = state.register_request(chan_send).await; // if an error occurs in any of the following, we want to abort the operation // but ? returns immediately, and we want to unregister the request before returning // so we bundle it all up in an async block and return a Result so we can handle errors let proceed = async { let notification = RequestNotification::new_aws(request_id, client, base); app_handle.emit("credential-request", ¬ification)?; let response = tokio::select! { r = chan_recv => r?, _ = waiter.wait_for_close() => { app_handle.emit("request-cancelled", request_id)?; return Err(HandlerError::Abandoned); }, }; match response.approval { Approval::Approved => { if response.base { let creds = state.get_aws_default().await?; Ok(Response::AwsBase(creds)) } else { let creds = state.get_aws_default_session().await?; Ok(Response::AwsSession(creds.clone())) } }, Approval::Denied => Err(HandlerError::Denied), } }; let result = match proceed.await { Ok(r) => Ok(r), Err(e) => { state.unregister_request(request_id).await; Err(e) } }; lease.release(); result }