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 12 out of 12 changed files in this pull request and generated 6 comments.

Comments suppressed due to low confidence (1)

src/proto.rs:197

• Making GenMessage #[non_exhaustive] and adding a pub(crate) field is a breaking API change for any downstream crate that previously constructed/destructured ttrpc::proto::GenMessage. If the crate intends to preserve semver compatibility, consider keeping GenMessage out of the public API surface (e.g., pub(crate)/#[doc(hidden)]) or making the new field publicly settable with a documented default, and call out the breaking change explicitly in release notes.

/// This struct is `#[non_exhaustive]` to allow adding internal fields without
/// breaking downstream construction. Use the constructor methods
/// ([`new_data`](Self::new_data), [`new_data_decrypted`](Self::new_data_decrypted),
/// [`new_response`](Self::new_response), [`new_close`](Self::new_close)) instead of
/// struct literals.

[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:702

• respond() applies transform_outbound, but it never re-validates the transformed payload size. If a transform expands the payload beyond MESSAGE_LENGTH_MAX, the server will attempt to send an oversized frame (and the peer will reject it). Add an explicit size check (or msg.check()?) after the transform, consistent with the other response path.

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

[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 (2)

src/proto.rs:197

• The doc comment for GenMessage is internally inconsistent: adding #[non_exhaustive] and a pub(crate) field does break downstream struct-literal construction, and the recommended constructor methods are pub(crate) + #[cfg(feature = "async")], so downstream users can’t follow this guidance (and sync builds won’t have them). Consider updating the docs to reflect that GenMessage is not intended to be constructed downstream (or make the constructors public and available without the async cfg).

/// Generic message of ttrpc.
///
/// This struct is `#[non_exhaustive]` to allow adding internal fields without
/// breaking downstream construction. Use the constructor methods
/// ([`new_data`](Self::new_data), [`new_data_decrypted`](Self::new_data_decrypted),
/// [`new_response`](Self::new_response), [`new_close`](Self::new_close)) instead of
/// struct literals.

src/asynchronous/server.rs:702

• HandlerContext::respond() applies transform_outbound() but never re-checks the post-transform payload size. A transform that expands data (e.g. decompression or header-wrapping) can exceed MESSAGE_LENGTH_MAX and lead to the peer rejecting/discarding the message. The main response path performs this check, but error/status responses sent via respond() currently do not.

    async fn respond(&self, stream_id: u32, resp: Response) -> Result<()> {
        let payload = resp
            .encode()
            .map_err(err_to_others_err!(e, "Encode Response failed."))?;
        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);
        self.tx

[Copilot]

Pull request overview

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

Comments suppressed due to low confidence (2)

src/asynchronous/server.rs:429

• check_oversize(resp.compute_size()) runs before transform_outbound(). If a PayloadTransform shrinks payloads (e.g., compression), this will incorrectly reject responses that would be under the limit on the wire after transformation. Size enforcement should be based on the bytes actually sent (post-transform).

                    Some(mut resp) => {
                        // Server: check size before sending to client
                        if let Err(e) = check_oversize(resp.compute_size() as usize, true) {
                            resp = e.into();
                        }

src/asynchronous/utils.rs:264

• This adds a new public field to TtrpcContext. Since TtrpcContext is a public struct, adding a field is a Rust semver-breaking change for any downstream code that constructs it with a struct literal—contradicting the PR description's “no breakage” claim. Consider documenting this as a breaking change and/or providing a stable construction/access pattern (e.g., constructors/getters) to reduce future churn.

    /// Opaque per-connection data from [`AcceptHook`](crate::extension::AcceptHook). Immutable after accept.
    /// See [`ConnectionData`](crate::extension::ConnectionData) for full contract.
    pub connection_data: Arc<ConnectionData>,
}

[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 (6)

src/extension.rs:383

• transform_inbound() matches on self.payload_transform by value, which attempts to move the Option<Arc<..>> out of &self and will not compile. Match on a borrow (e.g., .as_ref() or &self.payload_transform) instead.

        match self.payload_transform {
            Some(ref xform) => xform.transform_inbound(payload),
            None => Ok(payload),
        }

src/extension.rs:392

• transform_outbound() matches on self.payload_transform by value, which attempts to move the Option<Arc<..>> out of &self and will not compile. Match on a borrow (e.g., .as_ref() or &self.payload_transform) instead.

        match self.payload_transform {
            Some(ref xform) => xform.transform_outbound(payload),
            None => Ok(payload),
        }

src/extension.rs:402

• transform_inbound_msg() uses if let Some(ref xform) = self.payload_transform, which tries to move payload_transform out of &self and will not compile. Borrow the option (e.g., if let Some(xform) = self.payload_transform.as_ref()).

        if let Some(ref xform) = self.payload_transform {
            msg.payload = xform.transform_inbound(std::mem::take(&mut msg.payload))?;
            // Update header length to match decrypted payload
            msg.header.length = msg.payload.len() as u32;
        }

src/extension.rs:413

• transform_outbound_msg() uses if let Some(ref xform) = self.payload_transform, which tries to move payload_transform out of &self and will not compile. Borrow the option (e.g., if let Some(xform) = self.payload_transform.as_ref()).

        if let Some(ref xform) = self.payload_transform {
            msg.payload = xform.transform_outbound(std::mem::take(&mut msg.payload))?;
            // Update header length to match transformed payload
            msg.header.length = msg.payload.len() as u32;
        }

src/asynchronous/stream.rs:474

• StreamSender::transform_outbound() matches on self.payload_transform by value, which attempts to move the Option<Arc<..>> out of &self and will not compile. Match on a borrow (e.g., .as_ref()) instead.

        match self.payload_transform {
            Some(ref xform) => xform
                .transform_outbound(buf)
                .map_err(|e| Error::Others(format!("transform_outbound failed: {}", e))),
            None => Ok(buf),

src/asynchronous/stream.rs:520

• StreamReceiver::transform_inbound() matches on self.payload_transform by value, which attempts to move the Option<Arc<..>> out of &self and will not compile. Match on a borrow (e.g., .as_ref()) instead.

        match self.payload_transform {
            Some(ref xform) => xform
                .transform_inbound(payload)
                .map_err(|e| Error::Others(format!("transform_inbound failed: {}", e))),
            None => Ok(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/proto.rs:196

• The doc comment says #[non_exhaustive] allows adding internal fields "without breaking downstream compatibility", but #[non_exhaustive] itself is a (source) breaking change for downstream code that does exhaustive matches or struct literals. The comment should be reworded to avoid promising compatibility that isn’t guaranteed.

/// This struct is `#[non_exhaustive]` to allow adding internal fields
/// (such as `internal_flags`) without breaking downstream compatibility.
/// `GenMessage` is constructed internally by the ttrpc framework and is
/// not intended to be built directly by downstream code.

[Copilot]

Pull request overview

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

[Copilot]

Pull request overview

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