Merge pull request #21 from wiktor-k/more-useful-user-error

Relax the error type in handle
wiktor-k · Mar 15, 2024 · 2edc1fb · 2edc1fb
2 parents 8721f83 + 678d32c
commit 2edc1fb
Show file tree

Hide file tree

Showing 4 changed files with 13 additions and 17 deletions.
diff --git a/README.md b/README.md
@@ -15,15 +15,14 @@ This example starts listening on a Unix socket `ssh-agent.sock` and processes re
 use tokio::net::UnixListener;
 
 use ssh_agent_lib::agent::{Session, Agent};
-use ssh_agent_lib::error::AgentError;
 use ssh_agent_lib::proto::message::Message;
 
 #[derive(Default)]
 struct MyAgent;
 
 #[ssh_agent_lib::async_trait]
 impl Session for MyAgent {
-    async fn handle(&mut self, message: Message) -> Result<Message, AgentError> {
+    async fn handle(&mut self, message: Message) -> Result<Message, Box<dyn std::error::Error>> {
         match message {
             Message::SignRequest(request) => {
                 // get the signature by signing `request.data`

diff --git a/examples/key_storage.rs b/examples/key_storage.rs
@@ -3,7 +3,6 @@ use log::info;
 use tokio::net::UnixListener;
 
 use ssh_agent_lib::agent::{Agent, Session};
-use ssh_agent_lib::error::AgentError;
 use ssh_agent_lib::proto::message::{self, Message, SignRequest};
 use ssh_agent_lib::proto::private_key::{PrivateKey, RsaPrivateKey};
 use ssh_agent_lib::proto::public_key::PublicKey;
@@ -148,11 +147,8 @@ impl KeyStorage {
 
 #[async_trait]
 impl Session for KeyStorage {
-    async fn handle(&mut self, message: Message) -> Result<Message, AgentError> {
-        self.handle_message(message).or_else(|error| {
-            println!("Error handling message - {:?}", error);
-            Ok(Message::Failure)
-        })
+    async fn handle(&mut self, message: Message) -> Result<Message, Box<dyn std::error::Error>> {
+        self.handle_message(message)
     }
 }
 

diff --git a/src/agent.rs b/src/agent.rs
@@ -77,7 +77,7 @@ impl ListeningSocket for TcpListener {
 
 #[async_trait]
 pub trait Session: 'static + Sync + Send + Sized {
-    async fn handle(&mut self, message: Message) -> Result<Message, AgentError>;
+    async fn handle(&mut self, message: Message) -> Result<Message, Box<dyn std::error::Error>>;
 
     async fn handle_socket<S>(
         &mut self,
@@ -88,10 +88,13 @@ pub trait Session: 'static + Sync + Send + Sized {
     {
         loop {
             if let Some(incoming_message) = adapter.try_next().await? {
-                let response = self.handle(incoming_message).await.map_err(|e| {
-                    error!("Error handling message; error = {:?}", e);
-                    AgentError::User
-                })?;
+                let response = match self.handle(incoming_message).await {
+                    Ok(message) => message,
+                    Err(e) => {
+                        error!("Error handling message: {:?}", e);
+                        Message::Failure
+                    }
+                };
 
                 adapter.send(response).await?;
             } else {
@@ -118,12 +121,12 @@ pub trait Agent: 'static + Sync + Send + Sized {
                     tokio::spawn(async move {
                         let adapter = Framed::new(socket, MessageCodec);
                         if let Err(e) = session.handle_socket::<S>(adapter).await {
-                            error!("Agent protocol error; error = {:?}", e);
+                            error!("Agent protocol error: {:?}", e);
                         }
                     });
                 }
                 Err(e) => {
-                    error!("Failed to accept socket; error = {:?}", e);
+                    error!("Failed to accept socket: {:?}", e);
                     return Err(AgentError::IO(e));
                 }
             }

diff --git a/src/error.rs b/src/error.rs
@@ -3,7 +3,6 @@ use std::io;
 
 #[derive(Debug)]
 pub enum AgentError {
-    User,
     Proto(ProtoError),
     IO(io::Error),
 }
@@ -23,7 +22,6 @@ impl From<io::Error> for AgentError {
 impl std::fmt::Display for AgentError {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         match self {
-            AgentError::User => write!(f, "Agent: User error"),
             AgentError::Proto(proto) => write!(f, "Agent: Protocol error: {}", proto),
             AgentError::IO(error) => write!(f, "Agent: I/O error: {}", error),
         }