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 { let state = app_handle.state::(); 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 { let state = app_handle.state::(); 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", ¬ification)?; 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, s: &[u8]) { let len = s.len() as u32; buf.extend(len.to_be_bytes()); buf.extend(s); }