finish SSH key support

src-tauri/src/srv/agent.rs

@@ -0,0 +1,121 @@
+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);
+}

src-tauri/src/srv/creddy_server.rs

@@ -0,0 +1,132 @@
+use tauri::{AppHandle, Manager};
+use tokio::io::{AsyncReadExt, AsyncWriteExt};
+use tokio::sync::oneshot;
+
+use crate::clientinfo::{self, Client};
+use crate::errors::*;
+use crate::ipc::{Approval, RequestNotification};
+use crate::shortcuts::{self, ShortcutAction};
+use crate::state::AppState;
+use super::{
+    CloseWaiter,
+    Request,
+    Response,
+    Stream,
+};
+
+
+pub fn serve(app_handle: AppHandle) -> std::io::Result<()> {
+    super::serve("creddy-server", app_handle, handle)
+}
+
+
+async fn handle(
+    mut stream: Stream,
+    app_handle: AppHandle,
+    client_pid: u32
+) -> Result<(), HandlerError> {
+    // read from stream until delimiter is reached
+    let mut buf: Vec<u8> = Vec::with_capacity(1024); // requests are small, 1KiB is more than enough
+    let mut n = 0;
+    loop {
+        n += stream.read_buf(&mut buf).await?;
+        if let Some(&b'\n') = buf.last() {
+            break;
+        }
+        // sanity check, no request should ever be within a mile of 1MB
+        else if n >= (1024 * 1024) {
+            return Err(HandlerError::RequestTooLarge);
+        }
+    }
+
+    let client = clientinfo::get_client(client_pid, true)?;
+    let waiter = CloseWaiter { stream: &mut stream };
+
+
+    let req: Request = serde_json::from_slice(&buf)?;
+    let res = match req {
+        Request::GetAwsCredentials { name, base } => get_aws_credentials(
+            name, base, client, app_handle, waiter
+        ).await,
+        Request::InvokeShortcut(action) => invoke_shortcut(action).await,
+        Request::GetSshSignature(_) => return Err(HandlerError::Denied),
+    };
+
+    // doesn't make sense to send the error to the client if the client has already left
+    if let Err(HandlerError::Abandoned) = res {
+        return Err(HandlerError::Abandoned);
+    }
+
+    let res = serde_json::to_vec(&res).unwrap();
+    stream.write_all(&res).await?;
+    Ok(())
+}
+
+
+async fn invoke_shortcut(action: ShortcutAction) -> Result<Response, HandlerError> {
+    shortcuts::exec_shortcut(action);
+    Ok(Response::Empty)
+}
+
+
+async fn get_aws_credentials(
+    name: Option<String>,
+    base: bool,
+    client: Client,
+    app_handle: AppHandle,
+    mut waiter: CloseWaiter<'_>,
+) -> Result<Response, 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)?; // automate this conversion eventually?
+
+    let (chan_send, chan_recv) = oneshot::channel();
+    let request_id = state.register_request(chan_send).await;
+
+    // if an error occurs in any of the following, we want to abort the operation
+    // but ? returns immediately, and we want to unregister the request before returning
+    // so we bundle it all up in an async block and return a Result so we can handle errors
+    let proceed = async {
+        let notification = RequestNotification::new_aws(
+            request_id, client, name.clone(), base
+        );
+        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);
+            },
+        };
+
+        match response.approval {
+            Approval::Approved => {
+                if response.base {
+                    let creds = state.get_aws_base(name).await?;
+                    Ok(Response::AwsBase(creds))
+                }
+                else {
+                    let creds = state.get_aws_session(name).await?;
+                    Ok(Response::AwsSession(creds.clone()))
+                }
+            },
+            Approval::Denied => Err(HandlerError::Denied),
+        }
+    };
+
+    let result = match proceed.await {
+        Ok(r) => Ok(r),
+        Err(e) => {
+            state.unregister_request(request_id).await;
+            Err(e)
+        },
+    };
+
+    lease.release();
+    result
+}

src-tauri/src/srv/mod.rs

@@ -0,0 +1,170 @@
+use std::future::Future;
+
+use tauri::{
+    AppHandle,
+    async_runtime as rt,
+};
+use tokio::io::AsyncReadExt;
+use serde::{Serialize, Deserialize};
+use ssh_agent_lib::proto::message::SignRequest;
+
+use crate::credentials::{AwsBaseCredential, AwsSessionCredential};
+use crate::errors::*;
+use crate::shortcuts::ShortcutAction;
+
+pub mod creddy_server;
+pub mod agent;
+use platform::Stream;
+pub use platform::addr;
+
+
+#[derive(Debug, Serialize, Deserialize)]
+pub enum Request {
+    GetAwsCredentials { 
+        name: Option<String>,
+        base: bool,
+    },
+    GetSshSignature(SignRequest),
+    InvokeShortcut(ShortcutAction),
+}
+
+
+#[derive(Debug, Serialize, Deserialize)]
+pub enum Response {
+    AwsBase(AwsBaseCredential),
+    AwsSession(AwsSessionCredential),
+    Empty,
+}
+
+
+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 = 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))
+    }
+
+
+    pub fn addr(sock_name: &str) -> PathBuf {
+        let mut path = dirs::runtime_dir()
+            .unwrap_or_else(|| PathBuf::from("/tmp"));
+        path.push(format!("{sock_name}.sock"));
+        path
+    }
+}
+
+
+#[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 = 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))
+    }
+
+    pub fn addr(sock_name: &str) -> String {
+        format!(r"\\.\pipe\{sock_name}")
+    }
+}