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

use std::future::Future;

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

use crate::credentials::{AwsBaseCredential, AwsSessionCredential};
use crate::errors::*;
use crate::shortcuts::ShortcutAction;

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 {
    GetCredential {
        name: Option<String>,
        base: bool,
    },
    InvokeShortcut(ShortcutAction),
}


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


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


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


#[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))
    }
}