refactor(auth): connection hooks

feat(req,rep,pub,sub): support custom connection hooks

Author: Karrq

Cargo.lock

@@ -0,0 +1,160 @@
+//! Connection hooks for customizing connection setup.
+//!
+//! The [`ConnectionHook`] trait provides a way to intercept connections during setup,
+//! enabling custom authentication, handshakes, or protocol negotiations.
+//!
+//! # Built-in Hooks
+//!
+//! The [`token`] module provides ready-to-use token-based authentication hooks:
+//! - [`token::ServerHook`] - Server-side hook that validates client tokens
+//! - [`token::ClientHook`] - Client-side hook that sends a token to the server
+//!
+//! # Custom Hooks
+//!
+//! Implement [`ConnectionHook`] for custom authentication or protocol negotiation:
+//!
+//! ```rust,ignore
+//! use msg_socket::ConnectionHook;
+//! use std::io;
+//! use tokio::io::{AsyncRead, AsyncWrite, AsyncReadExt, AsyncWriteExt};
+//!
+//! struct MyAuth;
+//!
+//! impl<Io> ConnectionHook<Io> for MyAuth
+//! where
+//!     Io: AsyncRead + AsyncWrite + Send + Unpin + 'static,
+//! {
+//!     async fn on_connection(&self, mut io: Io) -> Result<Io, HookError> {
+//!         let mut buf = [0u8; 32];
+//!         io.read_exact(&mut buf).await?;
+//!         if &buf == b"expected_token_value_32_bytes!!!" {
+//!             io.write_all(b"OK").await?;
+//!             Ok(io)
+//!         } else {
+//!             Err(HookError::custom("invalid token"))
+//!         }
+//!     }
+//! }
+//! ```
+
+use std::{error::Error as StdError, fmt, future::Future, io, pin::Pin, sync::Arc};
+
+use tokio::io::{AsyncRead, AsyncWrite};
+
+pub mod token;
+
+/// Error type for connection hooks.
+///
+/// This enum provides two variants:
+/// - `Io` for standard I/O errors
+/// - `Custom` for hook-specific errors (type-erased)
+#[derive(Debug)]
+pub enum HookError {
+    /// An I/O error occurred.
+    Io(io::Error),
+    /// A custom hook-specific error.
+    Custom(Box<dyn StdError + Send + Sync + 'static>),
+}
+
+impl HookError {
+    /// Creates a custom error from any error type.
+    pub fn custom<E>(error: E) -> Self
+    where
+        E: StdError + Send + Sync + 'static,
+    {
+        Self::Custom(Box::new(error))
+    }
+
+    /// Creates a custom error from a string message.
+    pub fn message(msg: impl Into<String>) -> Self {
+        Self::Io(io::Error::other(msg.into()))
+    }
+}
+
+impl fmt::Display for HookError {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match self {
+            Self::Io(e) => write!(f, "IO error: {e}"),
+            Self::Custom(e) => write!(f, "Hook error: {e}"),
+        }
+    }
+}
+
+impl StdError for HookError {
+    fn source(&self) -> Option<&(dyn StdError + 'static)> {
+        match self {
+            Self::Io(e) => Some(e),
+            Self::Custom(e) => Some(e.as_ref()),
+        }
+    }
+}
+
+impl From<io::Error> for HookError {
+    fn from(error: io::Error) -> Self {
+        Self::Io(error)
+    }
+}
+
+/// Hook executed during connection setup.
+///
+/// For server sockets (Rep, Pub): called when a connection is accepted.
+/// For client sockets (Req, Sub): called after connecting.
+///
+/// The hook receives the raw IO stream and has full control over the handshake protocol.
+pub trait ConnectionHook<Io>: Send + Sync + 'static
+where
+    Io: AsyncRead + AsyncWrite + Send + Unpin + 'static,
+{
+    /// Called when a connection is established.
+    ///
+    /// # Arguments
+    /// * `io` - The raw IO stream for this connection
+    ///
+    /// # Returns
+    /// The IO stream on success (potentially wrapped/transformed), or an error to reject
+    /// the connection.
+    fn on_connection(&self, io: Io) -> impl Future<Output = Result<Io, HookError>> + Send;
+}
+
+// ============================================================================
+// Type-erased hook for internal use
+// ============================================================================
+
+/// Type-erased connection hook for internal use.
+///
+/// This trait allows storing hooks with different concrete types behind a single
+/// `Arc<dyn ConnectionHookErased<Io>>`.
+pub(crate) trait ConnectionHookErased<Io>: Send + Sync + 'static
+where
+    Io: AsyncRead + AsyncWrite + Send + Unpin + 'static,
+{
+    fn on_connection(
+        self: Arc<Self>,
+        io: Io,
+    ) -> Pin<Box<dyn Future<Output = Result<Io, HookError>> + Send + 'static>>;
+}
+
+impl<T, Io> ConnectionHookErased<Io> for T
+where
+    T: ConnectionHook<Io>,
+    Io: AsyncRead + AsyncWrite + Send + Unpin + 'static,
+{
+    fn on_connection(
+        self: Arc<Self>,
+        io: Io,
+    ) -> Pin<Box<dyn Future<Output = Result<Io, HookError>> + Send + 'static>> {
+        Box::pin(async move { ConnectionHook::on_connection(&*self, io).await })
+    }
+}
+
+// ============================================================================
+// Hook result type for driver tasks
+// ============================================================================
+
+/// The result of running a connection hook.
+///
+/// Contains the processed IO stream and associated address.
+pub(crate) struct HookResult<Io, A> {
+    pub(crate) stream: Io,
+    pub(crate) addr: A,
+}

