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 no new comments.

Suppressed comments (4)

src/asynchronous/server.rs:676

• respond() now returns an error on oversized responses (via check_oversize(...)), which causes the caller to fall back to an INTERNAL status. Previously, oversized responses were converted into an error Response and sent back to the client with an appropriate status. This is a behavioral regression and can also interact badly with transforms (the INTERNAL fallback may fail for the same reason). Handle oversize in respond() by converting the oversize error into a Response and sending that instead of returning Err.

        check_oversize(payload.len(), true)?;
        let msg = GenMessage::new_response(stream_id, payload);
        self.tx

src/asynchronous/transport/unix.rs:50

• Doc comment refers to a non-existent/incorrect feature name (security). The code is gated by the security_extension feature, so the comment should match the actual feature flag to avoid confusion for readers.

    /// Convert a `UnixListener` into a `Listener`.
    ///
    /// Uses `Box::pin(stream! {...})` directly (bypassing `Listener::new()`)
    /// so each accepted connection goes through `Socket::from(socket)`,
    /// which captures the raw fd when the `security` feature is enabled.

src/asynchronous/transport/tcp.rs:50

• Doc comment refers to a non-existent/incorrect feature name (security). The code is gated by the security_extension feature, so the comment should match the actual feature flag to avoid confusion for readers.

    /// Convert a `TcpListener` into a `Listener`.
    ///
    /// Uses `Box::pin(stream! {...})` directly (bypassing `Listener::new()`)
    /// so each accepted connection goes through `Socket::from(socket)`,
    /// which captures the raw fd when the `security` feature is enabled.

src/asynchronous/transport/vsock.rs:35

• Doc comment refers to a non-existent/incorrect feature name (security). The code is gated by the security_extension feature, so the comment should match the actual feature flag to avoid confusion for readers.

    /// Convert a `VsockListener` into a `Listener`.
    ///
    /// Uses `Box::pin(stream! {...})` directly (bypassing `Listener::new()`)
    /// so each accepted connection goes through `Socket::from(socket)`,
    /// which captures the raw fd when the `security` feature is enabled.

[jiangliu]

Some questions, thank you.

• Why use hooks ?
  Is it possible to encrypt the entire TTRPC layer, similar to encrypting HTTP over TLS?
  Or could it be supported at the protocol layer, such as by adding FLAG_ENCRYPT, FLAG_AUTH flags or MESSAGE_TYPE_ENCRYPT_DATA, MESSAGE_TYPE_AUTH or something similar?
  Is there a need that cannot be met through a more general approach?
• Will there be compatibility issues with the Go version?
• Is it possible to switch this feature's on/off using features?

Hi @wllenyj，
Updated the patchset with following changes:

• introduce feature flag security_extension to improve backward compatibility
• add support of sync in addition to async
• enhance github actions/Makefiles to support security extension
  Please help to review again:)

[Tim-Zhang]

Thanks for this work. I ran the full async, sync, and security extension tests, clippy, and the main feature combinations. They pass.

I found two design blockers and several bugs around header protection, blocking hooks, size checks, error responses, fd ownership, and existing error behavior. Please address the inline comments before merging.

I also added a few refactor suggestions to keep the sync and async paths consistent and easier to maintain.

[Copilot]

Pull request overview

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

Suppressed comments (4)

src/asynchronous/server.rs:463

• The DATA close-message validation rejects any close frame that carries a payload. With security_extension enabled, StreamSender::close_send() applies transform_outbound to close messages, which can legitimately produce a non-empty wire payload (e.g., AEAD tag/nonce) that decrypts to empty. This check will incorrectly fail encrypted close frames before they reach StreamReceiver::recv() for decryption.

                if (msg.header.flags & FLAG_REMOTE_CLOSED) == FLAG_REMOTE_CLOSED
                    && !msg.payload.is_empty()
                {

src/asynchronous/stream.rs:643

• StreamReceiver::recv() returns Eof for REMOTE_CLOSED|NO_DATA without verifying that the decrypted payload is empty. If you relax the server-side raw-payload check to allow encrypted close frames, this receiver should enforce the protocol invariant after transform_inbound to prevent a peer from smuggling data on a close frame.

                let msg = self.transform_data_msg(stream_msg)?;
                if (msg.header.flags & FLAG_REMOTE_CLOSED) == FLAG_REMOTE_CLOSED {
                    self.remote_closed = true;
                    if (msg.header.flags & FLAG_NO_DATA) == FLAG_NO_DATA {
                        return Err(Error::Eof);

src/asynchronous/server.rs:679

• On outbound transform failure, respond() builds a fallback Response but sends it without applying the connection's payload transform. On an encrypted connection, the client will try to decrypt this plaintext fallback and fail to decode the response. The fallback should go through conn_ctx.outbound() as well (and if that fails, bubble the error so the connection can be closed / handled).

        if let Err(e) = self.conn_ctx.outbound(&mut msg, true) {
            let fallback = Response::from(e);
            let fallback_payload = fallback
                .encode()
                .map_err(err_to_others_err!(e, "Encode fallback Response failed."))?;

src/asynchronous/server.rs:466

• Typo in the close-message error string: "connot" should be "cannot".

                        get_status(

[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,