finish SSH key support

src-tauri/src/srv/creddy_server.rs

@@ -0,0 +1,132 @@
+use tauri::{AppHandle, Manager};
+use tokio::io::{AsyncReadExt, AsyncWriteExt};
+use tokio::sync::oneshot;
+
+use crate::clientinfo::{self, Client};
+use crate::errors::*;
+use crate::ipc::{Approval, RequestNotification};
+use crate::shortcuts::{self, ShortcutAction};
+use crate::state::AppState;
+use super::{
+    CloseWaiter,
+    Request,
+    Response,
+    Stream,
+};
+
+
+pub fn serve(app_handle: AppHandle) -> std::io::Result<()> {
+    super::serve("creddy-server", app_handle, handle)
+}
+
+
+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<u8> = 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;
+        }
+        // sanity check, no request should ever be within a mile of 1MB
+        else if n >= (1024 * 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 { name, base } => get_aws_credentials(
+            name, base, client, app_handle, waiter
+        ).await,
+        Request::InvokeShortcut(action) => invoke_shortcut(action).await,
+        Request::GetSshSignature(_) => return Err(HandlerError::Denied),
+    };
+
+    // 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<Response, HandlerError> {
+    shortcuts::exec_shortcut(action);
+    Ok(Response::Empty)
+}
+
+
+async fn get_aws_credentials(
+    name: Option<String>,
+    base: bool,
+    client: Client,
+    app_handle: AppHandle,
+    mut waiter: CloseWaiter<'_>,
+) -> Result<Response, HandlerError> {
+    let state = app_handle.state::<AppState>();
+    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, name.clone(), base
+        );
+        app_handle.emit("credential-request", &notification)?;
+
+        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_base(name).await?;
+                    Ok(Response::AwsBase(creds))
+                }
+                else {
+                    let creds = state.get_aws_session(name).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
+}