feat: Add connection extension framework with symmetric client-server hooks

[jiangliu]

Summary

ttrpc currently provides no mechanism for applications to inject connection-level policies such as encryption, authentication, or file-descriptor passthrough. Each RPC message passes through serialization and transport layers with no opportunity for the application to inspect or transform the wire-format bytes, leaving security-critical concerns to be handled out-of-band or not at all.

This PR introduces a symmetric extension framework that exposes 10 injection points across the client and server lifecycle, enabling applications to implement composable security policies (e.g., AES-GCM encryption, Ed25519 authentication) without modifying ttrpc core.

Design Principles

• ttrpc provides mechanism, not policy. No encryption algorithms, identity types, or authorization logic in ttrpc itself — all decisions reside in the application layer.
• Symmetric hooks. Both AcceptHook (server) and ConnectHook (client) share the same HookOutput/HookError types and PayloadTransform trait.
• Zero overhead when unused. No hook → all connections accepted, payload_transform = None → plaintext pass-through.

10 Injection Points

#	Side	Location	Direction	Description
1	server	do_start() accept	—	AcceptHook::on_accept
2	server	handle_request()	inbound	Unary REQUEST transform_inbound
3	server	handle_msg() DATA	route	Streaming DATA routed without transform
4	server	handle_method() response	outbound	Unary RESPONSE transform_outbound
5	client	new_inner() connect	—	ConnectHook::on_connect
6	client	request() send	outbound	Unary REQUEST transform_outbound
7	client	handle_msg() recv	inbound	Unary RESPONSE transform_inbound
8	client	new_stream() send	outbound	Stream-init REQUEST transform_outbound
9	both	StreamSender::send()	outbound	Streaming DATA transform_outbound
10	both	StreamReceiver::recv()	inbound	Streaming DATA transform_inbound