msg-socket/src/hooks/mod.rs

@@ -0,0 +1,161 @@
+//! Token-based authentication hooks.
+//!
+//! This module provides ready-to-use connection hooks for simple token-based authentication:
+//!
+//! - [`ServerHook`] - Server-side hook that validates client tokens
+//! - [`ClientHook`] - Client-side hook that sends a token to the server
+//!
+//! # Example
+//!
+//! ```rust,ignore
+//! use msg_socket::{RepSocket, ReqSocket, tcp::Tcp};
+//! use msg_socket::hooks::token::{ServerHook, ClientHook};
+//! use bytes::Bytes;
+//!
+//! // Server side - validates incoming tokens
+//! let rep = RepSocket::new(Tcp::default())
+//!     .with_connection_hook(ServerHook::new(|token| {
+//!         // Custom validation logic
+//!         token == b"secret"
+//!     }));
+//!
+//! // Client side - sends token on connect
+//! let req = ReqSocket::new(Tcp::default())
+//!     .with_connection_hook(ClientHook::new(Bytes::from("secret")));
+//! ```
+
+use bytes::Bytes;
+use futures::SinkExt;
+use tokio::io::{AsyncRead, AsyncWrite};
+use tokio_stream::StreamExt;
+use tokio_util::codec::Framed;
+
+use crate::hooks::{ConnectionHook, HookError};
+use msg_wire::auth;
+
+/// Server-side authentication hook that validates incoming client tokens.
+///
+/// When a client connects, this hook:
+/// 1. Waits for the client to send an auth token
+/// 2. Validates the token using the provided validator function
+/// 3. Sends an ACK on success, or rejects the connection on failure
+///
+/// # Example
+///
+/// ```rust,ignore
+/// use msg_socket::hooks::token::ServerHook;
+///
+/// // Accept all tokens
+/// let hook = ServerHook::accept_all();
+///
+/// // Custom validation
+/// let hook = ServerHook::new(|token| token == b"my_secret_token");
+/// ```
+pub struct ServerHook<F> {
+    validator: F,
+}
+
+impl ServerHook<fn(&Bytes) -> bool> {
+    /// Creates a server hook that accepts all tokens.
+    pub fn accept_all() -> Self {
+        Self { validator: |_| true }
+    }
+}
+
+impl<F> ServerHook<F>
+where
+    F: Fn(&Bytes) -> bool + Send + Sync + 'static,
+{
+    /// Creates a new server hook with the given validator function.
+    ///
+    /// The validator receives the client's token and returns `true` to accept
+    /// the connection or `false` to reject it.
+    pub fn new(validator: F) -> Self {
+        Self { validator }
+    }
+}
+
+impl<Io, F> ConnectionHook<Io> for ServerHook<F>
+where
+    Io: AsyncRead + AsyncWrite + Send + Unpin + 'static,
+    F: Fn(&Bytes) -> bool + Send + Sync + 'static,
+{
+    async fn on_connection(&self, io: Io) -> Result<Io, HookError> {
+        let mut conn = Framed::new(io, auth::Codec::new_server());
+
+        // Wait for client authentication message
+        let msg = conn
+            .next()
+            .await
+            .ok_or_else(|| HookError::message("connection closed"))?
+            .map_err(HookError::custom)?;
+
+        let auth::Message::Auth(token) = msg else {
+            return Err(HookError::message("expected auth message"));
+        };
+
+        // Validate the token
+        if !(self.validator)(&token) {
+            conn.send(auth::Message::Reject).await?;
+            return Err(HookError::message("authentication rejected"));
+        }
+
+        // Send acknowledgment
+        conn.send(auth::Message::Ack).await?;
+
+        Ok(conn.into_inner())
+    }
+}
+
+/// Client-side authentication hook that sends a token to the server.
+///
+/// When connecting to a server, this hook:
+/// 1. Sends the configured token to the server
+/// 2. Waits for the server's ACK response
+/// 3. Returns an error if the server rejects the token
+///
+/// # Example
+///
+/// ```rust,ignore
+/// use msg_socket::hooks::token::ClientHook;
+/// use bytes::Bytes;
+///
+/// let hook = ClientHook::new(Bytes::from("my_secret_token"));
+/// ```
+pub struct ClientHook {
+    token: Bytes,
+}
+
+impl ClientHook {
+    /// Creates a new client hook with the given authentication token.
+    pub fn new(token: Bytes) -> Self {
+        Self { token }
+    }
+}
+
+impl<Io> ConnectionHook<Io> for ClientHook
+where
+    Io: AsyncRead + AsyncWrite + Send + Unpin + 'static,
+{
+    async fn on_connection(&self, io: Io) -> Result<Io, HookError> {
+        let mut conn = Framed::new(io, auth::Codec::new_client());
+
+        // Send authentication token
+        conn.send(auth::Message::Auth(self.token.clone())).await?;
+
+        conn.flush().await?;
+
+        // Wait for server acknowledgment
+        let ack = conn
+            .next()
+            .await
+            .ok_or_else(|| HookError::message("connection closed"))?
+            .map_err(HookError::custom)?;
+
+        if !matches!(ack, auth::Message::Ack) {
+            return Err(HookError::message("authentication denied"));
+        }
+
+        Ok(conn.into_inner())
+    }
+}

