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

122 lines
3.7 KiB
Rust

use futures::SinkExt;
use signature::Signer;
use ssh_agent_lib::agent::MessageCodec;
use ssh_agent_lib::proto::message::{
Message,
SignRequest,
};
use tauri::{AppHandle, Manager};
use tokio_stream::StreamExt;
use tokio::sync::oneshot;
use tokio_util::codec::Framed;
use crate::clientinfo;
use crate::errors::*;
use crate::ipc::{Approval, RequestNotification};
use crate::state::AppState;
use super::{CloseWaiter, Stream};
pub fn serve(app_handle: AppHandle) -> std::io::Result<()> {
super::serve("creddy-agent", app_handle, handle)
}
async fn handle(
stream: Stream,
app_handle: AppHandle,
client_pid: u32
) -> Result<(), HandlerError> {
let mut adapter = Framed::new(stream, MessageCodec);
while let Some(message) = adapter.try_next().await? {
match message {
Message::RequestIdentities => {
let resp = list_identities(app_handle.clone()).await?;
adapter.send(resp).await?;
},
Message::SignRequest(req) => {
// CloseWaiter could corrupt the framing, but this doesn't matter
// since we don't plan to pull any more frames out of the stream
let waiter = CloseWaiter { stream: adapter.get_mut() };
let resp = sign_request(req, app_handle.clone(), client_pid, waiter).await?;
adapter.send(resp).await?;
break;
},
_ => adapter.send(Message::Failure).await?,
};
}
Ok(())
}
async fn list_identities(app_handle: AppHandle) -> Result<Message, HandlerError> {
let state = app_handle.state::<AppState>();
let identities = state.list_ssh_identities().await?;
Ok(Message::IdentitiesAnswer(identities))
}
async fn sign_request(
req: SignRequest,
app_handle: AppHandle,
client_pid: u32,
mut waiter: CloseWaiter<'_>,
) -> 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 = tokio::select! {
r = chan_recv => r?,
_ = waiter.wait_for_close() => {
app_handle.emit("request-cancelled", request_id)?;
return Err(HandlerError::Abandoned);
},
};
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(_) = &res {
state.unregister_request(request_id).await;
}
lease.release();
res
}
fn encode_string(buf: &mut Vec<u8>, s: &[u8]) {
let len = s.len() as u32;
buf.extend(len.to_be_bytes());
buf.extend(s);
}