Key design decision: Streaming DATA messages are routed (not transformed) in handle_msg() (#3). The transform is applied exclusively in StreamSender::send() / StreamReceiver::recv() (#9/#10) to avoid double-transform.

streaming_client=false with initial payload

When streaming_client=false and the stream-init REQUEST carries a payload, handle_stream() creates a synthetic DATA message from the already-decrypted REQUEST payload (decrypted at #2). A pub(crate) internal_flags field on GenMessage carries a skip_transform bit so StreamReceiver::recv() (#10) passes it through without re-applying transform_inbound. This works for any PayloadTransform, including asymmetric transforms.

New Public API

// src/extension.rs
pub trait PayloadTransform: Send + Sync + Debug {
    fn transform_inbound(&self, data: Vec<u8>) -> Result<Vec<u8>, String>;
    fn transform_outbound(&self, data: Vec<u8>) -> Result<Vec<u8>, String>;
}

#[cfg(unix)]
pub trait AcceptHook: Send + Sync + Debug {
    fn on_accept(&self, fd: RawFd) -> Result<HookOutput, HookError>;
}

#[cfg(unix)]
pub trait ConnectHook: Send + Sync + Debug {
    fn on_connect(&self, fd: RawFd) -> Result<HookOutput, HookError>;
}

pub struct HookOutput {
    pub data: ConnectionData,
    pub payload_transform: Option<Box<dyn PayloadTransform>>,
}

pub enum HookError { Rejected(String), Timeout, Io(io::Error), Other(String) }

pub struct ConnectionContext { /* Arc-wrapped, immutable after accept */ }

Server registration

Server::new()
    .bind("unix:///run/agent.sock")?
    .set_accept_hook(Box::new(MyAcceptHook))
    .register_service(services)
    .start().await?;

Client registration

let socket = Socket::connect("unix:///run/agent.sock").await?;
let client = Client::with_hook(socket, Box::new(MyConnectHook));

Socket raw_fd Capture

Socket now stores the underlying RawFd (Unix only) so hooks can call getpeername() for peer identity inspection (e.g., vsock CID) and perform bidirectional handshake I/O. Platform-specific From impls (TcpStream, UnixStream, VsockStream) capture the fd via as_raw_fd(); the generic Socket::new() sets raw_fd = None.

Files Changed

File	Change
src/extension.rs	New — framework types, traits, ConnectionContext, 20 unit tests
tests/hook_integration.rs	New — 19 integration tests
src/proto.rs	GenMessage::internal_flags + helper constructors
src/lib.rs	Re-export HookError, HookOutput, AcceptHook, ConnectHook
src/asynchronous/server.rs	set_accept_hook(), on_accept() in accept loop, injection points #2/#3/#4
src/asynchronous/client.rs	with_hook(), ConnectHook in new_inner(), injection points #6/#7/#8
src/asynchronous/stream.rs	payload_transform field on sender/receiver, injection points #9/#10
src/asynchronous/transport/mod.rs	Socket gains raw_fd: Option<RawFd>, Socket::from() captures fd
src/asynchronous/transport/{tcp,unix,vsock}.rs	From impls capture as_raw_fd()
src/asynchronous/utils.rs	TtrpcContext::connection_data field

Test Coverage

87 tests (67 unit + 19 integration + 1 example):

• Payload transform roundtrip (XOR-0xA5A5 mock encryption): basic, empty, odd-length, large, header verification, error handling
• ConnectionContext construction, Arc sharing, transform pass-through
• Hook invocation, rejection, fd capture
• End-to-end unary + streaming with XOR transform
• streaming_client=false with initial payload (skip_transform path)
• streaming_server=false DATA message rejection
• Multiple concurrent connections with transform
• Multiple concurrent streams on one connection
• Server shutdown during active stream
• Unary request timeout

Backward Compatibility

Scenario	Behavior
No hook set	All connections accepted, empty data, no transform
payload_transform = None	Plaintext pass-through (no overhead)
Existing handlers	connection_data is empty HashMap — no breakage
Wire-protocol header	Unchanged — transform operates on payload bytes after framing

[Copilot]

Pull request overview

This PR adds a symmetric connection extension framework to ttrpc’s async client/server stack, enabling applications to attach per-connection metadata and optionally transform payload bytes (e.g., for encryption/auth) via hooks and a PayloadTransform.

Changes:

• Introduces a new extension module with hook/transform traits, connection context, and typed connection metadata propagation.
• Wires the connection context through async server/client/stream pipelines, applying transforms at defined injection points and adding a skip_transform internal flag for synthetic stream DATA.
• Extends async transport Socket to capture Unix raw_fd so hooks can inspect peers and perform handshake I/O.

Reviewed changes

Copilot reviewed 12 out of 12 changed files in this pull request and generated 6 comments.

Show a summary per file

File	Description
tests/hook_integration.rs	Adds end-to-end integration tests for hooks/transforms over Unix sockets.
src/extension.rs	New core extension framework types (hooks, transform trait, connection context) plus unit tests.
src/lib.rs	Exposes the new extension module and re-exports key extension APIs.
src/proto.rs	Adds non-wire internal_flags to GenMessage and helper constructors for DATA/RESPONSE/close.
src/asynchronous/utils.rs	Adds connection_data to TtrpcContext for handler access.
src/asynchronous/server.rs	Installs accept hook, propagates connection context, and applies transforms in server pipeline.
src/asynchronous/client.rs	Adds Client::with_hook, runs connect hook, and applies transforms in client pipeline.
src/asynchronous/stream.rs	Applies transforms for streaming DATA send/recv with skip-transform support.
src/asynchronous/transport/mod.rs	Refactors Socket to store Unix raw_fd and exposes as_raw_fd().
src/asynchronous/transport/unix.rs	Ensures accepted/connected Unix sockets capture raw_fd.
src/asynchronous/transport/tcp.rs	Ensures accepted/connected TCP sockets capture raw_fd.
src/asynchronous/transport/vsock.rs	Ensures accepted/connected vsock sockets capture raw_fd.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Pull request overview

Copilot reviewed 12 out of 12 changed files in this pull request and generated 6 comments.

[Copilot]

Pull request overview

Copilot reviewed 12 out of 12 changed files in this pull request and generated 5 comments.

Comments suppressed due to low confidence (1)

src/asynchronous/server.rs:692

• respond() applies transform_outbound, but it doesn't enforce MESSAGE_LENGTH_MAX on the transformed payload. A transform that expands data can cause oversized responses to be sent (especially for error/status responses that go through respond_with_status). Add check_oversize after the transform.

        let payload = self
            .conn_ctx
            .transform_outbound(payload)
            .map_err(|e| Error::Others(format!("transform_outbound failed: {}", e)))?;
        let msg = GenMessage::new_response(stream_id, payload);

[Copilot]

Pull request overview

Copilot reviewed 12 out of 12 changed files in this pull request and generated 1 comment.

Comments suppressed due to low confidence (1)

src/asynchronous/server.rs:694

• respond() applies transform_outbound but never enforces MESSAGE_LENGTH_MAX. Other outbound paths (e.g. the normal unary response path) check size after transform; without a check here, an oversized (or expansion-prone) transform can cause the peer to reject/discard the message and the caller to hang until timeout.

        let payload = self
            .conn_ctx
            .transform_outbound(payload)
            .map_err(|e| Error::Others(format!("transform_outbound failed: {}", e)))?;
        let msg = GenMessage::new_response(stream_id, payload);

[Copilot]

Pull request overview

Copilot reviewed 12 out of 12 changed files in this pull request and generated 2 comments.

Comments suppressed due to low confidence (1)

src/asynchronous/server.rs:694

• respond() applies transform_outbound but does not enforce MESSAGE_LENGTH_MAX after the transform. A transform (e.g., encryption) can expand the payload beyond the limit, leading to oversized frames being written and then rejected by the peer during read/oversize checks.

        let payload = self
            .conn_ctx
            .transform_outbound(payload)
            .map_err(|e| Error::Others(format!("transform_outbound failed: {}", e)))?;
        let msg = GenMessage::new_response(stream_id, payload);

[Copilot]

Pull request overview

Copilot reviewed 12 out of 12 changed files in this pull request and generated 2 comments.

Comments suppressed due to low confidence (2)

src/asynchronous/server.rs:695

• respond() applies transform_outbound but does not re-check the transformed payload size. A transform can expand data, so this can violate MESSAGE_LENGTH_MAX and produce frames the peer will reject or that can break framing.

        let payload = self
            .conn_ctx
            .transform_outbound(payload)
            .map_err(|e| Error::Others(format!("transform_outbound failed: {}", e)))?;
        let msg = GenMessage::new_response(stream_id, payload);

src/asynchronous/server.rs:455

• When transform_outbound fails, this path tries to send an INTERNAL status via respond_with_status(), which will call transform_outbound again and likely fail the same way—leaving the client waiting until timeout. Consider closing the stream/connection to unblock the client when the transform is unusable.

                                        self.respond_with_status(
                                            stream_id,
                                            get_status(Code::INTERNAL, format!("transform_outbound: {}", e)),
                                        )
                                        .await;

[Copilot]

Pull request overview

Copilot reviewed 12 out of 12 changed files in this pull request and generated 1 comment.

Comments suppressed due to low confidence (1)

src/asynchronous/server.rs:695

• respond() applies transform_outbound, but it does not enforce MESSAGE_LENGTH_MAX (or call GenMessage::check()) after the transform. A transform can expand payload size (e.g., AEAD tag, framing, compression expansion), which would cause the peer to reject/discard the message and potentially hang waiting for a response. Add a post-transform size check before sending.

        let payload = self
            .conn_ctx
            .transform_outbound(payload)
            .map_err(|e| Error::Others(format!("transform_outbound failed: {}", e)))?;
        let msg = GenMessage::new_response(stream_id, payload);

[Copilot]

Pull request overview

Copilot reviewed 23 out of 23 changed files in this pull request and generated 1 comment.

Suppressed comments (1)

src/asynchronous/server.rs:189

• The accept loop awaits server_ext.on_accept(&conn).await inline, so a slow/blocked hook will pause incoming.next() polling and prevent the server from accepting additional connections until the hook completes (or times out). This contradicts the stated goal that hooks should not stall acceptance, and can materially reduce connection throughput under load. Consider spawning a per-connection task that runs the hook and then calls spawn_connection_handler, so the main accept loop can continue accepting new connections immediately.

                                Ok(conn) => {
                                    // ── Injection Point 1/10: accept hook ──
                                    #[cfg(feature = "security_extension")]
                                    let conn_ctx = match server_ext.on_accept(&conn).await {
                                        Ok(output) => Arc::new(ConnectionContext::new(output)),

[Copilot]

Pull request overview

Copilot reviewed 23 out of 23 changed files in this pull request and generated no new comments.

Suppressed comments (1)

src/asynchronous/stream.rs:35

• The StreamMsg::Wire doc comment says transform_inbound has not been applied yet, but ClientReader::handle_msg() decrypts non-DATA messages before sending them into the per-stream channel as StreamMsg::Wire. This makes the comment misleading and can cause future changes to accidentally double-transform (or skip a needed transform) based on incorrect assumptions.

pub(crate) enum StreamMsg {
    /// Wire-format message — `transform_inbound` has NOT been applied to the
    /// payload yet. `StreamReceiver::recv()` will apply it.
    Wire(GenMessage),

[Copilot]

Pull request overview

Copilot reviewed 23 out of 23 changed files in this pull request and generated 5 comments.

Suppressed comments (2)

src/asynchronous/utils.rs:264

• Adding this field to a publicly constructible context is a source-breaking change: downstream tests or adapters using a TtrpcContext { ... } literal now fail even when security_extension is disabled. Deriving Default only helps newly updated literals; it does not preserve existing callers. The connection data should be exposed without changing the required fields of the existing public struct if backward compatibility is required.

    /// Opaque per-connection data from [`AcceptHook`](crate::security_extension::AcceptHook). Immutable after accept.
    /// See [`ConnectionData`](crate::security_extension::ConnectionData) for full contract.
    /// Empty when `security_extension` is not enabled.
    pub connection_data: Arc<ConnectionData>,

src/security_extension.rs:429

• The framework timeout is represented as HookError::Other, even though the public error API has a dedicated HookError::Timeout variant. This prevents logs/monitoring from distinguishing timeouts structurally as promised by the API.

                    Err(HookError::Other(format!(
                        "accept hook timed out after {:?}",
                        ACCEPT_HOOK_TIMEOUT
                    )))

[Copilot]

Pull request overview

Copilot reviewed 23 out of 23 changed files in this pull request and generated 1 comment.

Suppressed comments (5)

src/asynchronous/stream.rs:579

• Unlike send(), this transformed close-frame path never performs the post-transform size check. A transform that expands an empty payload beyond MESSAGE_LENGTH_MAX is therefore queued with an oversized frame instead of returning an error. Check the message after transformation before sending it.

        let mut msg = GenMessage::new_close(self.stream_id);
        // ── Close-message transform: authenticate header via AAD ──
        self.transform_msg(&mut msg)?;
        _send(&self.tx, msg).await?;

src/security_extension.rs:330

• This runtime requirement is incorrect: Tokio supports spawn_blocking on a current-thread runtime and executes the closure on an additional blocking thread. Keeping this warning unnecessarily excludes a supported deployment and contradicts the preceding statement that current-thread runtimes are handled safely.

    /// **Requirement**: the async server must run on a **multi-thread** tokio
    /// runtime; `spawn_blocking` panics on current-thread runtimes.

src/asynchronous/transport/mod.rs:118

• The captured descriptor belongs to a Tokio stream and remains nonblocking. Consequently, the synchronous hooks' documented handshake reads can return WouldBlock before peer data arrives, and SO_RCVTIMEO does not make I/O wait on an O_NONBLOCK descriptor. Provide readiness-aware/async hook I/O, or put the descriptor into blocking mode exclusively for the offloaded hook and restore nonblocking mode before starting Tokio I/O; a multi-step handshake integration test should exercise this race.

            let fd = socket.as_raw_fd();
            Self::with_raw_fd(socket, fd)

src/security_extension.rs:429

• The framework-generated timeout is returned as HookError::Other, so callers cannot distinguish it using the dedicated HookError::Timeout variant promised by this API. Return the structured timeout variant here.

                    Err(HookError::Other(format!(
                        "accept hook timed out after {:?}",
                        ACCEPT_HOOK_TIMEOUT
                    )))

src/asynchronous/server.rs:462

• The changed Rust sources are not rustfmt-clean (this spacing is one example, along with long unformatted signatures in the new transport and extension code). The check CI job runs cargo fmt --all -- --check, so this PR will fail that required job until cargo fmt --all is applied.

                Err(status) => self.respond_with_status( stream_id, status).await,

[Tim-Zhang]

Thanks for the update. Most of the earlier comments are now addressed.

I found three remaining correctness issues in the new hook and transform paths. I also added a few smaller comments about tests, documentation, and the no-feature path.

Please also update the PR description. It still shows the old boxed hook API, the old PayloadTransform methods without aad, and the removed internal_flags design.

[Copilot]

Pull request overview

Copilot reviewed 23 out of 23 changed files in this pull request and generated 4 comments.

Suppressed comments (3)

src/security_extension.rs:457

• The timeout does not bound hook resource usage once the blocking closure has started. The closure owns a duplicated descriptor, so dropping the original Socket does not interrupt a hook blocked on I/O, and spawn_blocking tasks cannot be aborted after starting. Repeated stalled handshakes can therefore retain one fd and blocking-pool thread each indefinitely. The timeout path needs to actively shut down the shared socket/cooperatively cancel the hook so the duplicate's I/O wakes before returning.

                _ = tokio::time::sleep(ACCEPT_HOOK_TIMEOUT) => {
                    // Cancel the blocking task. For queued tasks this prevents
                    // execution entirely. For already-running tasks, abort is
                    // best-effort (spawn_blocking cannot be forcibly killed);
                    // the hook's dup'd fd ensures no fd reuse occurs even if
                    // the task continues briefly after abort.
                    handle.abort();

src/sync/server.rs:468

• The synchronous hook runs directly on the sole listener_loop thread before the per-connection worker is spawned. A hook performing a stalled handshake prevents the server from accepting every subsequent connection, with no framework timeout on this path. Dispatch hook/connection setup to a per-connection worker or enforce a deadline without blocking the accept loop.

                    let conn_ctx = match connection_context_from_hook(
                        pipe_connection.id(),
                        accept_hook.as_ref(),
                    ) {

src/asynchronous/transport/mod.rs:96

• This declaration exceeds rustfmt's configured/default width and cargo fmt --all -- --check in the changed Makefile will reformat it, causing CI to fail until the generated formatting is committed.

    pub(crate) fn with_raw_fd(socket: impl AsyncRead + AsyncWrite + Send + Sync + 'static, fd: RawFd) -> Self {

[Copilot]

Pull request overview

Copilot reviewed 23 out of 23 changed files in this pull request and generated 2 comments.

Suppressed comments (4)

src/asynchronous/server.rs:464

• This branch says the connection will close, but it only logs the error. If respond() fails in transform_outbound, a timeout-free client request remains pending indefinitely on an otherwise open connection. Treat an outbound transform failure as connection-fatal (or otherwise notify all pending requests).

                        if let Err(e) = self.respond(stream_id, resp).await {
                            // respond() handles oversize internally via pre-check.
                            // Remaining failures (channel closed, transform error)
                            // are not recoverable — log and let connection close.
                            error!("respond failed for stream {}: {}", stream_id, e);
                        }

src/security_extension.rs:348

• This requirement is incorrect: Tokio supports spawn_blocking on both current-thread and multi-thread runtimes, using additional blocking threads in either case. The preceding paragraph correctly says current-thread runtimes are supported, so this contradictory warning should be removed.

    /// **Requirement**: the async server must run on a **multi-thread** tokio
    /// runtime; `spawn_blocking` panics on current-thread runtimes.

src/security_extension.rs:311

• This public API documentation describes the old pipeline and contradicts the implementation and the module-level table: inbound DATA is now transformed in the connection reader before routing, not in StreamReceiver::recv(). Update this propagation diagram so extension implementers do not design around a nonexistent receiver-side transform.

    /// **Important**: Streaming DATA messages are routed to the stream channel
    /// without transform in `handle_msg()` (#3). The transform is applied later
    /// in `StreamReceiver::recv()` (#10) / `StreamSender::send()` (#9) to avoid
    /// double-transform.

tests/hook_integration_async_unix.rs:367

• This roundtrip assertion would still pass if the framework skipped both outbound and inbound transforms, so it does not verify the behavior named transform_on_wire. Instrument the transform with invocation counters or inspect bytes through a transport proxy so the integration test fails when either injection point is bypassed.

    // Send request — should be XOR-encrypted on wire, decrypted by server
    let original_payload = b"encrypted_message_content";
    let req = build_echo_request(original_payload);
    let resp = client.request(req).await.unwrap();

[Copilot]

Pull request overview

Copilot reviewed 23 out of 23 changed files in this pull request and generated 2 comments.

Suppressed comments (4)

src/sync/utils.rs:89

• This compatibility wrapper deliberately replaces the real connection context with a plaintext context. Existing hand-written handlers can therefore silently leak responses on a transformed connection, despite the framework's guarantee that no payload bypasses the transform. The legacy path must be bound to the connection context or made unavailable when security extensions are active; deprecation alone does not prevent the runtime downgrade.

    let ctx = ConnectionContext::default();
    send_response(stream_id, res, tx, &ctx)

src/security_extension.rs:456

• After timeout, aborting an already-running spawn_blocking task does not stop it, and the duplicated descriptor keeps the socket open after the original Socket is dropped. A hook blocked in I/O can therefore retain one blocking-pool thread and one connection indefinitely. Add cooperative cancellation or shut down the socket on timeout so blocked hook I/O is actually interrupted before returning.

                _ = tokio::time::sleep(ACCEPT_HOOK_TIMEOUT) => {
                    // Cancel the blocking task. For queued tasks this prevents
                    // execution entirely. For already-running tasks, abort is
                    // best-effort (spawn_blocking cannot be forcibly killed);
                    // the hook's dup'd fd ensures no fd reuse occurs even if
                    // the task continues briefly after abort.
                    handle.abort();

src/sync/server.rs:55

• The sync accept hook runs directly in the sole listener loop. Since hooks are intended to perform handshake I/O, one slow or malicious peer can block all subsequent accepts indefinitely. Dispatch the hook/connection setup to a per-connection worker and enforce a timeout, as the async path does.

        Some(h) => h.on_accept(fd).map(|o| ConnectionContext::new(Some(o))),

src/sync/server.rs:543

• After decryption, x.length still contains the ciphertext length. This violates the MessageHeader payload-length invariant and exposes an incorrect length through TtrpcContext::mh whenever a transform changes size; the async pipeline already updates it. Set the header length from the transformed buffer before queueing the workload.

                                                Ok(data) => reader_ctx.inbound_buf(
                                                    data,
                                                    &serialize_aad(&x),
                                                    true,
                                                ),

[Copilot]

Pull request overview

Copilot reviewed 24 out of 24 changed files in this pull request and generated 6 comments.

Suppressed comments (5)

src/security_extension.rs:456

• JoinHandle::abort() cannot stop an already-running spawn_blocking hook, and the duplicated OwnedFd now keeps the connection alive after the outer Socket is dropped. A hook blocked on this descriptor can therefore retain both an fd and a blocking-pool thread indefinitely; repeated unauthenticated connections can exhaust those resources despite the 30-second timeout. The timeout needs cooperative cancellation or enforced I/O deadlines on the owned descriptor.

                _ = tokio::time::sleep(ACCEPT_HOOK_TIMEOUT) => {
                    // Cancel the blocking task. For queued tasks this prevents
                    // execution entirely. For already-running tasks, abort is
                    // best-effort (spawn_blocking cannot be forcibly killed);
                    // the hook's dup'd fd ensures no fd reuse occurs even if
                    // the task continues briefly after abort.
                    handle.abort();

src/asynchronous/server.rs:383

• Only DATA is transformed in the sequential reader, while REQUEST is transformed later in the spawned handler (handle_request, line 534). A stream-init REQUEST followed by DATA can therefore consume the DATA nonce first, and concurrent unary requests can also transform out of wire order, causing valid traffic to fail for stateful transforms such as counter-based AEAD. Apply inbound transforms to REQUEST and DATA before spawning, then remove the later REQUEST transform.

            let msg = if msg.header.type_ == MESSAGE_TYPE_DATA {
                let mut msg = msg;
                if let Err(e) = self.conn_ctx.inbound(&mut msg, false) {

src/asynchronous/server.rs:396

• On transform failure the stream ID has not been authenticated, so routing the error with msg.header.stream_id can fail an unrelated active stream or silently drop the error. Authentication failure should terminate the connection (and wake all stream handlers) rather than continue using header fields whose AAD check failed.

                    // Route the error to the stream handler so recv() returns
                    // immediately instead of waiting forever.
                    let stream_id = msg.header.stream_id;
                    let stream_tx = {
                        let guard = self.streams.lock().unwrap();
                        guard.get(&stream_id).cloned()
                    };
                    if let Some(tx) = stream_tx {
                        let _ = tx.send(Err(e)).await;

src/asynchronous/client.rs:161

• The stream entry is inserted before transformation, but this early ? leaves it in the map when the transform rejects a payload (for example, an oversized request). Because the connection can remain usable, repeated failures leak one sender per request. Remove the entry before returning the transform error.

        // ── Injection Point 6/10: unary REQUEST transform_outbound ──
        self.conn_ctx.transform_send(&mut msg, &self.req_tx, false).await?;

src/asynchronous/client.rs:221

• If transformation or enqueueing fails, this returns while the newly inserted stream sender remains in self.streams; no response can arrive to remove it. Failed stream creations can therefore accumulate stale entries. Remove stream_id from the map on this error path.

        // ── Injection Point 8/10: stream-init REQUEST transform_outbound ──
        self.conn_ctx.transform_send(&mut msg, &self.req_tx, false).await?;

[Copilot]

Pull request overview

Copilot reviewed 24 out of 24 changed files in this pull request and generated no new comments.

Suppressed comments (10)

Cargo.toml:69

• The feature can be enabled without async, but the feature-enabled ConnectionContext unconditionally uses tokio::sync::Mutex and defines transform_send against crate::asynchronous. Consequently, the supported sync-only command cargo check --features security_extension cannot compile. Gate the async-only field, initialization, and method with feature = "async" (as the no-op implementation already does), or make this feature explicitly enable async.

security_extension = []

src/asynchronous/client.rs:221

• A failed transform/enqueue leaves this newly inserted stream sender in self.streams, while new_stream returns no receiver whose drop could clean it up. Remove the entry before returning the send error to avoid a permanent per-failure leak.

        // ── Injection Point 8/10: stream-init REQUEST transform_outbound ──
        self.conn_ctx.transform_send(&mut msg, &self.req_tx, false).await?;

tests/common/mod.rs:228

• Use the same real close-frame type and flags as aad_close here; otherwise the rejection can be caused by changing type/flags rather than demonstrating that changing only stream_id is authenticated.

        // But if stream_id is changed, the close is rejected
        let aad_wrong = make_aad(99, 0x04, 0x00);

src/asynchronous/stream.rs:471

• This replaces the previous _send path, which used SendingMessage::new_with_result and awaited the writer's result. send() now returns success as soon as the frame is queued, even if the actual socket write fails, regressing the public stream API's error reporting. Preserve the writer acknowledgement while keeping transform-and-enqueue ordering serialized.

        let mut msg = GenMessage::new_data(self.stream_id, buf);
        // ── Injection Point 9/10: streaming DATA transform_outbound ──
        self.conn_ctx.transform_send(&mut msg, &self.tx, false).await?;

src/asynchronous/stream.rs:484

• Like send(), close_send() previously awaited the writer acknowledgement. It now marks the stream locally closed after enqueueing even when the close frame later fails to reach the transport, and callers cannot retry. Use a transform/send helper that returns the actual writer result before setting local_closed.

        let mut msg = GenMessage::new_close(self.stream_id);
        self.conn_ctx.transform_send(&mut msg, &self.tx, false).await?;
        self.local_closed.store(true, Ordering::Relaxed);

src/asynchronous/client.rs:161

• The stream entry is inserted before transformation, but an outbound transform or queue failure returns here without removing it. Since no request was sent, no response can remove the entry, so repeated transform failures permanently grow self.streams. Remove stream_id before propagating the error.

This issue also appears on line 220 of the same file.

        // ── Injection Point 6/10: unary REQUEST transform_outbound ──
        self.conn_ctx.transform_send(&mut msg, &self.req_tx, false).await?;

tests/hook_integration_async_unix.rs:1175

• This test accepts every possible outcome: the expected error, a successful receive, and a timeout. It therefore passes even if streaming_server=false stops rejecting DATA entirely. Require the receive to finish with the documented non-streaming-server error.

        Ok(Ok(_)) => {
            // If recv succeeded, the server may have returned a RESPONSE instead of DATA.
            // This can happen if the handler returns Some(Response) before echoing.
            // Acceptable — the test still exercises the streaming_server=false path.
        }
        Err(_) => {
            // Timeout: the server handler is still running in its echo loop.
            // This is expected since the server waits for more data or close signal.
        }

tests/common/mod.rs:214

• This does not model an actual close frame: ttrpc close frames use message type DATA (0x03) with FLAG_REMOTE_CLOSED | FLAG_NO_DATA (0x05), not a nonexistent CLOSE type 0x04 with zero flags. As written, the test does not exercise the AAD emitted by GenMessage::new_close.

This issue also appears on line 227 of the same file.

    fn detect_close_frame_aad() {
        // Close messages also carry AAD (stream_id || type_=CLOSE || flags=0).
        let xform = XorPayloadTransform;
        let aad_close = make_aad(7, 0x04, 0x00); // CLOSE type

src/asynchronous/server.rs:143

• The registration examples in the PR pass Box::new(MyAcceptHook), but this generic bound requires the argument type itself to implement AcceptHook; there is no forwarding implementation for Box<T>, so those examples do not compile. Either advertise passing the concrete hook directly or add boxed-trait forwarding/support consistently for the client and sync APIs too.

    pub fn set_accept_hook<H: AcceptHook + 'static>(mut self, hook: H) -> Self {
        let hook: Arc<dyn AcceptHook> = Arc::new(hook);
        self.accept_hook = Some(hook);

src/security_extension.rs:143

• The PR's “New Public API” snippet defines both transform methods with only data, while the actual public trait requires the additional aad argument (and also introduces max_overhead). Implementations copied from the advertised API will not compile. Update the PR/API documentation to match this signature, or keep the documented signature.

    fn transform_inbound(&self, data: Vec<u8>, aad: &[u8]) -> std::result::Result<Vec<u8>, String>;
    /// Encrypt / encode a payload. `aad` has the same semantics as
    /// [`transform_inbound`](Self::transform_inbound).
    fn transform_outbound(&self, data: Vec<u8>, aad: &[u8])
        -> std::result::Result<Vec<u8>, String>;

[Copilot]

Pull request overview

Copilot reviewed 24 out of 24 changed files in this pull request and generated 2 comments.

Suppressed comments (8)

src/security_extension.rs:481

• ConnectionContext cannot derive Debug: ConnectionData contains Box<dyn Any + Send + Sync>, and dyn Any does not implement Debug. Builds with security_extension therefore fail here (and TtrpcContext also requires this type to be Debug). Provide a manual implementation that reports only debuggable metadata.

    #[derive(Debug, Default)]

src/asynchronous/client.rs:221

• As in request, the stream-map entry is created before this fallible send and is never removed when transformation or writing fails. Since no StreamReceiver is constructed on that path, its Drop cleanup cannot run; explicitly remove stream_id before returning the error.

        // ── Injection Point 8/10: stream-init REQUEST transform_outbound ──
        self.conn_ctx.transform_send(&mut msg, &self.req_tx, false).await?;

src/asynchronous/stream.rs:388

• This changes the existing public StreamInner::new constructor to crate-private while also changing its signature, so downstream code using the exported StreamInner no longer compiles. Preserve the old public constructor with a default context and add a separate internal context-aware constructor for framework call sites.

    #[allow(clippy::too_many_arguments)]
    pub(crate) fn new(

src/asynchronous/server.rs:688

• For payloads that fit MESSAGE_LENGTH_MAX but exceed the transform-safe limit, this creates Error::Others, which converts to a Code::UNKNOWN response. Oversize input is otherwise reported as INVALID_ARGUMENT; use an RPC-status error here as well so enabling a transform does not change the status category.

            let err_resp: Response = Error::Others(err_msg).into();

src/security_extension.rs:468

• Aborting a running spawn_blocking task does not stop it, and the task owns the duplicated socket fd. A hook blocked in I/O can therefore survive the 30-second timeout indefinitely; repeated connections can retain unbounded blocking tasks and descriptors. The timeout path must actively interrupt/close the hook's I/O or use a cancellable execution model rather than relying on abort().

                    handle.abort();

src/asynchronous/client.rs:161

• This entry is inserted into streams before transform_send, but an outbound-transform or enqueue/write failure returns immediately without removing it. Repeated transform failures therefore retain dead senders for the lifetime of the client; remove the stream entry on this error path.

This issue also appears on line 220 of the same file.

        // ── Injection Point 6/10: unary REQUEST transform_outbound ──
        self.conn_ctx.transform_send(&mut msg, &self.req_tx, false).await?;

src/asynchronous/client.rs:161

• transform_send waits for the writer's actual socket-write result, but the request timeout is only started afterward. If the peer stops reading and the write blocks, a request with a finite timeout_nano can now hang indefinitely. Apply one deadline around both sending and receiving, or restore enqueue-only behavior for unary requests.

        // ── Injection Point 6/10: unary REQUEST transform_outbound ──
        self.conn_ctx.transform_send(&mut msg, &self.req_tx, false).await?;

src/sync/utils.rs:46

• The old oversize path converted check_oversize(..., true) into an INVALID_ARGUMENT response. Converting Error::Others instead changes the wire status to UNKNOWN, so plain connections now observe a regression for oversized responses. Build the fallback from an INVALID_ARGUMENT RPC status.

        let err_resp: Response = Error::Others(err_msg).into();

[Copilot]

Pull request overview

Copilot reviewed 24 out of 24 changed files in this pull request and generated 1 comment.

Suppressed comments (10)

src/asynchronous/server.rs:536

• The REQUEST transform runs only after handle_msg has validated the unauthenticated stream_id at line 453. An attacker who changes an authenticated odd stream ID to an even one bypasses transform_inbound entirely and gets a protocol error instead of an AAD failure; counter-based transforms can also become permanently desynchronized. Transform every inbound frame in the sequential reader before any header validation/routing, then avoid transforming REQUEST a second time here.

        // ── Injection Point 2/10: unary REQUEST transform_inbound ──
        let mut msg = msg;
        self.conn_ctx
            .inbound(&mut msg, true)
            .map_err(|e| get_status(Code::INVALID_ARGUMENT, format!("{}", e)))?;

src/asynchronous/server.rs:475

• A response transform failure is only logged, leaving both the connection and the client's pending request open. Requests with timeout_nano == 0 then wait forever because no response or disconnect reaches their channel. Treat an outbound transform failure as connection-fatal and notify/close pending streams rather than continuing the connection.

                        if let Err(e) = self.respond(stream_id, resp).await {
                            // respond() handles oversize internally via pre-check.
                            // Remaining failures (channel closed, transform error)
                            // are not recoverable per-request. The connection
                            // stays open until the next read error or client
                            // disconnect tears it down.
                            error!("respond failed for stream {}: {}", stream_id, e);

src/security_extension.rs:469

• JoinHandle::abort() cannot stop a spawn_blocking closure once it has started. Because the duplicated fd is owned by that closure, a hook blocked forever retains both a blocking-pool thread and the socket forever; repeated connections can exhaust threads and descriptors despite the advertised timeout. The timeout path needs cooperative cancellation that interrupts hook I/O and waits for cleanup, or a hook execution model that can actually be terminated.

                _ = tokio::time::sleep(ACCEPT_HOOK_TIMEOUT) => {
                    // Cancel the blocking task. For queued tasks this prevents
                    // execution entirely. For already-running tasks, abort is
                    // best-effort (spawn_blocking cannot be forcibly killed);
                    // the hook's dup'd fd ensures no fd reuse occurs even if
                    // the task continues briefly after abort.
                    handle.abort();
                    Err(HookError::Timeout)

src/asynchronous/stream.rs:402

• This constructor was previously public, and StreamInner is publicly re-exported. Reducing it to pub(crate) is a source-breaking API change for downstream stream adapters, contrary to the PR's backward-compatibility claim. Preserve the old public constructor (using a default context) and add a separate internal context-aware constructor, or explicitly version this as a breaking release.

impl StreamInner {
    #[allow(clippy::too_many_arguments)]
    pub(crate) fn new(
        stream_id: u32,

src/sync/utils.rs:180

• Adding an unconditional field to this public struct breaks every downstream TtrpcContext { ... } literal, including builds that do not enable security_extension. Gate the storage behind the opt-in feature while preserving the old default-feature layout/API, or treat this as a breaking release with migration guidance.

    /// Per-connection extension context (opaque data + optional payload transform).
    /// Immutable after accept. Default (empty data, no transform) when no hook is configured.
    pub conn_ctx: Arc<ConnectionContext>,

src/asynchronous/client.rs:161

• If transformation or enqueueing fails, this returns while the sender inserted at lines 155-158 remains in streams. No response can remove it, so repeated transform failures grow the map indefinitely.

        // ── Injection Point 6/10: unary REQUEST transform_outbound ──
        self.conn_ctx.transform_send(&mut msg, &self.req_tx, false, false).await?;

src/asynchronous/client.rs:221

• If transforming or enqueueing the stream-init frame fails, this returns before constructing StreamReceiver, whose Drop normally removes the map entry. The orphaned sender therefore remains in streams permanently.

        // ── Injection Point 8/10: stream-init REQUEST transform_outbound ──
        self.conn_ctx.transform_send(&mut msg, &self.req_tx, false, false).await?;

tests/hook_integration_async_unix.rs:1062

• This test is named and documented to verify delivery of the handler's final response, but it accepts EOF, closure, or decode failure as success. A regression that drops every final RESPONSE therefore passes. Require the second receive to return limited_done; closure before that response is precisely the failure this test should catch.

        Ok(Err(e)) => {
            // Server may close the stream before sending final response — that's acceptable
            let err_str = format!("{}", e);
            assert!(
                err_str.contains("Eof")
                    || err_str.contains("RemoteClosed")
                    || err_str.contains("closed")
                    || err_str.contains("Receiver")
                    || err_str.contains("Decode"),
                "Expected Eof/RemoteClosed/Decode, got: {}",
                err_str
            );
        }

src/asynchronous/server.rs:383

• The PR description explicitly says inbound DATA is routed unchanged in handle_msg and transformed in StreamReceiver::recv, with an internal skip-transform flag for synthetic data. This implementation instead transforms DATA in the reader and makes StreamReceiver pass it through, while GenMessage has no such flag. Update the PR's stated architecture/10 injection points to match this materially different contract, or change the implementation to the described design.

            // ── Inbound DATA transform in wire order ──
            // Apply transform before spawning so that concurrent stream
            // handlers process DATA with deterministic nonce sequencing
            // (e.g., AEAD). The connection reader loop calls handle_msg
            // sequentially, ensuring transforms execute in wire order.
            let msg = if msg.header.type_ == MESSAGE_TYPE_DATA {
                let mut msg = msg;
                if let Err(e) = self.conn_ctx.inbound(&mut msg, false) {

src/security_extension.rs:672

• The ordering lock is specifically justified for stateful nonce/counter transforms, but all integration concurrency tests use the stateless XOR transform, so they cannot detect transform-order versus wire-order races. Add a counter-based transform test with concurrent unary and streaming sends that fails if counters arrive out of sequence.

        /// Atomically: lock outbound → transform → enqueue on the wire channel.
        ///
        /// Serializes transform (which may assign AEAD nonces) and channel
        /// insertion so that wire order always matches nonce order.

[Tim-Zhang]

Reviewing the latest revision with the assumption that PR #318 will be integrated. The unary request cleanup covered by that PR is excluded here.

[Copilot]

Pull request overview

Copilot reviewed 24 out of 24 changed files in this pull request and generated 1 comment.

Suppressed comments (12)

src/asynchronous/client.rs:414

• get_resp_tx is called even when inbound returned an authentication error, so it uses the unauthenticated type, flags, and stream ID and may remove an arbitrary request entry before delivering the error. On transform failure, terminate the connection and fail all pending entries without routing through msg.header.

        tokio::spawn(async move {
            if let Some(resp_tx) = get_resp_tx(req_map, &msg.header).await {
                resp_tx
                    .send(result.map(|_| msg))

src/sync/client.rs:198

• A transform error is passed with the unauthenticated mh to trans_resp, which selects an in-flight request by that header's stream ID. A forged header can therefore fail a targeted request. Treat authentication failure as connection-fatal and notify all pending requests without consulting this header.

                            Ok(data) => receiver_ctx.inbound_buf(
                                data,
                                &serialize_aad(&mh),
                                false,
                            ),

src/sync/utils.rs:88

• This exported compatibility path deliberately uses a default context, so an existing or hand-written handler that calls response_to_channel on a secured connection sends plaintext. A deprecation warning does not enforce the framework's no-bypass security invariant. Apply transformation at the writer boundary, or prevent secured handlers from accessing an untransformed response channel.

    let ctx = ConnectionContext::default();
    send_response(stream_id, res, tx, &ctx)

src/sync/server.rs:541

• When this transform fails, the original unauthenticated header is retained and later dispatched; the worker responds using its stream_id. This lets a tampered header target an arbitrary stream and leaves a potentially desynchronized transform active. Close the connection on transform failure without enqueueing this workload.

                                                    let res = reader_ctx.inbound_buf(
                                                        data,
                                                        &serialize_aad(&x),
                                                        true,
                                                    );

src/asynchronous/stream.rs:371

• Changing this constructor from pub to pub(crate) removes an existing public API while StreamInner remains publicly re-exported. Downstream code that constructs custom streams will stop compiling. Preserve the old public constructor (using a default context) and add a separate crate-private context-aware constructor, or explicitly release this as a breaking change.

    #[allow(clippy::too_many_arguments)]
    pub(crate) fn new(

src/sync/utils.rs:180

• conn_ctx is a new required field on a public struct, so downstream sync tests and adapters that construct TtrpcContext with a struct literal will no longer compile. This is a source-breaking API change despite the PR's backward-compatibility claim; it needs a breaking release or a compatibility-preserving design.

    /// Per-connection extension context (opaque data + optional payload transform).
    /// Immutable after accept. Default (empty data, no transform) when no hook is configured.
    pub conn_ctx: Arc<ConnectionContext>,

src/asynchronous/server.rs:396

• These branches still use msg_type and msg_stream_id copied from a header whose authentication just failed. A tampered DATA header can therefore inject an error into an arbitrary active stream, despite the preceding comment saying unauthenticated fields must not be routed. Treat transform failure as connection-fatal and fail pending streams without consulting this header.

                if msg_type == MESSAGE_TYPE_DATA {
                    let stream_tx = {
                        let guard = self.streams.lock().unwrap();
                        guard.get(&msg_stream_id).cloned()
                    };

src/security_extension.rs:466

• JoinHandle::abort() cannot stop an already-running spawn_blocking task, so a stalled hook retains both a blocking-pool worker and the duplicated socket indefinitely, not merely “briefly.” A connection flood can exhaust threads and file descriptors despite the 30-second timeout. Use an actually interruptible/bounded hook execution model or enforce I/O cancellation at the framework boundary.

                _ = tokio::time::sleep(ACCEPT_HOOK_TIMEOUT) => {
                    // Cancel the blocking task. For queued tasks this prevents
                    // execution entirely. For already-running tasks, abort is
                    // best-effort (spawn_blocking cannot be forcibly killed);
                    // the hook's dup'd fd ensures no fd reuse occurs even if

src/asynchronous/utils.rs:264

• Adding a required field to this public struct is source-breaking for every downstream TtrpcContext struct literal; deriving Default does not make existing literals compile. This conflicts with the stated backward compatibility. Either treat the release as breaking or introduce connection data through a compatibility-preserving API.

    /// Opaque per-connection data from [`AcceptHook`](crate::security_extension::AcceptHook). Immutable after accept.
    /// See [`ConnectionData`](crate::security_extension::ConnectionData) for full contract.
    /// Empty when `security_extension` is not enabled.
    pub connection_data: Arc<ConnectionData>,

src/security_extension.rs:697

• This lock is acquired for every outbound frame even when payload_transform is None. Consequently, merely enabling the feature serializes all concurrent plaintext sends through an async mutex, contradicting the no-hook/no-transform overhead guarantee. Fast-path the None case and only lock when a stateful transform exists.

            let _guard = self.async_outbound_lock.lock().await;
            self.outbound(msg, rpc_error)?;

src/security_extension.rs:729

• The sync mutex is taken for every response even when no transform is configured, unnecessarily serializing plaintext handler responses before the already-thread-safe channel send. Only acquire this ordering lock when payload_transform is present.

            let _guard = self.sync_outbound_lock.lock().unwrap();
            buf = self.outbound_buf(buf, aad, true)?;

src/asynchronous/stream.rs:540

• The no-payload validation only runs when FLAG_NO_DATA is also set. A frame with FLAG_REMOTE_CLOSED and a non-empty payload but without FLAG_NO_DATA is now accepted, whereas the previous server-side validation rejected every close-with-data frame. Validate payload emptiness immediately whenever FLAG_REMOTE_CLOSED is present.

                if (msg.header.flags & FLAG_REMOTE_CLOSED) == FLAG_REMOTE_CLOSED {
                    self.remote_closed = true;
                    if (msg.header.flags & FLAG_NO_DATA) == FLAG_NO_DATA {
                        // Enforce protocol invariant: close frame must carry
                        // no payload after decryption. Prevents a peer from

[Tim-Zhang]

Two follow-up comments on hook resource lifetime and client metadata access.

[Tim-Zhang]

Looks good to me. Thanks a lot for your work on this! @jiangliu