msg-socket/src/hooks/token.rs

@@ -10,13 +10,15 @@
 #![doc(issue_tracker_base_url = "https://github.com/chainbound/msg-rs/issues/")]
 #![cfg_attr(docsrs, feature(doc_cfg, doc_auto_cfg))]
 #![cfg_attr(not(test), warn(unused_crate_dependencies))]
-use bytes::Bytes;
-use tokio::io::{AsyncRead, AsyncWrite};
-
-use msg_transport::Address;
 
 pub mod stats;
 
+pub mod hooks;
+pub use hooks::{ConnectionHook, HookError};
+
+// Re-export type-erased hook and HookResult for internal use
+pub(crate) use hooks::{ConnectionHookErased, HookResult};
+
 #[path = "pub/mod.rs"]
 mod pubs;
 pub use pubs::{PubError, PubOptions, PubSocket};
@@ -53,18 +55,6 @@ impl RequestId {
     }
 }
 
-/// An interface for authenticating clients, given their ID.
-pub trait Authenticator: Send + Sync + Unpin + 'static {
-    fn authenticate(&self, id: &Bytes) -> bool;
-}
-
-/// The result of an authentication attempt.
-pub(crate) struct AuthResult<S: AsyncRead + AsyncWrite, A: Address> {
-    id: Bytes,
-    addr: A,
-    stream: S,
-}
-
 /// The performance profile to tune socket options for.
 #[derive(Debug, Clone, Default, Copy, PartialEq, Eq)]
 pub enum Profile {