1 files changed, 279 insertions, 0 deletions
diff --git a/src/ssh/agent.rs b/src/ssh/agent.rs
new file mode 100644
index 0000000..c29ad62
--- /dev/null
+++ b/src/ssh/agent.rs
@@ -0,0 +1,279 @@
+// Copyright © 2022 Kim Altintop <kim@eagain.io>
+// SPDX-License-Identifier: GPL-2.0-only WITH openvpn-openssl-exception
+
+use std::{
+    env,
+    io::{
+        self,
+        ErrorKind::*,
+    },
+};
+
+use anyhow::Context;
+use ssh_encoding::{
+    CheckedSum,
+    Decode,
+    Encode,
+    Reader,
+    Writer,
+};
+use ssh_key::{
+    public::KeyData,
+    Algorithm,
+    HashAlg,
+    PublicKey,
+    Signature,
+};
+
+#[cfg(unix)]
+pub use std::os::unix::net::UnixStream;
+#[cfg(windows)]
+pub use uds_windows::UnixStram;
+
+const SSH_AUTH_SOCK: &str = "SSH_AUTH_SOCK";
+
+const MAX_AGENT_REPLY_LEN: usize = 256 * 1024;
+
+const SSH_AGENTC_REQUEST_IDENTITIES: u8 = 11;
+const SSH_AGENTC_SIGN_REQUEST: u8 = 13;
+const SSH_AGENT_FAILURE: u8 = 5;
+const SSH_AGENT_IDENTITIES_ANSWER: u8 = 12;
+const SSH_AGENT_RSA_SHA2_256: u32 = 2;
+const SSH_AGENT_RSA_SHA2_512: u32 = 4;
+const SSH_AGENT_SIGN_RESPONSE: u8 = 14;
+
+pub struct Client<T> {
+    conn: T,
+}
+
+impl Client<UnixStream> {
+    pub fn from_env() -> io::Result<Self> {
+        let path = env::var_os(SSH_AUTH_SOCK).ok_or_else(|| {
+            io::Error::new(
+                io::ErrorKind::AddrNotAvailable,
+                "SSH_AUTH_SOCK environment variable not set",
+            )
+        })?;
+        UnixStream::connect(path).map(Self::from)
+    }
+}
+
+impl From<UnixStream> for Client<UnixStream> {
+    fn from(conn: UnixStream) -> Self {
+        Self { conn }
+    }
+}
+
+impl<'a> From<&'a UnixStream> for Client<&'a UnixStream> {
+    fn from(conn: &'a UnixStream) -> Self {
+        Self { conn }
+    }
+}
+
+impl<T> Client<T>
+where
+    T: io::Read + io::Write,
+{
+    pub fn sign(&mut self, key: &PublicKey, msg: impl AsRef<[u8]>) -> io::Result<Signature> {
+        request(
+            &mut self.conn,
+            SignRequest {
+                key,
+                msg: msg.as_ref(),
+            },
+        )
+        .map(|SignResponse { sig }| sig)
+    }
+
+    pub fn list_keys(&mut self) -> io::Result<Vec<PublicKey>> {
+        request(&mut self.conn, RequestIdentities).map(|IdentitiesAnswer { keys }| keys)
+    }
+}
+
+trait Request: Encode<Error = crate::Error> {
+    type Response: Response;
+}
+
+trait Response: Decode<Error = crate::Error> {
+    const SUCCESS: u8;
+}
+
+fn request<I, T>(mut io: I, req: T) -> io::Result<T::Response>
+where
+    I: io::Read + io::Write,
+    T: Request,
+{
+    send(&mut io, req)?;
+    let resp = recv(&mut io)?;
+    let mut reader = resp.as_slice();
+    match u8::decode(&mut reader).map_err(|_| unknown_response())? {
+        x if x == T::Response::SUCCESS => T::Response::decode(&mut reader).map_err(decode),
+        SSH_AGENT_FAILURE => Err(agent_error()),
+        _ => Err(unknown_response()),
+    }
+}
+
+fn send<W, T>(mut io: W, req: T) -> io::Result<()>
+where
+    W: io::Write,
+    T: Encode<Error = crate::Error>,
+{
+    let len = req.encoded_len_prefixed().map_err(encode)?;
+    let mut buf = Vec::with_capacity(len);
+    req.encode_prefixed(&mut buf).map_err(encode)?;
+
+    io.write_all(&buf)?;
+    io.flush()?;
+
+    Ok(())
+}
+
+fn recv<R: io::Read>(mut io: R) -> io::Result<Vec<u8>> {
+    let want = {
+        let mut buf = [0; 4];
+        io.read_exact(&mut buf)?;
+        u32::from_be_bytes(buf) as usize
+    };
+
+    if want < 1 {
+        return Err(incomplete_response());
+    }
+    if want > MAX_AGENT_REPLY_LEN {
+        return Err(reponse_too_large());
+    }
+
+    let mut buf = vec![0; want];
+    io.read_exact(&mut buf)?;
+
+    Ok(buf)
+}
+
+struct SignRequest<'a> {
+    key: &'a PublicKey,
+    msg: &'a [u8],
+}
+
+impl Request for SignRequest<'_> {
+    type Response = SignResponse;
+}
+
+impl Encode for SignRequest<'_> {
+    type Error = crate::Error;
+
+    fn encoded_len(&self) -> Result<usize, Self::Error> {
+        Ok([
+            self.key.key_data().encoded_len_prefixed()?,
+            self.msg.encoded_len()?,
+            SSH_AGENTC_SIGN_REQUEST.encoded_len()?,
+            4, // flags
+        ]
+        .checked_sum()?)
+    }
+
+    fn encode(&self, writer: &mut impl Writer) -> Result<(), Self::Error> {
+        SSH_AGENTC_SIGN_REQUEST.encode(writer)?;
+        self.key.key_data().encode_prefixed(writer)?;
+        self.msg.encode(writer)?;
+        let flags = match self.key.algorithm() {
+            Algorithm::Rsa { hash } => match hash {
+                Some(HashAlg::Sha256) => SSH_AGENT_RSA_SHA2_256,
+                _ => SSH_AGENT_RSA_SHA2_512, // sane default
+            },
+            _ => 0,
+        };
+        flags.encode(writer)?;
+        Ok(())
+    }
+}
+
+struct SignResponse {
+    sig: Signature,
+}
+
+impl Response for SignResponse {
+    const SUCCESS: u8 = SSH_AGENT_SIGN_RESPONSE;
+}
+
+impl Decode for SignResponse {
+    type Error = crate::Error;
+
+    fn decode(reader: &mut impl Reader) -> Result<Self, Self::Error> {
+        let sig = reader.read_prefixed(Signature::decode)?;
+        Ok(Self { sig })
+    }
+}
+
+struct RequestIdentities;
+
+impl Request for RequestIdentities {
+    type Response = IdentitiesAnswer;
+}
+
+impl Encode for RequestIdentities {
+    type Error = crate::Error;
+
+    fn encoded_len(&self) -> Result<usize, Self::Error> {
+        Ok(SSH_AGENTC_REQUEST_IDENTITIES.encoded_len()?)
+    }
+
+    fn encode(&self, writer: &mut impl Writer) -> Result<(), Self::Error> {
+        Ok(SSH_AGENTC_REQUEST_IDENTITIES.encode(writer)?)
+    }
+}
+
+struct IdentitiesAnswer {
+    keys: Vec<PublicKey>,
+}
+
+impl Response for IdentitiesAnswer {
+    const SUCCESS: u8 = SSH_AGENT_IDENTITIES_ANSWER;
+}
+
+impl Decode for IdentitiesAnswer {
+    type Error = crate::Error;
+
+    fn decode(reader: &mut impl Reader) -> Result<Self, Self::Error> {
+        let nkeys = usize::decode(reader).context("nkeys")?;
+        let mut keys = Vec::with_capacity(nkeys);
+
+        for _ in 0..nkeys {
+            let key_data = reader.read_prefixed(KeyData::decode).context("key data")?;
+            let comment = String::decode(reader).context("comment")?;
+            keys.push(PublicKey::new(key_data, comment));
+        }
+
+        Ok(Self { keys })
+    }
+}
+
+fn e(kind: io::ErrorKind, msg: &str) -> io::Error {
+    io::Error::new(kind, msg)
+}
+
+fn ee(kind: io::ErrorKind, e: crate::Error) -> io::Error {
+    io::Error::new(kind, e)
+}
+
+fn incomplete_response() -> io::Error {
+    e(UnexpectedEof, "incomplete response")
+}
+
+fn reponse_too_large() -> io::Error {
+    e(Unsupported, "response payload too large")
+}
+
+fn encode(e: crate::Error) -> io::Error {
+    ee(InvalidData, e.context("failed to encode request"))
+}
+
+fn decode(e: crate::Error) -> io::Error {
+    ee(InvalidData, e.context("failed to decode response"))
+}
+
+fn agent_error() -> io::Error {
+    e(Other, "error response from agent")
+}
+
+fn unknown_response() -> io::Error {
+    e(Unsupported, "unknown response")
+}