creddy/src-tauri/src/server/mod.rs

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<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;
        }
        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<Response, HandlerError> {
    shortcuts::exec_shortcut(action);
    Ok(Response::Empty)
}


async fn get_aws_credentials(
    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, 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_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
}