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

use std::future::Future;

use tauri::{
    AppHandle,
    async_runtime as rt,
    Manager,
};
use tokio::io::AsyncReadExt;
use tokio::sync::oneshot;
use serde::{Serialize, Deserialize};

use crate::clientinfo::Client;
use crate::credentials::{
    AwsBaseCredential,
    AwsSessionCredential,
    Credential,
    DockerCredential,
};
use crate::errors::*;
use crate::ipc::{RequestNotification, RequestNotificationDetail, RequestResponse};
use crate::shortcuts::ShortcutAction;
use crate::state::AppState;

pub mod creddy_server;
pub mod agent;
use platform::Stream;


// These types match what's defined in creddy_cli, but they are separate types
// so that we avoid polluting the standalone CLI with a bunch of dependencies
// that would make it impossible to build a completely static-linked version
#[derive(Debug, Serialize, Deserialize)]
pub enum CliRequest {
    GetAwsCredential {
        name: Option<String>,
        base: bool,
    },
    GetDockerCredential {
        server_url: String,
    },
    SaveCredential {
        name: String,
        is_default: bool,
        credential: Credential,
    },
    InvokeShortcut(ShortcutAction),
}


#[derive(Debug, Serialize, Deserialize)]
pub enum CliResponse {
    Credential(CliCredential),
    Empty,
}


#[derive(Debug, Serialize, Deserialize)]
pub enum CliCredential {
    AwsBase(AwsBaseCredential),
    AwsSession(AwsSessionCredential),
    Docker(DockerCredential),
}


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),
            }
        }
    }
}


fn serve<H, F>(sock_name: &str, app_handle: AppHandle, handler: H) -> std::io::Result<()>
    where H: Copy + Send + Fn(Stream, AppHandle, u32) -> F + 'static,
          F: Send + Future<Output = Result<(), HandlerError>>,
{
    let (mut listener, addr) = platform::bind(sock_name)?;
    rt::spawn(async move {
        loop {
            let (stream, client_pid) = match platform::accept(&mut listener, &addr).await {
                Ok((s, c)) => (s, c),
                Err(e) => {
                    eprintln!("Error accepting request: {e}");
                    continue;
                },
            };
            let new_handle = app_handle.clone();
            rt::spawn(async move {
                handler(stream, new_handle, client_pid)
                    .await
                    .error_print_prefix("Error responding to request: ");
            });
        }
    });
    Ok(())
}


async fn send_credentials_request(
    detail: RequestNotificationDetail,
    app_handle: AppHandle,
    mut waiter: CloseWaiter<'_>
) -> Result<RequestResponse, 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)?;

    let (chan_send, chan_recv) = oneshot::channel();
    let request_id = state.register_request(chan_send).await;
    let notification = RequestNotification { id: request_id, detail };

    // the following could fail in various ways, but we want to make sure
    // the request gets unregistered on any failure, so we wrap this all
    // up in an async block so that we only have to handle the error case once
    let proceed = async {
        app_handle.emit("credential-request", &notification)?;
        tokio::select! {
            r = chan_recv => Ok(r?),
            _ = waiter.wait_for_close() => {
                app_handle.emit("request-cancelled", request_id)?;
                Err(HandlerError::Abandoned)
            },
        }
    };

    let res = proceed.await;
    if let Err(_) = &res {
        state.unregister_request(request_id).await;
    }

    lease.release();
    res
}


#[cfg(unix)]
mod platform {
    use std::io::ErrorKind;
    use std::path::PathBuf;
    use tokio::net::{UnixListener, UnixStream};
    use super::*;


    pub type Stream = UnixStream;

    pub fn bind(sock_name: &str) -> std::io::Result<(UnixListener, PathBuf)> {
        let path = creddy_cli::server_addr(sock_name);
        match std::fs::remove_file(&path) {
            Ok(_) => (),
            Err(e) if e.kind() == ErrorKind::NotFound => (),
            Err(e) => return Err(e),
        }

        let listener = UnixListener::bind(&path)?;
        Ok((listener, path))
    }

    pub async fn accept(listener: &mut UnixListener, _addr: &PathBuf) -> Result<(UnixStream, u32), HandlerError> {
        let (stream, _addr) = listener.accept().await?;
        let pid = stream.peer_cred()?
            .pid()
            .ok_or(ClientInfoError::PidNotFound)?
            as u32;

        Ok((stream, pid))
    }
}


#[cfg(windows)]
mod platform {
    use std::os::windows::io::AsRawHandle;
    use tokio::net::windows::named_pipe::{
        NamedPipeServer,
        ServerOptions,
    };
    use windows::Win32::{
        Foundation::HANDLE,
        System::Pipes::GetNamedPipeClientProcessId,
    };
    use super::*;


    pub type Stream = NamedPipeServer;

    pub fn bind(sock_name: &str) -> std::io::Result<(String, NamedPipeServer)> {
        let addr = creddy_cli::server_addr(sock_name);
        let listener = ServerOptions::new()
            .first_pipe_instance(true)
            .create(&addr)?;
        Ok((listener, addr))
    }

    pub async fn accept(listener: &mut NamedPipeServer, addr: &String) -> Result<(NamedPipeServer, u32), HandlerError> {
        // connect() just waits for a client to connect, it doesn't return anything
        listener.connect().await?;

        // unlike Unix sockets, a Windows NamedPipeServer *becomes* the open stream
        // once a client connects. If we want to keep listening, we have to construct
        // a new server and swap it in.
        let new_listener = ServerOptions::new().create(addr)?;
        let stream = std::mem::replace(listener, new_listener);

        let raw_handle = stream.as_raw_handle();
        let mut pid = 0u32;
        let handle = HANDLE(raw_handle as _);
        unsafe { GetNamedPipeClientProcessId(handle, &mut pid as *mut u32)? };
        Ok((stream, pid))
    }
}