WIP: Add ClientInfo type to ListeningSocket

Cargo.toml

@@ -32,6 +32,13 @@ thiserror = "1.0.58"
 #uuid = { version = "1.8.0", features = ["v4"] }
 subtle = { version = "2", default-features = false }
 
+[target.'cfg(windows)'.dependencies.windows]
+version = "0.54"
+features = [
+    "Win32_Foundation",
+    "Win32_System_Pipes",
+]
+
 [features]
 default = ["agent"]
 codec = ["tokio-util"]

examples/key_storage.rs

@@ -9,7 +9,7 @@ use rsa::BigUint;
 use sha1::Sha1;
 #[cfg(windows)]
 use ssh_agent_lib::agent::NamedPipeListener as Listener;
-use ssh_agent_lib::agent::Session;
+use ssh_agent_lib::agent::{ListeningSocket, Session};
 use ssh_agent_lib::error::AgentError;
 use ssh_agent_lib::proto::extension::{QueryResponse, RestrictDestination, SessionBind};
 use ssh_agent_lib::proto::{
@@ -235,7 +235,14 @@ impl KeyStorageAgent {
 }
 
 impl Agent for KeyStorageAgent {
-    fn new_session(&mut self) -> impl Session {
+    fn new_session<S>(&mut self, socket: &S::Stream) -> impl Session
+    where
+        S: ListeningSocket + std::fmt::Debug + Send,
+    {
+        if let Ok(client_info) = S::client_info(socket) {
+            println!("New connection from client: {:?}", client_info);
+        }
+
         KeyStorage {
             identities: Arc::clone(&self.identities),
         }

src/agent.rs

@@ -1,19 +1,18 @@
 //! Traits for implementing custom SSH agents
 
+pub mod listener;
+
 use std::fmt;
-use std::io;
 
 use async_trait::async_trait;
 use futures::{SinkExt, TryStreamExt};
 use ssh_key::Signature;
-use tokio::io::{AsyncRead, AsyncWrite};
-#[cfg(windows)]
-use tokio::net::windows::named_pipe::{NamedPipeServer, ServerOptions};
-use tokio::net::{TcpListener, TcpStream};
+use tokio::net::TcpListener;
 #[cfg(unix)]
-use tokio::net::{UnixListener, UnixStream};
+use tokio::net::UnixListener;
 use tokio_util::codec::Framed;
 
+pub use self::listener::*;
 use super::error::AgentError;
 use super::proto::message::{Request, Response};
 use crate::codec::Codec;
@@ -27,65 +26,6 @@ use crate::proto::RemoveIdentity;
 use crate::proto::SignRequest;
 use crate::proto::SmartcardKey;
 
-/// Type representing a socket that asynchronously returns a list of streams.
-#[async_trait]
-pub trait ListeningSocket {
-    /// Stream type that represents an accepted socket.
-    type Stream: fmt::Debug + AsyncRead + AsyncWrite + Send + Unpin + 'static;
-
-    /// Waits until a client connects and returns connected stream.
-    async fn accept(&mut self) -> io::Result<Self::Stream>;
-}
-
-#[cfg(unix)]
-#[async_trait]
-impl ListeningSocket for UnixListener {
-    type Stream = UnixStream;
-    async fn accept(&mut self) -> io::Result<Self::Stream> {
-        UnixListener::accept(self).await.map(|(s, _addr)| s)
-    }
-}
-
-#[async_trait]
-impl ListeningSocket for TcpListener {
-    type Stream = TcpStream;
-    async fn accept(&mut self) -> io::Result<Self::Stream> {
-        TcpListener::accept(self).await.map(|(s, _addr)| s)
-    }
-}
-
-/// Listener for Windows Named Pipes.
-#[cfg(windows)]
-#[derive(Debug)]
-pub struct NamedPipeListener(NamedPipeServer, std::ffi::OsString);
-
-#[cfg(windows)]
-impl NamedPipeListener {
-    /// Bind to a pipe path.
-    pub fn bind(pipe: impl Into<std::ffi::OsString>) -> std::io::Result<Self> {
-        let pipe = pipe.into();
-        Ok(NamedPipeListener(
-            ServerOptions::new()
-                .first_pipe_instance(true)
-                .create(&pipe)?,
-            pipe,
-        ))
-    }
-}
-
-#[cfg(windows)]
-#[async_trait]
-impl ListeningSocket for NamedPipeListener {
-    type Stream = NamedPipeServer;
-    async fn accept(&mut self) -> io::Result<Self::Stream> {
-        self.0.connect().await?;
-        Ok(std::mem::replace(
-            &mut self.0,
-            ServerOptions::new().create(&self.1)?,
-        ))
-    }
-}
-
 /// Represents one active SSH connection.
 ///
 /// This type is implemented by agents that want to handle incoming SSH agent
@@ -253,7 +193,9 @@ where
 #[async_trait]
 pub trait Agent: 'static + Sync + Send + Sized {
     /// Create new session object when a new socket is accepted.
-    fn new_session(&mut self) -> impl Session;
+    fn new_session<S>(&mut self, socket: &S::Stream) -> impl Session
+    where
+        S: ListeningSocket + fmt::Debug + Send;
 
     /// Listen on a socket waiting for client connections.
     async fn listen<S>(mut self, mut socket: S) -> Result<(), AgentError>
@@ -264,7 +206,7 @@ pub trait Agent: 'static + Sync + Send + Sized {
         loop {
             match socket.accept().await {
                 Ok(socket) => {
-                    let session = self.new_session();
+                    let session = self.new_session::<S>(&socket);
                     tokio::spawn(async move {
                         let adapter = Framed::new(socket, Codec::<Request, Response>::default());
                         if let Err(e) = handle_socket::<S>(session, adapter).await {
@@ -306,7 +248,10 @@ impl<T> Agent for T
 where
     T: Default + Session,
 {
-    fn new_session(&mut self) -> impl Session {
+    fn new_session<S>(&mut self, _socket: &S::Stream) -> impl Session
+    where
+        S: ListeningSocket + fmt::Debug + Send,
+    {
         Self::default()
     }
 }