refactor(auth): connection hooks by Karrq · Pull Request #177 · chainbound/msg-rs

Karrq · 2026-01-26T10:19:07Z

Closes #172 with a new ConnectionHook trait, which allows users of the library to inject custom custom logic when a connection is established.

The authentication system we implemented before was redone in form of a hook, and is provided in it's own module

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

claude · 2026-01-26T10:25:41Z

Code Review

I found 2 issues during the code review:

1. Compilation Error: Field name mismatch in SubOptions::default()

File: msg-socket/src/sub/mod.rs line 104

The Default implementation uses ingress_buffer_size but the struct field is named ingress_queue_size. This will cause a compilation error.

See the struct definition around line 56:

pub struct SubOptions {
    /// The maximum amount of incoming messages that will be buffered before being dropped due to
    /// a slow consumer.
    ingress_queue_size: usize,

But the Default impl at line 104 uses the wrong name:

impl Default for SubOptions {
    fn default() -> Self {
        Self {
            ingress_buffer_size: DEFAULT_QUEUE_SIZE,  // Should be ingress_queue_size

Fix: Change ingress_buffer_size to ingress_queue_size in the Default implementation.

2. CLAUDE.md not updated to reflect architecture changes

File: CLAUDE.md lines 53 and 96

The PR completely replaces the Authenticator trait with the new ConnectionHook trait, but CLAUDE.md still references the old Authenticator trait in two places:

Line 53: "Trait-based extensibility - Transport, Authenticator traits for pluggability"
Line 96: "Authentication System - Implement Authenticator trait for custom auth logic"

Since CLAUDE.md provides guidance about the architecture and key design patterns of the project, these references should be updated to reflect the new ConnectionHook trait system.

Suggested changes:

Line 53: Update to mention "ConnectionHook" instead of "Authenticator"
Line 96: Update the Authentication System section to describe the new token-based connection hooks (ServerHook/ClientHook) and the general ConnectionHook trait

See:

msg-rs/CLAUDE.md

Line 53 in f14e578

    
           - **Trait-based extensibility** - Transport, Authenticator traits for pluggability

and

msg-rs/CLAUDE.md

Line 96 in f14e578

- Implement `Authenticator` trait for custom auth logic

Karrq · 2026-01-26T14:03:20Z

@claude please review again

claude · 2026-01-26T14:03:37Z

Claude finished @Karrq's task in 1m 36s —— View job

Code Review Complete

Great work on this refactoring! The previous issues have been fixed:

✅ Fixed: SubOptions::default() now correctly uses ingress_queue_size (msg-socket/src/sub/mod.rs:104)
✅ Fixed: CLAUDE.md has been updated to reference ConnectionHook instead of Authenticator (lines 53, 96-97)

Overall Assessment

This is a well-executed refactoring that successfully replaces the authentication system with a more flexible ConnectionHook trait. The design is clean, type-safe, and provides good extensibility for users.

What's Good

Strong trait design - The ConnectionHook trait is well-designed with clear semantics:
- Takes ownership of the IO stream and returns it (or error to reject)
- Works for both client and server sockets with appropriate timing
- Properly uses async traits with Send bounds
Type erasure pattern - The internal ConnectionHookErased trait allows storing different hook implementations behind a single Arc<dyn> while preserving type safety at the socket level.
Error handling - HookError provides good flexibility with both IO errors and custom errors, while implementing proper std::error::Error traits.
Token auth implementation - The built-in token hooks (ServerHook/ClientHook) provide a good reference implementation and cover the common use case.
Consistent integration - All four socket types (Pub, Sub, Req, Rep) integrate hooks consistently via the builder pattern with_connection_hook().
Hook execution flow:
- Server sockets (Pub, Rep): Hooks run after accepting connections, spawned as background tasks
- Client sockets (Req, Sub): Hooks run after connecting, integrated into connection managers
- All hooks properly handle errors and update connection state
Documentation - Good module-level docs with examples in hooks/mod.rs and hooks/token.rs.

Minor Observations

These are not bugs, just observations for potential future improvements:

Hook task management (msg-socket/src/pub/driver.rs:35, msg-socket/src/rep/driver.rs:151):
- Both Pub and Rep drivers use JoinSet<Result<HookResult<T::Io, A>, HookError>> to spawn hook tasks
- This is correct and works well, though failed hooks properly decrement the active client counter
Dynamic dispatch for hooks (all socket types):
- Using Arc<dyn ConnectionHookErased<T::Io>> is appropriate here given the use case
- The comment in PubSocket (line 41-42) acknowledges this tradeoff for compressors too
- For most authentication/handshake operations, this overhead is negligible compared to network latency
Error message clarity (msg-socket/src/hooks/token.rs:90, 94):
- "connection closed" and "expected auth message" are clear enough
- Could optionally include more context (e.g., peer address) but not necessary
API consistency:
- All sockets consistently use with_connection_hook<H>(hook: H) builder pattern
- The hook is taken from the socket and moved to the driver during bind/connect
- Good separation of concerns

Test Coverage

The examples demonstrate the hooks working correctly:

reqrep_auth.rs - Shows basic usage with ServerHook::accept_all() and ClientHook
pubsub_auth.rs - More complex example with custom validation logic
Both examples are well-commented

Consider adding integration tests in msg-socket/tests/ to verify:

Hook rejection behavior (what happens when auth fails)
Multiple concurrent connections with hooks
Hook errors are properly propagated

Code Quality

Architecture: Clean separation between the public ConnectionHook trait and internal ConnectionHookErased trait
Type safety: Proper use of generics and trait bounds throughout
Error handling: Comprehensive error types with proper conversions
Async design: Proper use of Future trait and Pin<Box<dyn Future>> for type erasure
Documentation: Good inline docs and examples

Summary

This PR successfully achieves its goal of replacing the authentication system with a more flexible hook-based approach. The implementation is solid, well-tested through examples, and maintains consistency across all socket types. The design allows users to implement custom authentication, handshakes, or protocol negotiation while providing convenient built-in token-based auth.

Recommendation: ✅ Approve with the understanding that the minor observations above are not blockers.

thedevbirb · 2026-01-26T14:15:39Z

msg-socket/src/hooks/mod.rs

@@ -0,0 +1,160 @@
+//! Connection hooks for customizing connection setup.


Nit but if the folder is for connection hooks only, I'd rename it as such, since this might hint there are some other type of hooks that could be implemented. Otherwise, these docs can be adjusted slightly.

Hmmm, maybe I can move it under connection? so it's even more explicit?
tbh I didn't leave out other hook types, but at the same time we can refactor if/when we do that

I made it explicit in the doc that this could be expanded to other hook types

msg-socket/src/hooks/mod.rs

thedevbirb · 2026-01-26T14:21:59Z

msg-socket/src/hooks/mod.rs

+    /// A custom hook-specific error.
+    Custom(Box<dyn StdError + Send + Sync + 'static>),


Isn't it a bit too loose? Have you thought about adding a generic?

Originally I wanted the generic too, but it was a pain to propagate everywhere
We already type erase the hook itself, and tbh we don't really use this error so far.
It's mostly to differentiate it from the underlying io error

Now we still have hooks::Error accepting a generic arg, and then we have the type erased one for internal use

msg-socket/src/hooks/mod.rs

thedevbirb · 2026-01-26T14:33:39Z

msg-socket/src/pub/mod.rs

+        let mut pub_socket =
+            PubSocket::new(Tcp::default()).with_connection_hook(ServerHook::accept_all());


Is it necessary to be specified if it's an accept_all policy? Same for other occurrences here.

This replaces the logic of the test that was there before, which I think was just showing how you'd do auth, but the implementation actually didn't verify anything.
So here we are showing that we could do authentication using a connection hook

msg-socket/src/rep/mod.rs

thedevbirb · 2026-01-26T14:38:37Z

msg-socket/src/sub/driver.rs

+    /// A joinset of connection tasks.
+    pub(super) connection_tasks: JoinMap<A, Result<T::Io, HookError>>,


Why now connection tasks may return HookError? Should this be renamed?

Split the join map so we avoid this

thedevbirb · 2026-01-26T14:43:02Z

msg-socket/src/sub/socket.rs

+        // Take the transport from the existing driver and rebuild with the hook
+        if let Some(mut driver) = self.driver.take() {
+            driver.hook = Some(hook_arc.clone());
+            self.driver = Some(driver);
+        }


I would expect it to make clear that this should be run before the driver is taken and spawn by connecting.

Yes this is what I mentioned in the standup today: there's a bit of unclarity in how we want to handle this pattern, where realistically items like hook (and compression as it stands today) shouldn't be able to be changed after, and in general a socket shouldn't be able to spawn multiple drivers

I decided that this decision should be made in another PR

Added assertions

fix: better hook task spans docs: consistent "connection hook" phrasing refactor(sub): split connection and (old auth) hooks task sets fix: ensure connection hook not set when driver is running

mempirate

Very nice! Could you add a more involved example of like a negotiation?

msg-socket/src/hooks/mod.rs

mempirate · 2026-01-27T10:01:12Z

msg-socket/src/hooks/mod.rs

+    ) -> Pin<Box<dyn Future<Output = ErasedHookResult<Io>> + Send + 'static>>;
+}
+
+impl<T, Io> ConnectionHookErased<Io> for T


Nice pattern! Simplifies implementation for any consumers of this API

docs(hook): remove specific socket type mention

Karrq added 2 commits January 26, 2026 11:14

refactor(auth): connection hooks

2f89a1d

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

chore: fmt

0ff635d

Karrq requested review from mempirate and thedevbirb January 26, 2026 10:19

Karrq force-pushed the refactor/auth branch from 6e74279 to 0ff635d Compare January 26, 2026 10:19

Merge remote-tracking branch 'origin/main' into refactor/auth

f14e578

Karrq added 3 commits January 26, 2026 11:35

fix: typo

e55bbe2

docs: update CLAUDE.md

08feb80

deps: remove unused tokio-util

0941d6a

thedevbirb requested changes Jan 26, 2026

View reviewed changes

Karrq added 3 commits January 27, 2026 10:41

refactor: review

c57c8f1

fix: better hook task spans docs: consistent "connection hook" phrasing refactor(sub): split connection and (old auth) hooks task sets fix: ensure connection hook not set when driver is running

chore: fmt

c7365b3

docs: fix refs in test

245ca62

mempirate reviewed Jan 27, 2026

View reviewed changes

Karrq added 2 commits January 27, 2026 14:03

feat: challenge-response auth example

cfa8331

docs(hook): remove specific socket type mention

Merge branch 'main' into refactor/auth

24be1fe

mempirate approved these changes Jan 27, 2026

View reviewed changes

thedevbirb mentioned this pull request Jan 27, 2026

Disconnection hooks #179

Open

thedevbirb approved these changes Jan 28, 2026

View reviewed changes

chore: fmt

c8dee23

Karrq merged commit f013820 into main Jan 28, 2026
13 checks passed

Karrq deleted the refactor/auth branch January 28, 2026 08:49

		@@ -0,0 +1,160 @@
		//! Connection hooks for customizing connection setup.

		/// A custom hook-specific error.
		Custom(Box<dyn StdError + Send + Sync + 'static>),

		let mut pub_socket =
		PubSocket::new(Tcp::default()).with_connection_hook(ServerHook::accept_all());

		/// A joinset of connection tasks.
		pub(super) connection_tasks: JoinMap<A, Result<T::Io, HookError>>,

Conversation

Karrq commented Jan 26, 2026

Uh oh!

claude bot commented Jan 26, 2026

Code Review

1. Compilation Error: Field name mismatch in SubOptions::default()

2. CLAUDE.md not updated to reflect architecture changes

Uh oh!

Karrq commented Jan 26, 2026

Uh oh!

claude bot commented Jan 26, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Code Review Complete

Overall Assessment

What's Good

Minor Observations

Test Coverage

Code Quality

Summary

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Karrq Jan 26, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

mempirate left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

claude bot commented Jan 26, 2026 •

edited

Loading

Karrq Jan 26, 2026 •

edited

Loading