Merge pull request #65 from wiktor-k/wiktor/add-proto-dumper

Add protocol dumper example

Author: wiktor-k

examples/pgp-wrapper.rs

@@ -33,7 +33,7 @@
 //!
 //! Works perfectly in conjunction with `openpgp-card-agent.rs`!
 
-use std::{cell::RefCell, pin::Pin};
+use std::cell::RefCell;
 
 use chrono::DateTime;
 use clap::Parser;
@@ -55,11 +55,11 @@ use tokio::runtime::Runtime;
 struct WrappedKey {
     public_key: PublicKey,
     pubkey: KeyData,
-    client: RefCell<Pin<Box<dyn Session>>>,
+    client: RefCell<Box<dyn Session>>,
 }
 
 impl WrappedKey {
-    fn new(pubkey: KeyData, client: Pin<Box<dyn Session>>) -> Self {
+    fn new(pubkey: KeyData, client: Box<dyn Session>) -> Self {
         let KeyData::Ed25519(key) = pubkey.clone() else {
             panic!("The first key was not ed25519!");
         };
@@ -206,11 +206,11 @@ fn main() -> testresult::TestResult {
     let (client, identities) = rt.block_on(async move {
         #[cfg(unix)]
         let mut client =
-            connect(Binding::FilePath(std::env::var("SSH_AUTH_SOCK")?.into()).try_into()?).await?;
+            connect(Binding::FilePath(std::env::var("SSH_AUTH_SOCK")?.into()).try_into()?)?;
 
         #[cfg(windows)]
         let mut client =
-            connect(Binding::NamedPipe(std::env::var("SSH_AUTH_SOCK")?.into()).try_into()?).await?;
+            connect(Binding::NamedPipe(std::env::var("SSH_AUTH_SOCK")?.into()).try_into()?)?;
 
         let identities = client.request_identities().await?;
 

examples/proto-dumper.rs

@@ -0,0 +1,124 @@
+//! This example illustrates a couple of features: First, it
+//! implements a forwarder, exposing an SSH agent socket and
+//! forwarding to a different one. Secondly it shows how to work with
+//! low-level handling of messages instead of parsed high-level
+//! structures.
+//!
+//! Run with
+//! RUST_LOG=info cargo run --example proto-dumper -- --target unix://$SSH_AUTH_SOCK -H unix:///tmp/test.sock
+
+use clap::Parser;
+use service_binding::Binding;
+use ssh_agent_lib::{
+    agent::bind,
+    agent::Agent,
+    agent::Session,
+    async_trait,
+    client::connect,
+    error::AgentError,
+    proto::{Request, Response},
+};
+use ssh_encoding::Encode;
+
+struct DumpAndForward {
+    target: Box<dyn Session>,
+    session: u64,
+    id: u64,
+}
+
+#[async_trait]
+impl Session for DumpAndForward {
+    async fn handle(&mut self, message: Request) -> Result<Response, AgentError> {
+        use std::io::Write;
+
+        self.id += 1;
+        let req_file = format!("req-{}-{}.bin", self.session, self.id);
+        log::info!("Writing request {message:?} to {req_file}");
+
+        let mut req = std::fs::File::create(req_file)?;
+        let mut buf = vec![];
+        message.encode(&mut buf).map_err(AgentError::other)?;
+        req.write_all(&buf)?;
+        drop(req);
+
+        let response = self.target.handle(message).await?;
+
+        let resp_file = format!("resp-{}-{}.bin", self.session, self.id);
+        log::info!("Writing response {response:?} to {resp_file}");
+        let mut resp = std::fs::File::create(resp_file)?;
+        let mut buf = vec![];
+        response.encode(&mut buf).map_err(AgentError::other)?;
+        resp.write_all(&buf)?;
+        drop(resp);
+
+        Ok(response)
+    }
+}
+
+struct Forwarder {
+    target: Binding,
+    id: u64,
+}
+
+#[cfg(unix)]
+impl Agent<tokio::net::UnixListener> for Forwarder {
+    fn new_session(&mut self, _socket: &tokio::net::UnixStream) -> impl Session {
+        self.create_new_session()
+    }
+}
+
+impl Agent<tokio::net::TcpListener> for Forwarder {
+    fn new_session(&mut self, _socket: &tokio::net::TcpStream) -> impl Session {
+        self.create_new_session()
+    }
+}
+
+#[cfg(windows)]
+impl Agent<ssh_agent_lib::agent::NamedPipeListener> for Forwarder {
+    fn new_session(
+        &mut self,
+        _socket: &tokio::net::windows::named_pipe::NamedPipeServer,
+    ) -> impl Session {
+        self.create_new_session()
+    }
+}
+
+impl Forwarder {
+    fn create_new_session(&mut self) -> impl Session {
+        self.id += 1;
+        DumpAndForward {
+            target: connect(self.target.clone().try_into().unwrap()).unwrap(),
+            session: self.id,
+            id: 0,
+        }
+    }
+}
+
+#[derive(Debug, Parser)]
+struct Args {
+    /// Target SSH agent to which we will proxy all requests.
+    #[clap(long)]
+    target: Binding,
+
+    /// Source that we will bind to.
+    #[clap(long, short = 'H')]
+    host: Binding,
+}
+
+#[tokio::main]
+async fn main() -> Result<(), Box<dyn std::error::Error>> {
+    env_logger::init();
+
+    let args = Args::parse();
+
+    bind(
+        args.host.try_into()?,
+        Forwarder {
+            target: args.target,
+            id: 0,
+        },
+    )
+    .await?;
+
+    Ok(())
+}

examples/ssh-agent-client.rs

@@ -5,11 +5,11 @@ use ssh_agent_lib::client::connect;
 async fn main() -> Result<(), Box<dyn std::error::Error>> {
     #[cfg(unix)]
     let mut client =
-        connect(Binding::FilePath(std::env::var("SSH_AUTH_SOCK")?.into()).try_into()?).await?;
+        connect(Binding::FilePath(std::env::var("SSH_AUTH_SOCK")?.into()).try_into()?)?;
 
     #[cfg(windows)]
     let mut client =
-        connect(Binding::NamedPipe(std::env::var("SSH_AUTH_SOCK")?.into()).try_into()?).await?;
+        connect(Binding::NamedPipe(std::env::var("SSH_AUTH_SOCK")?.into()).try_into()?)?;
 
     eprintln!(
         "Identities that this agent knows of: {:#?}",

src/client.rs

@@ -37,18 +37,18 @@ where
 }
 
 /// Wrap a stream into an SSH agent client.
-pub async fn connect(
+pub fn connect(
     stream: service_binding::Stream,
-) -> Result<std::pin::Pin<Box<dyn crate::agent::Session>>, Box<dyn std::error::Error>> {
+) -> Result<Box<dyn crate::agent::Session>, Box<dyn std::error::Error>> {
     match stream {
         #[cfg(unix)]
         service_binding::Stream::Unix(stream) => {
             let stream = tokio::net::UnixStream::from_std(stream)?;
-            Ok(Box::pin(Client::new(stream)))
+            Ok(Box::new(Client::new(stream)))
         }
         service_binding::Stream::Tcp(stream) => {
             let stream = tokio::net::TcpStream::from_std(stream)?;
-            Ok(Box::pin(Client::new(stream)))
+            Ok(Box::new(Client::new(stream)))
         }
         #[cfg(windows)]
         service_binding::Stream::NamedPipe(pipe) => {
@@ -65,9 +65,9 @@ pub async fn connect(
                     Err(e) => Err(e)?,
                 }
 
-                tokio::time::sleep(std::time::Duration::from_millis(50)).await;
+                std::thread::sleep(std::time::Duration::from_millis(50));
             };
-            Ok(Box::pin(Client::new(stream)))
+            Ok(Box::new(Client::new(stream)))
         }
         #[cfg(not(windows))]
         service_binding::Stream::NamedPipe(_) => Err(ProtoError::IO(std::io::Error::other(