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

use std::io::ErrorKind;

use futures::SinkExt;
use signature::Signer;
use ssh_agent_lib::agent::MessageCodec;
use ssh_agent_lib::proto::message::{
    Message,
    Identity,
    SignRequest,
};
use tokio::net::{UnixListener, UnixStream};
use tauri::{
    AppHandle,
    Manager,
    async_runtime as rt,
};
use tokio_util::codec::Framed;
use tokio_stream::StreamExt;
use tokio::sync::oneshot;

use crate::clientinfo;
use crate::credentials::{Credential, SshKey};
use crate::errors::*;
use crate::ipc::{Approval, RequestNotification};
use crate::state::AppState;


pub struct Agent {
    listener: UnixListener,
    app_handle: AppHandle,
}

impl Agent {
    pub fn start(app_handle: AppHandle) -> std::io::Result<()> {
        match std::fs::remove_file("/tmp/creddy-agent.sock") {
            Ok(_) => (),
            Err(e) if e.kind() == ErrorKind::NotFound => (),
            Err(e) => return Err(e),
        }

        let listener = UnixListener::bind("/tmp/creddy-agent.sock")?;
        let srv = Agent { listener, app_handle };
        rt::spawn(srv.serve());
        Ok(())
    }

    async fn serve(self) {
        loop {
            self.try_serve()
                .await
                .error_print_prefix("Error accepting request: ");
        }
    }

    async fn try_serve(&self) -> Result<(), HandlerError> {
        let (stream, _addr) = self.listener.accept().await?;
        let new_handle = self.app_handle.clone();
        let client_pid = get_client_pid(&stream)?;
        rt::spawn(async move {
            let adapter = Framed::new(stream, MessageCodec);
            handle_framed(adapter, new_handle, client_pid)
                .await
                .error_print_prefix("Error responding to request: ");
        });
        Ok(())
    }
}


async fn handle_framed(
    mut adapter: Framed<UnixStream, MessageCodec>,
    app_handle: AppHandle,
    client_pid: u32,
) -> Result<(), HandlerError> {
    while let Some(message) = adapter.try_next().await? {
        let resp = match message {
            Message::RequestIdentities => list_identities(app_handle.clone()).await?,
            Message::SignRequest(req) => sign_request(req, app_handle.clone(), client_pid).await?,
            _ => Message::Failure,
        };

        adapter.send(resp).await?;
    }

    Ok(())
}


async fn list_identities(app_handle: AppHandle) -> Result<Message, HandlerError> {
    let state = app_handle.state::<AppState>();
    let identities: Vec<Identity> = state.list_ssh_identities().await?;
    Ok(Message::IdentitiesAnswer(identities))
}


async fn sign_request(req: SignRequest, app_handle: AppHandle, client_pid: u32) -> Result<Message, HandlerError> {
    let state = app_handle.state::<AppState>();
        let rehide_ms = {
        let config = state.config.read().await;
        config.rehide_ms
    };
    let client = clientinfo::get_client(client_pid, false)?;
    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 proceed = async {
        let key_name = state.ssh_name_from_pubkey(&req.pubkey_blob).await?;
        let notification = RequestNotification::new_ssh(request_id, client, key_name.clone());
        app_handle.emit("credential-request", &notification)?;

        let response = chan_recv.await?;
        if let Approval::Denied = response.approval {
            return Ok(Message::Failure);
        }

        let key = state.sshkey_by_name(&key_name).await?;
        let sig = Signer::sign(&key.private_key, &req.data);
        let key_type = key.algorithm.as_str().as_bytes();

        let payload_len = key_type.len() + sig.as_bytes().len() + 8;
        let mut payload = Vec::with_capacity(payload_len);
        encode_string(&mut payload, key.algorithm.as_str().as_bytes());
        encode_string(&mut payload, sig.as_bytes());

        Ok(Message::SignResponse(payload))
    };

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

    lease.release();
    res
}



fn get_client_pid(stream: &UnixStream) -> std::io::Result<u32> {
    let cred = stream.peer_cred()?;
    Ok(cred.pid().unwrap() as u32)
}


fn encode_string(buf: &mut Vec<u8>, s: &[u8]) {
    let len = s.len() as u32;
    buf.extend(len.to_be_bytes());
    buf.extend(s);
}