Extract listen from the Agent

Signed-off-by: Wiktor Kwapisiewicz <wiktor@metacode.biz>

Author: wiktor-k

README.md

@@ -18,7 +18,7 @@ use tokio::net::UnixListener as Listener;
 #[cfg(windows)]
 use ssh_agent_lib::agent::NamedPipeListener as Listener;
 use ssh_agent_lib::error::AgentError;
-use ssh_agent_lib::agent::{Session, Agent};
+use ssh_agent_lib::agent::{Session, listen};
 use ssh_agent_lib::proto::{Identity, SignRequest};
 use ssh_key::{Algorithm, Signature};
 
@@ -50,7 +50,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
 
     let _ = std::fs::remove_file(socket); // remove the socket if exists
 
-    MyAgent.listen(Listener::bind(socket)?).await?;
+    listen(Listener::bind(socket)?, MyAgent::default()).await?;
     Ok(())
 }
 ```

examples/key_storage.rs

@@ -9,14 +9,13 @@ use rsa::BigUint;
 use sha1::Sha1;
 #[cfg(windows)]
 use ssh_agent_lib::agent::NamedPipeListener as Listener;
-use ssh_agent_lib::agent::{ListeningSocket, Session};
+use ssh_agent_lib::agent::{listen, Agent, Session};
 use ssh_agent_lib::error::AgentError;
 use ssh_agent_lib::proto::extension::{QueryResponse, RestrictDestination, SessionBind};
 use ssh_agent_lib::proto::{
     message, signature, AddIdentity, AddIdentityConstrained, AddSmartcardKeyConstrained,
     Credential, Extension, KeyConstraint, RemoveIdentity, SignRequest, SmartcardKey,
 };
-use ssh_agent_lib::Agent;
 use ssh_key::{
     private::{KeypairData, PrivateKey},
     public::PublicKey,
@@ -237,11 +236,21 @@ impl KeyStorageAgent {
     }
 }
 
-impl Agent for KeyStorageAgent {
-    fn new_session<S>(&mut self, _socket: &S::Stream) -> impl Session
-    where
-        S: ListeningSocket + std::fmt::Debug + Send,
-    {
+#[cfg(unix)]
+impl Agent<Listener> for KeyStorageAgent {
+    fn new_session(&mut self, _socket: &tokio::net::UnixStream) -> impl Session {
+        KeyStorage {
+            identities: Arc::clone(&self.identities),
+        }
+    }
+}
+
+#[cfg(windows)]
+impl Agent<Listener> for KeyStorageAgent {
+    fn new_session(
+        &mut self,
+        _socket: &tokio::net::windows::named_pipe::NamedPipeServer,
+    ) -> impl Session {
         KeyStorage {
             identities: Arc::clone(&self.identities),
         }
@@ -263,8 +272,6 @@ async fn main() -> Result<(), AgentError> {
     #[cfg(windows)]
     std::fs::File::create("server-started")?;
 
-    KeyStorageAgent::new()
-        .listen(Listener::bind(socket)?)
-        .await?;
+    listen(Listener::bind(socket)?, KeyStorageAgent::new()).await?;
     Ok(())
 }

examples/openpgp-card-agent.rs

@@ -16,6 +16,11 @@
 
 use std::{sync::Arc, time::Duration};
 
+#[cfg(windows)]
+use ssh_agent_lib::agent::NamedPipeListener as Listener;
+
+#[cfg(not(windows))]
+use tokio::net::UnixListener as Listener;
 use card_backend_pcsc::PcscBackend;
 use clap::Parser;
 use openpgp_card::{
@@ -27,16 +32,16 @@ use retainer::{Cache, CacheExpiration};
 use secrecy::{ExposeSecret, SecretString};
 use service_binding::Binding;
 use ssh_agent_lib::{
-    agent::Session,
+    agent::{bind, Session, Agent},
     error::AgentError,
     proto::{AddSmartcardKeyConstrained, Identity, KeyConstraint, SignRequest, SmartcardKey},
-    Agent,
 };
 use ssh_key::{
     public::{Ed25519PublicKey, KeyData},
     Algorithm, Signature,
 };
 use testresult::TestResult;
+use tokio::net::TcpListener;
 
 struct CardAgent {
     pwds: Arc<Cache<String, SecretString>>,
@@ -51,8 +56,30 @@ impl CardAgent {
     }
 }
 
-impl Agent for CardAgent {
-    fn new_session(&mut self) -> impl Session {
+#[cfg(unix)]
+impl Agent<Listener> for CardAgent {
+    fn new_session(&mut self, _socket: &tokio::net::UnixStream) -> impl Session {
+        CardSession {
+            pwds: Arc::clone(&self.pwds),
+        }
+    }
+}
+
+#[cfg(unix)]
+impl Agent<TcpListener> for CardAgent {
+    fn new_session(&mut self, _socket: &tokio::net::TcpStream) -> impl Session {
+        CardSession {
+            pwds: Arc::clone(&self.pwds),
+        }
+    }
+}
+
+#[cfg(windows)]
+impl Agent<Listener> for CardAgent {
+    fn new_session(
+        &mut self,
+        _socket: &tokio::net::windows::named_pipe::NamedPipeServer,
+    ) -> impl Session {
         CardSession {
             pwds: Arc::clone(&self.pwds),
         }
@@ -201,6 +228,6 @@ async fn main() -> TestResult {
     env_logger::init();
 
     let args = Args::parse();
-    CardAgent::new().bind(args.host.try_into()?).await?;
+    bind(args.host.try_into()?, CardAgent::new()).await?;
     Ok(())
 }

src/agent.rs

@@ -249,69 +249,105 @@ where
     }
 }
 
-/// Type representing an agent listening for incoming connections.
-#[async_trait]
-pub trait Agent: 'static + Sync + Send + Sized {
-    /// Create new session object when a new socket is accepted.
-    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>
-    where
-        S: ListeningSocket + fmt::Debug + Send,
-    {
-        log::info!("Listening; socket = {:?}", socket);
-        loop {
-            match socket.accept().await {
-                Ok(socket) => {
-                    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 {
-                            log::error!("Agent protocol error: {:?}", e);
-                        }
-                    });
-                }
-                Err(e) => {
-                    log::error!("Failed to accept socket: {:?}", e);
-                    return Err(AgentError::IO(e));
-                }
-            }
-        }
-    }
+/// Factory of sessions for the given type of sockets.
+pub trait Agent<S>: 'static + Send + Sync
+where
+    S: ListeningSocket + fmt::Debug + Send,
+{
+    /// Create a [`Session`] object for a given `socket`.
+    fn new_session(&mut self, socket: &S::Stream) -> impl Session;
+}
 
-    /// Bind to a service binding listener.
-    async fn bind(mut self, listener: service_binding::Listener) -> Result<(), AgentError> {
-        match listener {
-            #[cfg(unix)]
-            service_binding::Listener::Unix(listener) => {
-                self.listen(UnixListener::from_std(listener)?).await
-            }
-            service_binding::Listener::Tcp(listener) => {
-                self.listen(TcpListener::from_std(listener)?).await
+/// Listen for connections on a given socket and use session factory
+/// to create new session for each accepted socket.
+pub async fn listen<S>(mut socket: S, mut sf: impl Agent<S>) -> Result<(), AgentError>
+where
+    S: ListeningSocket + fmt::Debug + Send,
+{
+    log::info!("Listening; socket = {:?}", socket);
+    loop {
+        match socket.accept().await {
+            Ok(socket) => {
+                let session = sf.new_session(&socket);
+                tokio::spawn(async move {
+                    let adapter = Framed::new(socket, Codec::<Request, Response>::default());
+                    if let Err(e) = handle_socket::<S>(session, adapter).await {
+                        log::error!("Agent protocol error: {:?}", e);
+                    }
+                });
             }
-            #[cfg(windows)]
-            service_binding::Listener::NamedPipe(pipe) => {
-                self.listen(NamedPipeListener::bind(pipe)?).await
+            Err(e) => {
+                log::error!("Failed to accept socket: {:?}", e);
+                return Err(AgentError::IO(e));
             }
-            #[cfg(not(windows))]
-            service_binding::Listener::NamedPipe(_) => Err(AgentError::IO(std::io::Error::other(
-                "Named pipes supported on Windows only",
-            ))),
         }
     }
 }
 
-impl<T> Agent for T
+#[cfg(unix)]
+impl<T> Agent<tokio::net::UnixListener> for T
+where
+    T: Default + Send + Sync + Session,
+{
+    fn new_session(&mut self, _socket: &tokio::net::UnixStream) -> impl Session {
+        Self::default()
+    }
+}
+
+impl<T> Agent<tokio::net::TcpListener> for T
 where
-    T: Default + Session,
+    T: Default + Send + Sync + Session,
 {
-    fn new_session<S>(&mut self, _socket: &S::Stream) -> impl Session
-    where
-        S: ListeningSocket + fmt::Debug + Send,
-    {
+    fn new_session(&mut self, _socket: &tokio::net::TcpStream) -> impl Session {
         Self::default()
     }
 }
+
+#[cfg(windows)]
+impl<T> Agent<NamedPipeListener> for T
+where
+    T: Default + Send + Sync + Session,
+{
+    fn new_session(
+        &mut self,
+        _socket: &tokio::net::windows::named_pipe::NamedPipeServer,
+    ) -> impl Session {
+        Self::default()
+    }
+}
+
+/// Bind to a service binding listener.
+#[cfg(unix)]
+pub async fn bind<SF>(listener: service_binding::Listener, sf: SF) -> Result<(), AgentError>
+where
+    SF: Agent<tokio::net::UnixListener> + Agent<tokio::net::TcpListener>,
+{
+    match listener {
+        #[cfg(unix)]
+        service_binding::Listener::Unix(listener) => {
+            listen(UnixListener::from_std(listener)?, sf).await
+        }
+        service_binding::Listener::Tcp(listener) => {
+            listen(TcpListener::from_std(listener)?, sf).await
+        }
+        _ => Err(AgentError::IO(std::io::Error::other(
+            "Unsupported type of a listener.",
+        ))),
+    }
+}
+
+/// Bind to a service binding listener.
+#[cfg(windows)]
+pub async fn bind<SF>(listener: service_binding::Listener, sf: SF) -> Result<(), AgentError>
+where
+    SF: Agent<NamedPipeListener> + Agent<tokio::net::TcpListener>,
+{
+    match listener {
+        service_binding::Listener::Tcp(listener) => {
+            listen(TcpListener::from_std(listener)?, sf).await
+        }
+        service_binding::Listener::NamedPipe(pipe) => {
+            listen(NamedPipeListener::bind(pipe)?, sf).await
+        }
+    }
+}

src/lib.rs

@@ -15,6 +15,3 @@ pub mod error;
 
 #[cfg(feature = "agent")]
 pub use async_trait::async_trait;
-
-#[cfg(feature = "agent")]
-pub use self::agent::Agent;