diff --git a/.github/workflows/bvt.yml b/.github/workflows/bvt.yml index b3011f5c..8427476e 100644 --- a/.github/workflows/bvt.yml +++ b/.github/workflows/bvt.yml @@ -33,6 +33,16 @@ jobs: # https://github.com/actions/runner-images/issues/6668 shell: bash + security-extension: + name: Security Extension Tests + runs-on: ubuntu-latest + steps: + - name: Checkout + uses: actions/checkout@v3 + - name: Run tests with security_extension + run: | + cargo test --features "async,sync,security_extension" --verbose + deny: runs-on: ubuntu-latest strategy: diff --git a/Cargo.toml b/Cargo.toml index 7f7a6aa8..6f24a7a3 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -66,6 +66,7 @@ protobuf-codegen = { workspace = true } default = ["sync"] async = ["async-trait", "async-stream", "tokio", "futures", "tokio-vsock"] sync = [] +security_extension = [] [package.metadata.docs.rs] all-features = true diff --git a/Makefile b/Makefile index 3ec99be6..c2c0a441 100644 --- a/Makefile +++ b/Makefile @@ -19,14 +19,24 @@ build: debug # Tests and linters # +# The `security_extension` feature is only supported on Unix platforms +# (see the compile_error! in src/lib.rs), so avoid `--all-features` on +# Windows. Sub-crate Makefiles (compiler, ttrpc-codegen) pre-set FEATURES +# before including this file, as they don't have these features. +ifeq ($(OS),Windows_NT) +FEATURES ?= --features sync,async +else +FEATURES ?= --all-features +endif + .PHONY: test test: - cargo test --all-features --verbose - + cargo test $(FEATURES) --verbose + .PHONY: check check: cargo fmt --all -- --check - cargo clippy --all-targets --all-features -- -D warnings + cargo clippy --all-targets $(FEATURES) -- -D warnings .PHONY: check-all check-all: diff --git a/compiler/Makefile b/compiler/Makefile index bb69e69a..291e6c00 100644 --- a/compiler/Makefile +++ b/compiler/Makefile @@ -1 +1,2 @@ +FEATURES = --all-features include ../Makefile diff --git a/src/asynchronous/client.rs b/src/asynchronous/client.rs index 524d6432..c522b866 100644 --- a/src/asynchronous/client.rs +++ b/src/asynchronous/client.rs @@ -15,6 +15,9 @@ use async_trait::async_trait; use tokio::{self, sync::mpsc, task}; use crate::error::{get_rpc_status, Error, Result}; +use crate::ConnectionContext; +#[cfg(feature = "security_extension")] +use crate::security_extension::ConnectHook; use crate::proto::{ Code, Codec, GenMessage, Message, MessageHeader, Request, Response, FLAG_NO_DATA, FLAG_REMOTE_CLOSED, FLAG_REMOTE_OPEN, MESSAGE_TYPE_DATA, MESSAGE_TYPE_RESPONSE, @@ -34,6 +37,7 @@ pub struct Client { req_tx: MessageSender, next_stream_id: Arc, streams: Arc>>, + conn_ctx: Arc, } impl Client { @@ -41,7 +45,7 @@ impl Client { let socket = Socket::connect(sockaddr) .await .map_err(err_to_others_err!(e, "Socket::connect error "))?; - Ok(Self::new(socket)) + Self::new_inner(socket, None) } #[cfg(unix)] @@ -54,45 +58,120 @@ impl Client { /// Initialize a new [`Client`]. pub fn new(stream: Socket) -> Client { - let (req_tx, rx): (MessageSender, MessageReceiver) = mpsc::channel(100); + Self::new_inner(stream, None) + .expect("new_inner without hook cannot fail") + } + + /// Initialize a new [`Client`] with a connection hook. + /// + /// The hook is invoked synchronously during construction, receiving the + /// socket's raw file descriptor so it can inspect peer identity (e.g., + /// via `getpeername`). Its [`HookOutput`](crate::security_extension::HookOutput) + /// — connection metadata and optional payload transform — is stored in the + /// client's [`ConnectionContext`](crate::security_extension::ConnectionContext) + /// before this method returns. + /// + /// # Blocking behavior + /// + /// The hook may perform synchronous I/O (e.g., a cryptographic handshake). + /// If called from an async context on a `current_thread` runtime, wrap the + /// call in [`tokio::task::spawn_blocking`] to avoid stalling the executor: + /// + /// ```ignore + /// let client = tokio::task::spawn_blocking(|| { + /// Client::with_hook(stream, hook) + /// }).await??; + /// ``` + /// + /// # Errors + /// + /// Returns an error if the connect hook rejects or otherwise fails, + /// allowing the caller to handle the failure. The socket is consumed + /// regardless; on failure the caller should open a new connection. + #[cfg(feature = "security_extension")] + pub fn with_hook(stream: Socket, hook: H) -> Result { + Self::new_inner(stream, Some(Box::new(hook))) + } + /// Returns the per-connection metadata from the [`ConnectHook`]. + /// + /// This is the [`ConnectionData`](crate::security_extension::ConnectionData) + /// returned by the hook during connection establishment. Empty (default) + /// when no hook was configured. + #[cfg(feature = "security_extension")] + pub fn connection_data(&self) -> &crate::security_extension::ConnectionData { + &self.conn_ctx.data + } + + fn new_inner( + stream: Socket, + #[cfg(feature = "security_extension")] hook: Option>, + #[cfg(not(feature = "security_extension"))] _hook: Option<()>, + ) -> Result { + // ── Injection Point 5/10: connect hook ── + // Call connect hook if set, create ConnectionContext from output + #[cfg(feature = "security_extension")] + let conn_ctx = match hook { + Some(h) => match stream.as_raw_fd() { + Some(fd) => match h.on_connect(fd) { + Ok(output) => Arc::new(ConnectionContext::new(Some(output))), + Err(e) => { + return Err(Error::Others(format!( + "client connect hook failed (fd={}): {}", + fd, e + ))); + } + }, + None => { + return Err(Error::Others( + "client connect hook configured but socket has no raw fd; construct the Socket via Socket::connect or Socket::from()".to_string(), + )); + } + }, + None => Arc::new(ConnectionContext::default()), + }; + #[cfg(not(feature = "security_extension"))] + let conn_ctx = Arc::new(ConnectionContext::default()); + + let (req_tx, rx): (MessageSender, MessageReceiver) = mpsc::channel(100); let req_map = Arc::new(Mutex::new(HashMap::new())); let delegate = ClientBuilder { rx: Some(rx), streams: req_map.clone(), + conn_ctx: conn_ctx.clone(), }; let conn = Connection::new(stream, delegate); - // Long-running receiver task tokio::spawn(async move { conn.run().await }); - Client { + Ok(Client { req_tx, next_stream_id: Arc::new(AtomicU32::new(1)), streams: req_map, - } + conn_ctx, + }) } - /// Requsts a unary request and returns with response. + /// Requests a unary request and returns with response. pub async fn request(&self, req: Request) -> Result { let timeout_nano = req.timeout_nano; let stream_id = self.next_stream_id.fetch_add(2, Ordering::Relaxed); - let msg: GenMessage = Message::new_request(stream_id, req)? + let mut msg: GenMessage = Message::new_request(stream_id, req)? .try_into() .map_err(|e: protobuf::Error| Error::Others(e.to_string()))?; let (tx, mut rx): (ResultSender, ResultReceiver) = mpsc::channel(100); - self.streams .lock() .map_err(|_| Error::Others("Failed to acquire lock on streams".to_string()))? .insert(stream_id, tx); - self.req_tx - .send(SendingMessage::new(msg)) - .await - .map_err(|_| Error::LocalClosed)?; + // ── Injection Point 6/10: unary REQUEST transform_outbound ── + if let Err(e) = self.conn_ctx.transform_send(&mut msg, &self.req_tx, false, false).await { + self.streams.lock().unwrap().remove(&stream_id); + return Err(e); + } let result = if timeout_nano == 0 { rx.recv().await.ok_or(Error::RemoteClosed)? @@ -151,10 +230,11 @@ impl Client { .map_err(|_| Error::Others("Failed to acquire lock on streams".to_string()))? .insert(stream_id, tx); - self.req_tx - .send(SendingMessage::new(msg)) - .await - .map_err(|e| Error::Others(format!("Send packet to sender error {e:?}")))?; + // ── Injection Point 8/10: stream-init REQUEST transform_outbound ── + if let Err(e) = self.conn_ctx.transform_send(&mut msg, &self.req_tx, false, false).await { + self.streams.lock().unwrap().remove(&stream_id); + return Err(e); + } Ok(StreamInner::new( stream_id, @@ -164,6 +244,7 @@ impl Client { streaming_server, Kind::Client, self.streams.clone(), + self.conn_ctx.clone(), )) } } @@ -172,6 +253,7 @@ impl Client { struct ClientBuilder { rx: Option, streams: Arc>>, + conn_ctx: Arc, } impl Builder for ClientBuilder { @@ -184,11 +266,11 @@ impl Builder for ClientBuilder { ClientReader { shutdown_waiter: waiter, streams: self.streams.clone(), + conn_ctx: self.conn_ctx.clone(), }, ClientWriter { rx: self.rx.take().unwrap(), shutdown_notifier: notifier, - streams: self.streams.clone(), }, ) @@ -286,6 +368,7 @@ async fn get_resp_tx( struct ClientReader { streams: Arc>>, shutdown_waiter: shutdown::Waiter, + conn_ctx: Arc, } #[async_trait] @@ -326,13 +409,62 @@ impl ReaderDelegate for ClientReader { async fn handle_msg(&self, msg: GenMessage) { let req_map = self.streams.clone(); + let conn_ctx = self.conn_ctx.clone(); + + // ── Inbound transform in wire order ── + // Apply transform here (in the connection read loop) before spawning + // a handler task. This ensures deterministic nonce sequencing for + // stateful transforms (e.g., AEAD) regardless of task scheduling. + let mut msg = msg; + let result = conn_ctx.inbound(&mut msg, false); + tokio::spawn(async move { if let Some(resp_tx) = get_resp_tx(req_map, &msg.header).await { resp_tx - .send(Ok(msg)) + .send(result.map(|_| msg)) .await .unwrap_or_else(|_e| error!("The request has returned")); } }); } } + +#[cfg(all(test, feature = "security_extension"))] +mod tests { + use super::*; + use crate::security_extension::{ConnectHook, ConnectionData, HookError, HookOutput}; + + #[derive(Debug)] + struct DummyConnectHook; + + impl ConnectHook for DummyConnectHook { + fn on_connect( + &self, + _fd: std::os::unix::io::RawFd, + ) -> std::result::Result { + Ok(HookOutput { + data: ConnectionData::new(), + payload_transform: None, + }) + } + } + + /// Constructing a Socket via Socket::new() leaves raw_fd == None. + /// Client::with_hook must fail in that case instead of silently skipping + /// the hook and using an untransformed connection. + #[test] + fn with_hook_requires_raw_fd() { + let (client, _server) = tokio::io::duplex(64); + let socket = Socket::new(client); + let err = match Client::with_hook(socket, DummyConnectHook) { + Ok(_) => panic!("hook configured but no fd -> should fail"), + Err(e) => e, + }; + let err_str = format!("{}", err); + assert!( + err_str.contains("socket has no raw fd"), + "error should tell caller to use Socket::connect / Socket::from: {}", + err_str + ); + } +} diff --git a/src/asynchronous/mod.rs b/src/asynchronous/mod.rs index 819c5b06..9e43980d 100644 --- a/src/asynchronous/mod.rs +++ b/src/asynchronous/mod.rs @@ -20,6 +20,7 @@ pub use self::stream::{ SSSender, ServerStream, ServerStreamReceiver, ServerStreamSender, StreamInner, StreamReceiver, StreamSender, }; +pub(crate) use self::stream::SendingMessage; #[doc(inline)] pub use crate::r#async::client::Client; #[doc(inline)] diff --git a/src/asynchronous/server.rs b/src/asynchronous/server.rs index dca04e33..0ebf831e 100644 --- a/src/asynchronous/server.rs +++ b/src/asynchronous/server.rs @@ -14,7 +14,7 @@ use std::time::Duration; use async_trait::async_trait; use futures::StreamExt as _; -use protobuf::Message as _; +use protobuf::Message as PbMessage; use tokio::{ self, select, spawn, sync::mpsc::{channel, Sender}, @@ -26,9 +26,12 @@ use crate::asynchronous::stream::SendingMessage; use crate::asynchronous::transport::{Listener, Socket}; use crate::context; use crate::error::{get_status, Error, Result}; +use crate::ConnectionContext; +#[cfg(feature = "security_extension")] +use crate::security_extension::{AcceptHook, ServerExtensionConfig}; use crate::proto::{ check_oversize, Code, Codec, GenMessage, Message, MessageHeader, Request, Response, Status, - FLAG_NO_DATA, FLAG_REMOTE_CLOSED, MESSAGE_TYPE_DATA, MESSAGE_TYPE_REQUEST, + FLAG_NO_DATA, MESSAGE_TYPE_DATA, MESSAGE_TYPE_REQUEST, }; use crate::r#async::connection::*; use crate::r#async::shutdown; @@ -63,6 +66,11 @@ pub struct Server { shutdown: shutdown::Notifier, stop_listen_tx: Option>>, + + // ── Connection Extension Framework ── + // See crate::security_extension for full architecture documentation. + #[cfg(feature = "security_extension")] + accept_hook: Option>, } impl Default for Server { @@ -72,6 +80,8 @@ impl Default for Server { services: Arc::new(HashMap::new()), shutdown: shutdown::with_timeout(DEFAULT_SERVER_SHUTDOWN_TIMEOUT).0, stop_listen_tx: None, + #[cfg(feature = "security_extension")] + accept_hook: None, } } } @@ -119,6 +129,21 @@ impl Server { Ok(self.add_listener(listener)) } + /// Register a hook called on every new accepted connection. + /// Replaces any previously registered hook. + /// + /// Note (Unix): the hook requires the accepted transport + /// [`Socket`](crate::asynchronous::transport::Socket) to carry a captured + /// raw fd (`Socket::as_raw_fd() != None`). This is true for listeners + /// created via `Server::bind()` / `Listener::from()`, + /// but not for sockets produced by `Listener::new()` / `Socket::new()`. + #[cfg(feature = "security_extension")] + pub fn set_accept_hook(mut self, hook: H) -> Self { + let hook: Arc = Arc::new(hook); + self.accept_hook = Some(hook); + self + } + pub fn register_service(mut self, new: HashMap) -> Server { let services = Arc::get_mut(&mut self.services).unwrap(); services.extend(new); @@ -136,6 +161,7 @@ impl Server { self.do_start(incoming).await } + // ── Connection Extension: Injection Point 1/10 (accept hook) ── async fn do_start(&mut self, mut incoming: Listener) -> Result<()> { let services = self.services.clone(); @@ -144,6 +170,11 @@ impl Server { let (stop_listen_tx, mut stop_listen_rx) = channel(1); self.stop_listen_tx = Some(stop_listen_tx); + #[cfg(feature = "security_extension")] + let server_ext = Arc::new(ServerExtensionConfig { + accept_hook: self.accept_hook.clone(), + }); + spawn(async move { loop { select! { @@ -152,12 +183,33 @@ impl Server { // Accept a new connection match conn { Ok(conn) => { - // spawn a connection handler, would not block - spawn_connection_handler( - conn, - services.clone(), - shutdown_waiter.clone(), - ).await; + // Spawn hook + handler setup per-connection + // so the accept loop can immediately return + // to accepting new connections. + #[cfg(feature = "security_extension")] + let server_ext = server_ext.clone(); + let services = services.clone(); + let shutdown_waiter = shutdown_waiter.clone(); + spawn(async move { + // ── 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)), + Err(e) => { + log::warn!("accept hook failed for connection: {:?}", e); + return; + } + }; + #[cfg(not(feature = "security_extension"))] + let conn_ctx = Arc::new(ConnectionContext::default()); + + spawn_connection_handler( + conn, + services, + shutdown_waiter, + conn_ctx, + ).await; + }); } Err(e) => { error!("incoming conn fail {:?}", e) @@ -216,11 +268,13 @@ async fn spawn_connection_handler( conn: Socket, services: Arc>, shutdown_waiter: shutdown::Waiter, + conn_ctx: Arc, ) { let delegate = ServerBuilder { services, streams: Arc::new(Mutex::new(HashMap::new())), shutdown_waiter, + conn_ctx, }; let conn = Connection::new(conn, delegate); spawn(async move { @@ -237,6 +291,7 @@ struct ServerBuilder { services: Arc>, streams: Arc>>, shutdown_waiter: shutdown::Waiter, + conn_ctx: Arc, } impl Builder for ServerBuilder { @@ -255,6 +310,7 @@ impl Builder for ServerBuilder { streams: self.streams.clone(), server_shutdown: self.shutdown_waiter.clone(), handler_shutdown: disconnect_notifier, + conn_ctx: self.conn_ctx.clone(), }, ServerWriter { rx, @@ -284,6 +340,7 @@ struct ServerReader { streams: Arc>>, server_shutdown: shutdown::Waiter, handler_shutdown: shutdown::Notifier, + conn_ctx: Arc, } #[async_trait] @@ -315,6 +372,27 @@ impl ReaderDelegate for ServerReader { //Check if it is already shutdown no need select wait if !handler_shutdown_waiter.is_shutdown() { let (wait_tx, wait_rx) = tokio::sync::oneshot::channel::<()>(); + + // ── Inbound transform for ALL frames in wire order ── + // Authenticate every inbound frame sequentially before any header + // validation or routing. This prevents: + // 1. Bypassing AAD verification by sending invalid stream_ids + // 2. Desynchronizing stateful transforms (counter-based nonce) + // The connection reader loop calls handle_msg sequentially, + // ensuring transforms execute in wire order. + let mut msg = msg; + let is_request = msg.header.type_ == MESSAGE_TYPE_REQUEST; + if let Err(e) = self.conn_ctx.inbound(&mut msg, is_request) { + // Transform failure: drop the frame silently. + // Do NOT use any unauthenticated header fields (type, + // stream_id) for routing — an attacker could tamper with + // them to redirect errors to unrelated active streams. + // Affected stream handlers will time out via their own + // request deadline or be cleaned up on disconnect. + error!("transform inbound failed (frame dropped): {}", e); + return; + } + spawn(async move { select! { _ = context.handle_msg(msg, wait_tx) => {} @@ -337,6 +415,7 @@ impl ServerReader { services: self.services.clone(), streams: self.streams.clone(), _handler_shutdown_waiter: self.handler_shutdown.subscribe(), + conn_ctx: self.conn_ctx.clone(), } } } @@ -347,11 +426,12 @@ struct HandlerContext { streams: Arc>>, // Used for waiting handler exit. _handler_shutdown_waiter: shutdown::Waiter, + conn_ctx: Arc, } impl HandlerContext { async fn handle_err(&self, header: MessageHeader, e: Error) { - Self::respond(self.tx.clone(), header.stream_id, e.into()) + self.respond(header.stream_id, e.into()) .await .map_err(|e| { error!("respond error got error {:?}", e); @@ -362,8 +442,7 @@ impl HandlerContext { let stream_id = msg.header.stream_id; if (stream_id % 2) != 1 { - Self::respond_with_status( - self.tx.clone(), + self.respond_with_status( stream_id, get_status(Code::INVALID_ARGUMENT, "stream id must be odd"), ) @@ -375,60 +454,37 @@ impl HandlerContext { MESSAGE_TYPE_REQUEST => match self.handle_request(msg, wait_tx).await { Ok(opt_msg) => match opt_msg { 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(); } - - Self::respond(self.tx.clone(), stream_id, resp) - .await - .map_err(|e| { - error!("respond got error {:?}", e); - }) - .ok(); + 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); + } } None => { - let mut header = MessageHeader::new_data(stream_id, 0); - header.set_flags(FLAG_REMOTE_CLOSED | FLAG_NO_DATA); - let msg = GenMessage { - header, - payload: Vec::new(), - }; - - self.tx - .send(SendingMessage::new(msg)) - .await - .map_err(err_to_others_err!(e, "Send packet to sender error ")) - .ok(); + let mut msg = GenMessage::new_close(stream_id); + if let Err(e) = self.conn_ctx.transform_send(&mut msg, &self.tx, false, false).await { + error!("transform close message failed: {}", e); + } } }, - Err(status) => Self::respond_with_status(self.tx.clone(), stream_id, status).await, + Err(status) => self.respond_with_status(stream_id, status).await, }, MESSAGE_TYPE_DATA => { // no need to wait data message handling drop(wait_tx); - // TODO(wllenyj): Compatible with golang behavior. - if (msg.header.flags & FLAG_REMOTE_CLOSED) == FLAG_REMOTE_CLOSED - && !msg.payload.is_empty() - { - Self::respond_with_status( - self.tx.clone(), - stream_id, - get_status( - Code::INVALID_ARGUMENT, - format!( - "Stream id {stream_id}: data close message connot include data" - ), - ), - ) - .await; - return; - } + + // DATA transform already applied in ServerReader::handle_msg() + // (wire order, before spawn). let stream_tx = self.streams.lock().unwrap().get(&stream_id).cloned(); if let Some(stream_tx) = stream_tx { if let Err(e) = stream_tx.send(Ok(msg)).await { - Self::respond_with_status( - self.tx.clone(), + self.respond_with_status( stream_id, get_status( Code::INVALID_ARGUMENT, @@ -438,8 +494,7 @@ impl HandlerContext { .await; } } else { - Self::respond_with_status( - self.tx.clone(), + self.respond_with_status( stream_id, get_status(Code::INVALID_ARGUMENT, "Stream is no longer active"), ) @@ -465,6 +520,9 @@ impl HandlerContext { //} // self.last_stream_id = header.stream_id; + // ── REQUEST transform_inbound already applied in ServerReader::handle_msg() ── + // (wire order, before any header validation or routing) + let req_msg = Message::::try_from(msg) .map_err(|e| get_status(Code::INVALID_ARGUMENT, e.to_string()))?; @@ -503,6 +561,7 @@ impl HandlerContext { mh: req_msg.header, metadata: context::from_pb(&req.metadata), timeout_nano: req.timeout_nano, + connection_data: self.conn_ctx.data.clone(), }; let get_unknown_status_and_log_err = |e| { @@ -560,18 +619,28 @@ impl HandlerContext { true, Kind::Server, self.streams.clone(), + self.conn_ctx.clone(), ); let ctx = TtrpcContext { mh: req_msg.header, metadata: context::from_pb(&req.metadata), timeout_nano: req.timeout_nano, + connection_data: self.conn_ctx.data.clone(), }; let task = spawn(async move { stream.handler(ctx, si).await }); if !no_data { - // Fake the first data message. + // "Fake" the first DATA message from the stream-init REQUEST payload. + // + // `req.payload` has already been decrypted by `handle_request()`. + // The payload is decrypted exactly once regardless of transform type. + // + // For the common duplex streaming case (`streaming_client = true`, + // FLAG_NO_DATA set), this block is skipped and all DATA messages + // arrive via the normal `handle_msg` route where they are + // transformed in the connection reader (wire order). let msg = GenMessage { header: MessageHeader::new_data(stream_id, req.payload.len() as u32), payload: req.payload, @@ -586,23 +655,41 @@ impl HandlerContext { .map_err(|e| get_status(Code::UNKNOWN, e)) } - async fn respond(tx: MessageSender, stream_id: u32, resp: Response) -> Result<()> { + 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 msg = GenMessage { - header: MessageHeader::new_response(stream_id, payload.len() as u32), - payload, + + // Pre-check: ensure raw payload fits within the transform-safe limit. + // This avoids calling transform_outbound on data that would be rejected + // post-transform, which would advance stateful transforms (e.g., AEAD + // nonce counters) without producing a sendable message. + let max_len = self.conn_ctx.max_raw_payload_len(); + let payload = if payload.len() > max_len { + // Original too large — build a small error response that is + // guaranteed to fit after transform (single transform call). + let err_msg = format!( + "response payload {} bytes exceeds safe limit {} bytes (after transform overhead)", + payload.len(), + max_len + ); + let err_resp: Response = + Error::RpcStatus(get_status(Code::INVALID_ARGUMENT, err_msg)).into(); + err_resp + .encode() + .map_err(err_to_others_err!(e, "Encode error Response failed."))? + } else { + payload }; - tx.send(SendingMessage::new(msg)) - .await - .map_err(err_to_others_err!(e, "Send packet to sender error ")) + + let mut msg = GenMessage::new_response(stream_id, payload); + self.conn_ctx.transform_send(&mut msg, &self.tx, true, false).await } - async fn respond_with_status(tx: MessageSender, stream_id: u32, status: Status) { + async fn respond_with_status(&self, stream_id: u32, status: Status) { let mut resp = Response::new(); resp.set_status(status); - Self::respond(tx, stream_id, resp) + self.respond(stream_id, resp) .await .map_err(|e| { error!("respond with status got error {:?}", e); diff --git a/src/asynchronous/stream.rs b/src/asynchronous/stream.rs index 4324532c..e3a78f50 100644 --- a/src/asynchronous/stream.rs +++ b/src/asynchronous/stream.rs @@ -9,18 +9,22 @@ use std::marker::PhantomData; use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::{Arc, Mutex}; +use crate::ConnectionContext; + use tokio::sync::mpsc; use super::Client; use crate::error::{Error, Result}; use crate::proto::{ - Code, Codec, GenMessage, MessageHeader, Response, FLAG_NO_DATA, FLAG_REMOTE_CLOSED, - MESSAGE_TYPE_DATA, MESSAGE_TYPE_RESPONSE, + check_oversize, Code, Codec, GenMessage, Response, FLAG_NO_DATA, + FLAG_REMOTE_CLOSED, MESSAGE_TYPE_DATA, MESSAGE_TYPE_RESPONSE, }; pub type MessageSender = mpsc::Sender; pub type MessageReceiver = mpsc::Receiver; +/// Internal message type for stream channels. +/// pub type ResultSender = mpsc::Sender>; pub type ResultReceiver = mpsc::Receiver>; @@ -350,16 +354,6 @@ async fn _recv(rx: &mut ResultReceiver) -> Result { }) } -async fn _send(tx: &MessageSender, msg: GenMessage) -> Result<()> { - let (res_tx, res_rx) = tokio::sync::oneshot::channel(); - tx.send(SendingMessage::new_with_result(msg, res_tx)) - .await - .map_err(|e| Error::Others(format!("Send data packet to sender error {:?}", e)))?; - res_rx - .await - .map_err(|e| Error::Others(format!("Failed to wait send result {:?}", e)))? -} - #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub enum Kind { Client, @@ -373,7 +367,8 @@ pub struct StreamInner { } impl StreamInner { - pub fn new( + #[allow(clippy::too_many_arguments)] + pub(crate) fn new( stream_id: u32, tx: MessageSender, rx: ResultReceiver, @@ -382,6 +377,7 @@ impl StreamInner { recveivable: bool, kind: Kind, streams: Arc>>, + conn_ctx: Arc, ) -> Self { Self { sender: StreamSender { @@ -390,6 +386,7 @@ impl StreamInner { sendable, local_closed: Arc::new(AtomicBool::new(false)), kind, + conn_ctx, }, receiver: StreamReceiver { rx, @@ -419,6 +416,11 @@ impl StreamInner { } } +/// Sends streaming DATA messages with optional payload transform. +/// +/// Part of the extension framework's Injection Point 9/10. +/// Each `send()` call applies `transform_outbound` to the payload +/// before framing it as a DATA message. #[derive(Clone, Debug)] pub struct StreamSender { tx: MessageSender, @@ -426,8 +428,16 @@ pub struct StreamSender { sendable: bool, local_closed: Arc, kind: Kind, + conn_ctx: Arc, } +/// Receives streaming DATA messages. +/// +/// Inbound transforms are applied in the connection read loop (wire order) +/// before messages are routed here. `recv()` therefore receives payloads +/// that are already decrypted and only needs to decode/pass through. +/// +/// `handle_msg()` routes DATA messages to this receiver. #[derive(Debug)] pub struct StreamReceiver { rx: ResultReceiver, @@ -445,21 +455,20 @@ impl Drop for StreamReceiver { } impl StreamSender { + /// Send a streaming DATA message. + /// + /// Applies `transform_outbound` (Injection Point 9/10) to `buf`, + /// then frames it as a DATA message and sends to the transport. pub async fn send(&self, buf: Vec) -> Result<()> { debug_assert!(self.sendable); if self.local_closed.load(Ordering::Relaxed) { debug_assert_eq!(self.kind, Kind::Client); return Err(Error::LocalClosed); } - let header = MessageHeader::new_data(self.stream_id, buf.len() as u32); - let msg = GenMessage { - header, - payload: buf, - }; - - msg.check()?; - _send(&self.tx, msg).await?; + 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, true).await?; Ok(()) } @@ -470,19 +479,20 @@ impl StreamSender { if self.local_closed.load(Ordering::Relaxed) { return Err(Error::LocalClosed); } - let mut header = MessageHeader::new_data(self.stream_id, 0); - header.set_flags(FLAG_REMOTE_CLOSED | FLAG_NO_DATA); - let msg = GenMessage { - header, - payload: Vec::new(), - }; - _send(&self.tx, msg).await?; + let mut msg = GenMessage::new_close(self.stream_id); + self.conn_ctx.transform_send(&mut msg, &self.tx, false, true).await?; self.local_closed.store(true, Ordering::Relaxed); Ok(()) } } impl StreamReceiver { + /// Receive the next streaming message. + /// + /// All inbound transforms are applied in the connection read loop + /// (wire order) before messages reach this receiver, so `recv()` + /// only decodes/passes through the pre-decoded payload. + /// Returns `Err(Error::Eof)` when the remote side closes the stream. pub async fn recv(&mut self) -> Result> { if self.remote_closed { return Err(Error::RemoteClosed); @@ -509,12 +519,26 @@ impl StreamReceiver { "received data from non-streaming server.".to_string(), )); } + // Close-flag checks on pre-decoded payload. + // Transform was applied in wire order by the connection reader. + // A tampered close would have failed AEAD verification there + // and never reached this point. 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 + // smuggling data on a close frame. + if !msg.payload.is_empty() { + return Err(Error::Others(format!( + "stream {}: close message cannot include data", + self.stream_id + ))); + } return Err(Error::Eof); } } + check_oversize(msg.payload.len(), false)?; msg.payload } _ => { diff --git a/src/asynchronous/transport/mod.rs b/src/asynchronous/transport/mod.rs index f340c37f..c660efd8 100644 --- a/src/asynchronous/transport/mod.rs +++ b/src/asynchronous/transport/mod.rs @@ -1,5 +1,7 @@ use std::io::{Error as IoError, Result as IoResult}; use std::pin::Pin; +#[cfg(feature = "security_extension")] +use std::os::unix::io::RawFd; use futures::stream::{BoxStream, Stream, StreamExt as _}; use tokio::io::{AsyncRead, AsyncWrite}; @@ -8,7 +10,16 @@ trait AsyncReadWrite: AsyncRead + AsyncWrite {} impl AsyncReadWrite for T {} pub struct Listener(BoxStream<'static, IoResult>); -pub struct Socket(Pin>); +/// A type-erased async socket. +/// +/// On Unix with the `security_extension` feature, stores the raw fd for +/// [`AcceptHook`](crate::security_extension::AcceptHook) access. +/// See `crate::security_extension` module doc for full architecture. +pub struct Socket { + inner: Pin>, + #[cfg(feature = "security_extension")] + raw_fd: Option, +} macro_rules! io_other { ($fmt_str:literal, $($args:expr),*) => { @@ -32,6 +43,12 @@ mod vsock; mod windows; impl Listener { + /// Create a listener from a generic async stream of sockets. + /// + /// **Note**: This uses [`Socket::new`] which does **not** capture the + /// raw fd. For platform-specific listeners (Unix, TCP, vsock), prefer + /// the `From` impls which capture the raw fd via + /// `Socket::from()` — required by [`AcceptHook`](crate::security_extension::AcceptHook). pub fn new( listener: impl Stream> + Send + 'static, ) -> Self { @@ -66,8 +83,44 @@ impl Listener { } impl Socket { + /// Create a socket from any `AsyncRead + AsyncWrite` type. pub fn new(socket: impl AsyncRead + AsyncWrite + Send + Sync + 'static) -> Self { - Self(Box::pin(socket)) + Self { + inner: Box::pin(socket), + #[cfg(feature = "security_extension")] + raw_fd: None, + } + } + + #[cfg(feature = "security_extension")] + pub(crate) fn with_raw_fd(socket: impl AsyncRead + AsyncWrite + Send + Sync + 'static, fd: RawFd) -> Self { + Self { + inner: Box::pin(socket), + raw_fd: Some(fd), + } + } + + #[cfg(feature = "security_extension")] + pub fn as_raw_fd(&self) -> Option { + self.raw_fd + } + + /// Create a socket from a stream, capturing the raw fd on Unix when + /// `security_extension` is enabled. Eliminates per-transport cfg branching + /// in `From for Socket` impls. + #[cfg(unix)] + pub(crate) fn from_fd_aware( + socket: S, + ) -> Self { + #[cfg(feature = "security_extension")] + { + let fd = socket.as_raw_fd(); + Self::with_raw_fd(socket, fd) + } + #[cfg(not(feature = "security_extension"))] + { + Self::new(socket) + } } pub async fn connect(addr: impl AsRef) -> IoResult { @@ -118,7 +171,7 @@ impl AsyncRead for Socket { cx: &mut std::task::Context<'_>, buf: &mut tokio::io::ReadBuf<'_>, ) -> std::task::Poll> { - self.get_mut().0.as_mut().poll_read(cx, buf) + self.get_mut().inner.as_mut().poll_read(cx, buf) } } @@ -128,21 +181,21 @@ impl AsyncWrite for Socket { cx: &mut std::task::Context<'_>, buf: &[u8], ) -> std::task::Poll> { - self.get_mut().0.as_mut().poll_write(cx, buf) + self.get_mut().inner.as_mut().poll_write(cx, buf) } fn poll_flush( self: Pin<&mut Self>, cx: &mut std::task::Context<'_>, ) -> std::task::Poll> { - self.get_mut().0.as_mut().poll_flush(cx) + self.get_mut().inner.as_mut().poll_flush(cx) } fn poll_shutdown( self: Pin<&mut Self>, cx: &mut std::task::Context<'_>, ) -> std::task::Poll> { - self.get_mut().0.as_mut().poll_shutdown(cx) + self.get_mut().inner.as_mut().poll_shutdown(cx) } fn poll_write_vectored( @@ -150,10 +203,10 @@ impl AsyncWrite for Socket { cx: &mut std::task::Context<'_>, bufs: &[std::io::IoSlice<'_>], ) -> std::task::Poll> { - self.get_mut().0.as_mut().poll_write_vectored(cx, bufs) + self.get_mut().inner.as_mut().poll_write_vectored(cx, bufs) } fn is_write_vectored(&self) -> bool { - self.0.is_write_vectored() + self.inner.is_write_vectored() } } diff --git a/src/asynchronous/transport/tcp.rs b/src/asynchronous/transport/tcp.rs index 25eafd8d..8a5a96f6 100644 --- a/src/asynchronous/transport/tcp.rs +++ b/src/asynchronous/transport/tcp.rs @@ -41,12 +41,20 @@ impl Socket { } impl From for Listener { + /// 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_extension` feature is enabled. fn from(listener: TcpListener) -> Self { - Self::new(stream! { + Self(Box::pin(stream! { loop { - yield listener.accept().await.map(|(socket, _)| socket); + match listener.accept().await { + Ok((socket, _)) => yield Ok(Socket::from(socket)), + Err(e) => yield Err(e), + } } - }) + })) } } @@ -60,7 +68,7 @@ impl TryFrom for Listener { impl From for Socket { fn from(socket: TcpStream) -> Self { - Self::new(socket) + Socket::from_fd_aware(socket) } } diff --git a/src/asynchronous/transport/unix.rs b/src/asynchronous/transport/unix.rs index 37df3b82..089d2c8c 100644 --- a/src/asynchronous/transport/unix.rs +++ b/src/asynchronous/transport/unix.rs @@ -41,12 +41,20 @@ impl Socket { } impl From for Listener { + /// 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_extension` feature is enabled. fn from(listener: UnixListener) -> Self { - Self::new(stream! { + Self(Box::pin(stream! { loop { - yield listener.accept().await.map(|(socket, _)| socket); + match listener.accept().await { + Ok((socket, _)) => yield Ok(Socket::from(socket)), + Err(e) => yield Err(e), + } } - }) + })) } } @@ -60,7 +68,7 @@ impl TryFrom for Listener { impl From for Socket { fn from(socket: UnixStream) -> Self { - Self::new(socket) + Socket::from_fd_aware(socket) } } diff --git a/src/asynchronous/transport/vsock.rs b/src/asynchronous/transport/vsock.rs index 38e3ee72..a41e22ea 100644 --- a/src/asynchronous/transport/vsock.rs +++ b/src/asynchronous/transport/vsock.rs @@ -28,18 +28,26 @@ impl Socket { } impl From for Listener { + /// 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_extension` feature is enabled. fn from(listener: VsockListener) -> Self { - Self::new(stream! { + Self(Box::pin(stream! { loop { - yield listener.accept().await.map(|(socket, _)| socket); + match listener.accept().await { + Ok((socket, _)) => yield Ok(Socket::from(socket)), + Err(e) => yield Err(e), + } } - }) + })) } } impl From for Socket { fn from(socket: VsockStream) -> Self { - Self::new(socket) + Socket::from_fd_aware(socket) } } diff --git a/src/asynchronous/utils.rs b/src/asynchronous/utils.rs index ce315edb..ca547818 100644 --- a/src/asynchronous/utils.rs +++ b/src/asynchronous/utils.rs @@ -5,10 +5,12 @@ // use std::collections::HashMap; +use std::sync::Arc; use async_trait::async_trait; use crate::error::Result; +use crate::security_extension::ConnectionData; use crate::proto::{MessageHeader, Request, Response}; /// Handle request in async mode. @@ -247,11 +249,19 @@ pub trait StreamHandler { } /// The context of ttrpc (async). -#[derive(Debug)] +/// +/// Implements [`Default`] so test/mock code can construct a context with +/// `..Default::default()` and only specify the fields they need. +#[derive(Debug, Default)] pub struct TtrpcContext { pub mh: MessageHeader, pub metadata: HashMap>, pub timeout_nano: i64, + + /// 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, } pub(crate) fn get_path(service: &str, method: &str) -> String { diff --git a/src/lib.rs b/src/lib.rs index 9742b7a0..077ff5fd 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -24,6 +24,9 @@ //! //! - `async`: Enables async server and client. //! - `sync`: Enables traditional sync server and client (default enabled). +//! - `security_extension`: Enables the connection extension framework (AcceptHook, +//! ConnectHook, PayloadTransform) for per-connection security policies +//! such as encryption and authentication. Unix only. //! //! # Socket address //! @@ -33,24 +36,31 @@ //! - `unix://@/run/some.sock`: Abstract Unix domain socket. //! - `vsock://vsock://8:1024`: [vsock](https://man7.org/linux/man-pages/man7/vsock.7.html). //! -//! For mscOS, ttrpc-rust **only** supports normal Unix domain socket: +//! For macOS, ttrpc-rust **only** supports normal Unix domain socket: //! //! - `unix:///run/some.sock`: Normal Unix domain socket. //! #![cfg_attr(docsrs, feature(doc_cfg))] +// The security_extension feature requires Unix platform. +#[cfg(all(not(unix), feature = "security_extension"))] +compile_error!("The 'security_extension' feature is only supported on Unix platforms."); + #[macro_use] extern crate log; #[macro_use] pub mod error; +#[cfg(feature = "sync")] #[macro_use] mod common; #[macro_use] mod macros; +pub mod security_extension; + pub mod context; pub mod proto; @@ -60,12 +70,27 @@ pub use self::proto::{Code, MessageHeader, Request, Response, Status}; #[doc(inline)] pub use crate::error::{get_status, Error, Result}; +// Core extension types are always available. +#[doc(inline)] +pub use crate::security_extension::{ConnectionData, ConnectionDataExt, PayloadTransform}; + +#[cfg(feature = "security_extension")] +#[doc(inline)] +pub use crate::security_extension::{ + AcceptHook, ConnectHook, ConnectionContext, HookError, HookOutput, +}; + +#[cfg(not(feature = "security_extension"))] +#[doc(hidden)] +pub use crate::security_extension::ConnectionContext; + cfg_sync! { pub mod sync; #[doc(hidden)] + #[allow(deprecated)] pub use sync::response_to_channel; #[doc(inline)] - pub use sync::{MethodHandler, TtrpcContext}; + pub use sync::{send_response, MethodHandler, TtrpcContext}; pub use sync::Client; #[doc(inline)] pub use sync::Server; diff --git a/src/proto.rs b/src/proto.rs index ff3e3f7c..d46e7326 100644 --- a/src/proto.rs +++ b/src/proto.rs @@ -189,6 +189,9 @@ impl MessageHeader { } /// Generic message of ttrpc. +/// +/// Constructed internally by the ttrpc framework and is not intended to be +/// built directly by downstream code. #[derive(Default, Debug, Clone, PartialEq, Eq)] pub struct GenMessage { pub header: MessageHeader, @@ -209,6 +212,32 @@ impl From for GenMessageError { #[cfg(feature = "async")] impl GenMessage { + /// Create a DATA message. + pub(crate) fn new_data(stream_id: u32, payload: Vec) -> Self { + Self { + header: MessageHeader::new_data(stream_id, payload.len() as u32), + payload, + } + } + + /// Create a RESPONSE message. + pub(crate) fn new_response(stream_id: u32, payload: Vec) -> Self { + Self { + header: MessageHeader::new_response(stream_id, payload.len() as u32), + payload, + } + } + + /// Create a DATA close message (FLAG_REMOTE_CLOSED | FLAG_NO_DATA). + pub(crate) fn new_close(stream_id: u32) -> Self { + let mut header = MessageHeader::new_data(stream_id, 0); + header.set_flags(FLAG_REMOTE_CLOSED | FLAG_NO_DATA); + Self { + header, + payload: Vec::new(), + } + } + /// Encodes a MessageHeader to writer. pub async fn write_to( &self, diff --git a/src/security_extension.rs b/src/security_extension.rs new file mode 100644 index 00000000..3dee9e4f --- /dev/null +++ b/src/security_extension.rs @@ -0,0 +1,1590 @@ +// Copyright 2026 Alibaba Cloud. All rights reserved. +// +// SPDX-License-Identifier: Apache-2.0 +// + +//! Connection Extension Framework. +//! +//! # Design Principle +//! +//! **ttrpc provides mechanism, not policy.** +//! +//! ttrpc is a lightweight RPC transport framework. It provides generic extension +//! points -- hooks and opaque data attachment -- for applications to implement +//! connection-level security, authorization, or any other cross-cutting concern. +//! ttrpc itself contains **zero** domain logic: no roles, no authorization +//! policies, no encryption algorithms, no identity types. +//! +//! All policy decisions and concrete implementations reside exclusively in +//! the application layer (e.g., kata-agent). +//! +//! # Extension Points +//! +//! ```text +//! +------------------------------+ +------------------------------+ +//! | ttrpc provides | | application decides | +//! | | | | +//! accept() | AcceptHook | | CID check, mTLS, token auth, | +//! | -> accept / reject | | or nothing -- app's choice | +//! | -> attach metadata | | | +//! | | | | +//! request | PayloadTransform | | AES-GCM, ChaCha20, zstd | +//! (in) | (optional, per-conn) | | compression, or no-op | +//! | | | | +//! dispatch | TtrpcContext | | | +//! | .connection_data | | Handler reads metadata, | +//! | | | app enforces its own authz | +//! | | | | +//! response | PayloadTransform | | AES-GCM, ChaCha20, zstd | +//! (out) | (optional, per-conn) | | compression, or no-op | +//! +------------------------------+ +------------------------------+ +//! ``` +//! +//! # What ttrpc Knows vs. Doesn't Know +//! +//! | ttrpc knows | ttrpc does NOT know | +//! |------------------------------------|---------------------------------| +//! | A hook exists and can reject conns | CID, roles, identities | +//! | Connections can carry opaque data | What the data means | +//! | Payloads can be transformed in/out | Encryption, compression, codecs | +//! | Handlers can read connection data | Authorization policies | +//! | Stream IDs, message framing | Business logic of any kind | +//! +//! # Backward Compatibility +//! +//! | Scenario | Behavior | +//! |----------------------------|---------------------------------------------------| +//! | No hook set | All connections accepted, empty data, no xform | +//! | Hook returns `Ok` | Connection accepted with app-defined metadata | +//! | Hook returns `Err` | Connection rejected and closed immediately | +//! | `payload_transform = None` | Plaintext pass-through (no overhead) | +//! | Existing handlers | `connection_data` is empty HashMap -- no breakage | +//! +//! # 10 Injection Points +//! +//! ttrpc has 10 internal injection points where the extension framework hooks +//! into the message processing pipeline. They ensure **no payload bypasses +//! the transform regardless of message type or direction**. +//! +//! | # | Side | Location | Direction | Message Type | +//! |----|--------|--------------------------|-----------|---------------------| +//! | 1 | server | `do_start()` accept | - | New connection | +//! | 2 | server | `handle_request()` | inbound | Unary REQUEST | +//! | 3 | server | `handle_msg()` DATA | route | Streaming DATA (1) | +//! | 4 | server | `handle_msg()` response | outbound | Unary response | +//! | 5 | client | `new_inner()` connect | - | New connection | +//! | 6 | client | `request()` send | outbound | Unary REQUEST | +//! | 7 | client | `handle_msg()` recv | inbound | Unary response | +//! | 8 | client | `new_stream()` send | outbound | Stream-init REQUEST | +//! | 9 | both | `StreamSender::send()` | outbound | Streaming DATA | +//! | 10 | both | connection reader | inbound | Streaming DATA | +//! +//! (1) Streaming DATA inbound transform is applied in the **connection reader** +//! (`handle_msg()`) before spawning a handler task. This ensures deterministic +//! nonce sequencing for stateful transforms (e.g., AEAD) when multiple streams +//! are active concurrently. `StreamReceiver::recv()` passes through the +//! already-decrypted payload without re-applying transform. +//! +//! **Note**: The `streaming_client = false` with non-empty initial payload +//! path is fully supported — `handle_stream()` sends the already-decrypted +//! payload directly, so `StreamReceiver::recv()` passes it through without +//! re-applying `transform_inbound`. This works for **any** `PayloadTransform` +//! implementation, including asymmetric transforms where +//! `transform_outbound(transform_inbound(x))` is not guaranteed to equal `x`. +//! +//! # Socket raw_fd +//! +//! [`Socket`](crate::asynchronous::transport::Socket) stores the underlying +//! file descriptor so [`AcceptHook`] can access it for: +//! - `getpeername()` to inspect peer identity (e.g., vsock CID) +//! - Bidirectional handshake I/O (e.g., ECDH key exchange) +//! +//! The fd is captured in the platform-specific `From` impls (vsock, unix, tcp) +//! before the socket is type-erased into the inner `Box`. + +// ── Always-compiled imports ── +use crate::error::Error; +use crate::proto::{check_oversize, MessageHeader}; +use std::sync::Arc; + +// ── Always-compiled core types ────────────────────────────────────────────── + +/// Helper to extract typed values from [`ConnectionData`]. +/// +/// The trait is always available; the concrete `ConnectionData` type and its +/// `ConnectionDataExt` impl live in the feature-gated sub-modules below. +pub trait ConnectionDataExt { + fn get_typed(&self, key: &str) -> Option<&T>; +} + +/// Optional bidirectional payload transform, attached per-connection. +/// +/// Applied by ttrpc at the wire boundary for **all** message types: +/// - `transform_inbound`: after reading from transport, before dispatching +/// - `transform_outbound`: after handler returns, before writing to transport +/// +/// Application-defined. Typical uses: +/// - Encryption/decryption (AES-256-GCM, ChaCha20-Poly1305) +/// - No-op (plaintext, the default) +/// +/// **Compression** should be performed at the application layer (before +/// serialization into `Request.payload`), not as a `PayloadTransform`. +/// This avoids payload-expansion edge cases and keeps the transform layer +/// focused on security with predictable, fixed-size overhead. +pub trait PayloadTransform: Send + Sync + std::fmt::Debug { + /// Decrypt / decode a payload. `aad` is the authenticated-but-unencrypted + /// header data (`stream_id || type_ || flags`); implementations using AEAD + /// ciphers should pass it as Additional Authenticated Data so that header + /// tampering is detected. + fn transform_inbound(&self, data: Vec, aad: &[u8]) -> std::result::Result, String>; + /// Encrypt / encode a payload. `aad` has the same semantics as + /// [`transform_inbound`](Self::transform_inbound). + fn transform_outbound(&self, data: Vec, aad: &[u8]) + -> std::result::Result, String>; + + /// Maximum number of bytes that `transform_outbound` may add to a payload. + /// + /// Used by the framework to pre-check whether a raw payload will fit + /// within [`MESSAGE_LENGTH_MAX`](crate::proto::MESSAGE_LENGTH_MAX) after + /// transform, *without* invoking the transform (which may advance + /// internal state such as AEAD nonce counters). + /// + /// The default (64) covers common AEAD schemes (AES-GCM: 28 bytes, + /// ChaCha20-Poly1305: 28 bytes) with comfortable margin. Override for + /// transforms with larger expansion. + fn max_overhead(&self) -> usize { + 64 + } +} + +/// Serialize the AAD portion of a message header: `stream_id(4B BE) || type_(1B) || flags(1B)`. +/// +/// `length` is intentionally excluded because it changes after transform +/// (e.g., AES-GCM adds nonce + tag). Only the immutable routing/control +/// fields are authenticated. +/// +/// Always compiled so that call sites (sync and async) need no cfg gates. +/// The noop `ConnectionContext` ignores `aad` when the feature is disabled, +/// making the computation harmless. +#[allow(dead_code)] +pub(crate) fn serialize_aad(header: &MessageHeader) -> [u8; 6] { + let mut buf = [0u8; 6]; + buf[0..4].copy_from_slice(&header.stream_id.to_be_bytes()); + buf[4] = header.type_; + buf[5] = header.flags; + buf +} + +// Re-export: both module variants define ConnectionContext and ConnectionData. +// Since only one `mod hooks` exists at compile time, no cfg gate needed. +pub use hooks::{ConnectionContext, ConnectionData}; + +// Feature-only items re-exported separately. +#[cfg(all(feature = "async", feature = "security_extension"))] +pub(crate) use hooks::ServerExtensionConfig; +#[cfg(feature = "security_extension")] +pub use hooks::{AcceptHook, ConnectHook, HookError, HookOutput}; + +// ── Feature-gated hook types ─────────────────────────────────────────────── +// +// All hook-related items live in this inner module behind a single cfg gate. +// Public re-exports below the module make them accessible as +// `crate::security_extension::HookError` etc. +#[cfg(feature = "security_extension")] +mod hooks { + use super::*; + #[cfg(feature = "async")] + use crate::asynchronous::transport::Socket; + use std::any::Any; + use std::collections::HashMap; + use std::fmt; + use std::os::unix::io::RawFd; + #[cfg(feature = "async")] + use std::os::unix::io::{AsRawFd as _, FromRawFd as _, OwnedFd}; + #[cfg(feature = "async")] + use tokio::{task, time::Duration}; + + /// Type-erased per-connection data store. + /// + /// Applications attach whatever types they need. ttrpc stores and forwards + /// them but never inspects the contents. + /// + /// `ConnectionData` is **immutable** after connection establishment. + /// It is wrapped in [`std::sync::Arc`] and shared read-only across all request + /// handlers on the connection. If the application needs mutable per-connection + /// state (e.g., request counters, rate-limit tracking), it should maintain that + /// externally, keyed by a connection identity (e.g., session_id). + pub type ConnectionData = HashMap>; + + impl super::ConnectionDataExt for ConnectionData { + fn get_typed(&self, key: &str) -> Option<&T> { + self.get(key)?.downcast_ref::() + } + } + + /// Framework-level timeout applied to [`AcceptHook::on_accept`] when invoked + /// from the async server's accept loop. The hook runs on the blocking + /// threadpool (`tokio::task::spawn_blocking`); if it does not complete within + /// this window, ttrpc treats the connection as failed and drops it. + /// + /// Hooks are expected to set their own tighter per-I/O timeouts (e.g., + /// `SO_RCVTIMEO` / `SO_SNDTIMEO`) for finer-grained control; this value is a + /// last-resort safety net against a misbehaving or malicious peer. + #[cfg(feature = "async")] + pub const ACCEPT_HOOK_TIMEOUT: Duration = Duration::from_secs(30); + + /// Structured error returned by [`AcceptHook::on_accept`] and [`ConnectHook::on_connect`]. + /// + /// On the server side, ttrpc logs the error and closes the connection. + /// On the client side, hook failures are propagated back to the caller + /// and the connection is not used. Structured errors enable proper audit + /// logging and monitoring. + pub enum HookError { + /// Application rejected the connection (e.g., bad CID, failed auth). + Rejected(String), + /// Handshake or I/O operation timed out. + Timeout, + /// I/O error during handshake. + Io(std::io::Error), + /// Other error. + Other(String), + } + + impl fmt::Debug for HookError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + HookError::Rejected(msg) => write!(f, "Rejected({})", msg), + HookError::Timeout => write!(f, "Timeout"), + HookError::Io(e) => write!(f, "Io({:?})", e), + HookError::Other(msg) => write!(f, "Other({})", msg), + } + } + } + + impl fmt::Display for HookError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + HookError::Rejected(msg) => write!(f, "connection rejected: {}", msg), + HookError::Timeout => write!(f, "handshake timeout"), + HookError::Io(e) => write!(f, "I/O error: {}", e), + HookError::Other(msg) => write!(f, "{}", msg), + } + } + } + + /// Output produced by [`AcceptHook`] for an accepted connection. + /// + /// # Data Propagation Path + /// + /// ```text + /// accept(fd) + /// | + /// +-> AcceptHook::on_accept(fd) -> HookOutput { data, payload_transform } + /// | + /// +-> ServerReader { conn_data, payload_transform: Arc<...> } + /// | | + /// | +-- handle_msg(msg) + /// | | +- MESSAGE_TYPE_REQUEST -> handle_request() + /// | | | +-> transform_inbound(msg.payload) <- unary request (#2) + /// | | | + /// | | +- MESSAGE_TYPE_DATA -> transform_inbound -> route to stream_rx + /// | | | (wire order, #10) + /// | | +-> StreamReceiver::recv() + /// | | +-> pass-through (already decrypted) + /// | | + /// | +-- handle_method(method, req) + /// | | +- TtrpcContext { connection_data: Arc } + /// | | | +-> handler reads ctx.connection_data.get_typed::(key) + /// | | | + /// | | +-> transform_outbound(response.payload) <- unary response (#4) + /// | | + /// | +-- handle_stream(stream, req) + /// | +-> StreamSender { payload_transform: Arc<...> } + /// | +-> transform_outbound(buf) on .send() <- streaming data (#9) + /// | + /// +-> all responses/data written to transport + /// ``` + /// + /// **Important**: Streaming DATA inbound transform is applied in the + /// connection reader (`handle_msg()`) before spawning a handler task, + /// ensuring deterministic nonce sequencing for stateful transforms. + /// `StreamReceiver::recv()` passes through the already-decrypted payload. + pub struct HookOutput { + /// Opaque metadata attached to this connection for its lifetime. + /// Forwarded to method/stream handlers via `TtrpcContext.connection_data`. + /// This data is immutable after `on_accept()` returns. + pub data: ConnectionData, + + /// Optional payload transform for this connection. + /// `None` = pass-through (no transformation). + /// Stored as `Arc` internally so it can be shared with stream handlers. + pub payload_transform: Option>, + } + + /// Called once per accepted connection. + /// + /// The hook receives the raw fd of the accepted connection and returns: + /// - `Ok(HookOutput)`: accept the connection, attach metadata and optional transform + /// - `Err(HookError)`: reject the connection; ttrpc closes it immediately + /// + /// The hook may inspect the peer address, perform a handshake, look up a + /// certificate, or do anything else the application requires. + /// + /// The hook has exclusive use of the fd -- ttrpc has not started its message + /// loop yet. The hook may read/write on the fd for handshake purposes. + /// + /// # Async server execution model + /// + /// On the async server, each accepted connection's hook is dispatched via + /// [`tokio::task::spawn_blocking`] and wrapped in a framework-level timeout + /// (`ACCEPT_HOOK_TIMEOUT`). This guarantees: + /// - The accept-loop worker is never blocked by a slow handshake, so other + /// connections keep being accepted in parallel. + /// - Current-thread runtimes cannot deadlock (the hook runs on a separate + /// blocking thread). + /// - A malicious peer cannot stall the accept loop beyond the timeout. + /// + /// Hooks are still expected to set their own per-I/O timeouts + /// (`SO_RCVTIMEO` / `SO_SNDTIMEO`) for finer-grained control; the + /// framework-level timeout is a last-resort safety net. + pub trait AcceptHook: Send + Sync + std::fmt::Debug { + fn on_accept(&self, fd: RawFd) -> std::result::Result; + } + + impl AcceptHook for Box { + fn on_accept(&self, fd: RawFd) -> std::result::Result { + (**self).on_accept(fd) + } + } + + /// Client-side connection hook, symmetric to [`AcceptHook`]. + /// + /// Called after a client connection is established, before RPC messages are sent. + /// The hook receives the raw fd and can perform: + /// - Peer identity verification + /// - Handshake I/O (e.g., ECDH client-side) + /// - Attaching connection metadata + /// - Installing a [`PayloadTransform`] + /// + /// If the hook returns an error, [`Client::with_hook`](crate::asynchronous::Client::with_hook) + /// returns the error and the connection is **not** used, preventing a silent + /// downgrade to an untransformed connection. + /// + /// **WARNING**: The hook **must** set I/O timeouts before any read/write on the fd. + pub trait ConnectHook: Send + Sync + std::fmt::Debug { + fn on_connect(&self, fd: RawFd) -> std::result::Result; + } + + impl ConnectHook for Box { + fn on_connect(&self, fd: RawFd) -> std::result::Result { + (**self).on_connect(fd) + } + } + + /// Server-wide extension configuration. Immutable after server build. + /// + /// Holds server-scoped hooks (currently only the [`AcceptHook`]) that apply to + /// every accepted connection. Shared via `Arc` from + /// `Server` down to the accept loop. + /// + /// Per-connection state (data, payload transform) lives in + /// [`ConnectionContext`], which is constructed from the hook output after a + /// connection is accepted. + #[cfg(feature = "async")] + #[derive(Debug, Default)] + pub(crate) struct ServerExtensionConfig { + /// Server-side hook called on each accepted connection. + pub(crate) accept_hook: Option>, + } + + #[cfg(feature = "async")] + impl ServerExtensionConfig { + /// Invoke the accept hook for a newly accepted connection. + /// + /// The (synchronous) hook runs on tokio's blocking threadpool via + /// [`task::spawn_blocking`] and is wrapped in a framework-level + /// [`ACCEPT_HOOK_TIMEOUT`]. This keeps the async accept loop responsive: + /// - The runtime worker that accepts new connections is never blocked by + /// a slow handshake. + /// - A current-thread runtime cannot deadlock (the hook runs on a + /// separate blocking thread). + /// - A malicious peer cannot stall the accept loop beyond the timeout. + /// + /// Returns: + /// - `Ok(None)`: no hook set → accept with empty defaults + /// - `Ok(Some(output))`: hook accepted → use hook output + /// - `Err(e)`: hook rejected the connection, timed out, or could not run + pub(crate) async fn on_accept( + &self, + conn: &Socket, + ) -> Result, HookError> { + let hook = match self.accept_hook { + Some(ref h) => Arc::clone(h), + None => return Ok(None), + }; + let fd = conn.as_raw_fd().ok_or_else(|| { + HookError::Other( + "accept hook requires a raw fd; use Server::bind() or construct Socket from a platform stream (UnixStream, TcpStream, VsockStream)".to_string(), + ) + })?; + // Dup the fd so the blocking task owns its own copy. This prevents + // fd reuse: if the outer task drops the connection on timeout, the + // hook's dup'd fd remains valid (and is not reassigned by the OS) + // until the hook finishes and drops it. + let owned_fd: OwnedFd = unsafe { + let dup_fd = libc::fcntl(fd, libc::F_DUPFD_CLOEXEC, 0); + if dup_fd < 0 { + return Err(HookError::Io(std::io::Error::last_os_error())); + } + OwnedFd::from_raw_fd(dup_fd) + }; + // spawn_blocking moves the hook execution off the async worker so + // other connections keep being accepted. A slow or malicious peer + // cannot stall the accept loop beyond `ACCEPT_HOOK_TIMEOUT`. + let mut handle = task::spawn_blocking(move || { + let hook_fd = owned_fd.as_raw_fd(); + // owned_fd stays alive for the closure's lifetime, ensuring + // the dup'd fd is not reused by the OS until the hook returns. + let _guard = owned_fd; + hook.on_accept(hook_fd) + }); + tokio::select! { + result = &mut handle => { + match result { + Ok(Ok(output)) => Ok(Some(output)), + Ok(Err(e)) => Err(e), + Err(join_err) => Err(HookError::Other(format!( + "accept hook task failed: {}", + join_err + ))), + } + } + _ = 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) + } + } + } + } + + /// Bundles all per-connection extension data into a single propagation unit. + /// + /// By passing a single `Arc` through client and server + /// internals, we avoid scattering feature gates across every struct and + /// function. Server-wide configuration (e.g., the accept hook) lives in + /// `ServerExtensionConfig`. + #[derive(Debug, Default)] + pub struct ConnectionContext { + /// Opaque per-connection metadata. Immutable after accept. + pub data: Arc, + /// Optional payload transform. `None` = pass-through. + pub payload_transform: Option>, + /// Serializes outbound transform + enqueue to prevent wire-order + /// races with stateful transforms (e.g., AEAD nonce counters). + #[cfg(feature = "async")] + async_outbound_lock: std::sync::Arc>, + /// Sync-path outbound lock (analogous to `async_outbound_lock` for tokio paths). + #[cfg(feature = "sync")] + sync_outbound_lock: std::sync::Mutex<()>, + } + + impl ConnectionContext { + /// Create a context from hook output. No hook → empty defaults. + pub fn new(output: Option) -> Self { + match output { + Some(o) => Self { + data: Arc::new(o.data), + payload_transform: o.payload_transform.map(Arc::from), + #[cfg(feature = "async")] + async_outbound_lock: Arc::new(tokio::sync::Mutex::new(())), + #[cfg(feature = "sync")] + sync_outbound_lock: std::sync::Mutex::new(()), + }, + None => Self::default(), + } + } + + /// Maximum raw (pre-transform) payload size that is guaranteed to fit + /// within [`MESSAGE_LENGTH_MAX`](crate::proto::MESSAGE_LENGTH_MAX) + /// after `transform_outbound`. + /// + /// Use this to pre-check payload size *before* invoking the transform, + /// avoiding state advancement (e.g., AEAD nonce counters) on payloads + /// that would be rejected post-transform. + pub fn max_raw_payload_len(&self) -> usize { + match self.payload_transform { + Some(ref xform) => { + crate::proto::MESSAGE_LENGTH_MAX.saturating_sub(xform.max_overhead()) + } + None => crate::proto::MESSAGE_LENGTH_MAX, + } + } + + /// Apply inbound payload transform (if any). + /// Returns the transformed payload, or the original if no transform is set. + pub fn transform_inbound(&self, payload: Vec, aad: &[u8]) -> Result, String> { + match self.payload_transform { + Some(ref xform) => xform.transform_inbound(payload, aad), + None => Ok(payload), + } + } + + /// Apply outbound payload transform (if any). + /// Returns the transformed payload, or the original if no transform is set. + pub fn transform_outbound(&self, payload: Vec, aad: &[u8]) -> Result, String> { + match self.payload_transform { + Some(ref xform) => xform.transform_outbound(payload, aad), + None => Ok(payload), + } + } + + /// Inbound pipeline on a raw buffer: apply payload transform and enforce + /// the post-transform size limit. + /// + /// Intended for sync paths that shuttle a `Vec` rather than a full + /// `GenMessage` (e.g., the sync server reader and sync client sender / + /// receiver threads). For message-level pipelines that also update + /// `header.length`, use [`ConnectionContext::inbound`]. + /// + /// `rpc_error` selects the error flavor (see [`ConnectionContext::inbound`]). + pub fn inbound_buf( + &self, + data: Vec, + aad: &[u8], + rpc_error: bool, + ) -> Result, Error> { + let transformed = match self.payload_transform { + Some(ref xform) => xform + .transform_inbound(data, aad) + .map_err(|e| Error::Others(format!("transform_inbound failed: {}", e)))?, + None => data, + }; + check_oversize(transformed.len(), rpc_error)?; + Ok(transformed) + } + + /// Outbound pipeline on a raw buffer: apply payload transform and enforce + /// the post-transform size limit. + /// + /// Counterpart to [`ConnectionContext::inbound_buf`] for the send path. + /// `rpc_error` selects the error flavor (see [`ConnectionContext::inbound`]). + pub fn outbound_buf( + &self, + data: Vec, + aad: &[u8], + rpc_error: bool, + ) -> Result, Error> { + let transformed = match self.payload_transform { + Some(ref xform) => { + let max_len = self.max_raw_payload_len(); + if data.len() > max_len { + return Err(Error::Others(format!( + "payload {} bytes exceeds safe limit {} bytes \ + (MESSAGE_LENGTH_MAX - transform overhead)", + data.len(), + max_len + ))); + } + xform + .transform_outbound(data, aad) + .map_err(|e| Error::Others(format!("transform_outbound failed: {}", e)))? + } + None => data, + }; + check_oversize(transformed.len(), rpc_error)?; + Ok(transformed) + } + + /// Inbound pipeline: apply payload transform in-place and enforce the + /// post-transform size limit on a `GenMessage`. + /// + /// Updates `msg.header.length` to match the transformed payload. Then + /// delegates to `check_oversize` for the Layer 2 size + /// guard, which rejects payloads that expanded past `MESSAGE_LENGTH_MAX` + /// (e.g., from a decompression transform). + /// + /// The header's routing fields (`stream_id`, `type_`, `flags`) are + /// serialized and passed as AAD to the transform for authentication. + /// + /// `rpc_error` selects the error flavor returned by the size check: + /// - `true` (server paths): returns `Error::RpcStatus(INVALID_ARGUMENT)` + /// - `false` (client paths): returns `Error::Others(...)` + pub fn inbound( + &self, + msg: &mut crate::proto::GenMessage, + rpc_error: bool, + ) -> Result<(), Error> { + if let Some(ref xform) = self.payload_transform { + let aad = serialize_aad(&msg.header); + msg.payload = xform + .transform_inbound(std::mem::take(&mut msg.payload), &aad) + .map_err(|e| Error::Others(format!("transform_inbound failed: {}", e)))?; + msg.header.length = msg.payload.len() as u32; + } + check_oversize(msg.payload.len(), rpc_error) + } + + /// Outbound pipeline: apply payload transform in-place and enforce the + /// post-transform size limit on a `GenMessage`. + /// + /// Updates `msg.header.length` to match the transformed payload. Then + /// delegates to `check_oversize` for the Layer 2 size + /// guard. + /// + /// The header's routing fields (`stream_id`, `type_`, `flags`) are + /// serialized and passed as AAD to the transform for authentication. + /// + /// `rpc_error` selects the error flavor returned by the size check + /// (see [`ConnectionContext::inbound`] for semantics). + pub fn outbound( + &self, + msg: &mut crate::proto::GenMessage, + rpc_error: bool, + ) -> Result<(), Error> { + if let Some(ref xform) = self.payload_transform { + let max_len = self.max_raw_payload_len(); + if msg.payload.len() > max_len { + return Err(Error::Others(format!( + "payload {} bytes exceeds safe limit {} bytes \ + (MESSAGE_LENGTH_MAX - transform overhead); \ + reduce payload size", + msg.payload.len(), + max_len + ))); + } + let aad = serialize_aad(&msg.header); + msg.payload = xform + .transform_outbound(std::mem::take(&mut msg.payload), &aad) + .map_err(|e| Error::Others(format!("transform_outbound failed: {}", e)))?; + msg.header.length = msg.payload.len() as u32; + } + check_oversize(msg.payload.len(), rpc_error) + } + + /// Atomically: reserve capacity → lock outbound → transform → send. + /// + /// Reserves channel capacity first (cancellable await), then acquires + /// the outbound lock and transforms (non-cancellable). This prevents + /// nonce desynchronization if the future is cancelled after advancing + /// the stateful transform but before the frame reaches the channel. + /// + /// When `await_ack` is true, waits for the writer task to confirm the + /// frame reached the transport (needed for streaming paths that depend + /// on write success before updating state). When false, returns after + /// enqueue — suitable for unary request/response paths where the caller + /// applies its own timeout on the reply. + #[cfg(feature = "async")] + pub async fn transform_send( + &self, + msg: &mut crate::proto::GenMessage, + tx: &tokio::sync::mpsc::Sender, + rpc_error: bool, + await_ack: bool, + ) -> Result<(), Error> { + // Reserve capacity first — this is the only cancellable await point. + // If cancelled here, no nonce has been advanced. + let permit = tx + .reserve() + .await + .map_err(|e| Error::Others(format!("reserve channel capacity failed: {e}")))?; + + let _guard = self.async_outbound_lock.lock().await; + self.outbound(msg, rpc_error)?; + let taken = std::mem::take(msg); + + // From here on: no await until the frame is in the channel. + // permit.send() is synchronous — cannot be cancelled. + if await_ack { + let (result_tx, result_rx) = tokio::sync::oneshot::channel(); + permit.send(crate::asynchronous::SendingMessage::new_with_result( + taken, result_tx, + )); + drop(_guard); + result_rx + .await + .map_err(|_| Error::Others("writer task dropped result channel".to_string()))? + } else { + permit.send(crate::asynchronous::SendingMessage::new(taken)); + Ok(()) + } + } + + /// Sync-path equivalent of [`transform_send`](Self::transform_send). + /// Serializes `outbound_buf` + `tx.send` under a std mutex to prevent + /// nonce ordering races across concurrent handler threads. + #[cfg(feature = "sync")] + pub fn send_response_sync( + &self, + mut buf: Vec, + aad: &[u8], + tx: &std::sync::mpsc::Sender<(crate::proto::MessageHeader, Vec)>, + mh: crate::proto::MessageHeader, + ) -> Result<(), Error> { + let _guard = self.sync_outbound_lock.lock().unwrap(); + buf = self.outbound_buf(buf, aad, true)?; + let mh = crate::proto::MessageHeader { + length: buf.len() as u32, + ..mh + }; + tx.send((mh, buf)) + .map_err(|e| Error::Others(format!("send to wire channel failed: {e}"))) + } + } +} // mod hooks + +// ── Noop types (feature disabled) ────────────────────────────────────────── +// +// When security_extension is disabled, these types replace the real ones with +// minimal-cost stubs (an Arc is still allocated per connection +// for API uniformity, but its fields are unused). All items inside have no cfg gates. +#[cfg(not(feature = "security_extension"))] +mod hooks { + use super::*; + + /// Zero-sized placeholder when `security_extension` is disabled. + /// `Arc` costs ~24 bytes (ArcInner header only, no HashMap). + #[derive(Clone, Debug, Default)] + pub struct ConnectionData; + + impl super::ConnectionDataExt for ConnectionData { + fn get_typed(&self, _key: &str) -> Option<&T> { + None + } + } + + /// No-op connection context used when `security_extension` is disabled. + /// + /// All pipeline methods (`inbound`, `outbound`, `inbound_buf`, `outbound_buf`) + /// skip payload transform and only enforce the message size limit. + /// + /// The real [`ConnectionContext`] (available with `--features security_extension`) + /// adds per-connection data and payload encryption on top of the same method + /// signatures. + #[derive(Clone, Debug, Default)] + #[doc(hidden)] + pub struct ConnectionContext { + /// Always empty when security_extension is disabled. + pub data: Arc, + /// Always `None` when security_extension is disabled. + pub payload_transform: Option>, + } + + #[allow(dead_code)] + impl ConnectionContext { + /// No-op transform always reports the full limit (no overhead). + pub fn max_raw_payload_len(&self) -> usize { + crate::proto::MESSAGE_LENGTH_MAX + } + + /// Inbound pipeline: size-check only (no transform). + pub fn inbound( + &self, + msg: &mut crate::proto::GenMessage, + rpc_error: bool, + ) -> Result<(), Error> { + check_oversize(msg.payload.len(), rpc_error) + } + + /// Outbound pipeline: size-check only (no transform). + pub fn outbound( + &self, + msg: &mut crate::proto::GenMessage, + rpc_error: bool, + ) -> Result<(), Error> { + check_oversize(msg.payload.len(), rpc_error) + } + + /// Inbound buffer pipeline: size-check only (no transform). + pub fn inbound_buf( + &self, + data: Vec, + _aad: &[u8], + rpc_error: bool, + ) -> Result, Error> { + check_oversize(data.len(), rpc_error)?; + Ok(data) + } + + /// Outbound buffer pipeline: size-check only (no transform). + pub fn outbound_buf( + &self, + data: Vec, + _aad: &[u8], + rpc_error: bool, + ) -> Result, Error> { + check_oversize(data.len(), rpc_error)?; + Ok(data) + } + + /// Size-check + enqueue (no transform, no lock when security_extension is disabled). + #[cfg(feature = "async")] + pub async fn transform_send( + &self, + msg: &mut crate::proto::GenMessage, + tx: &tokio::sync::mpsc::Sender, + rpc_error: bool, + await_ack: bool, + ) -> Result<(), Error> { + self.outbound(msg, rpc_error)?; + let permit = tx + .reserve() + .await + .map_err(|e| Error::Others(format!("reserve channel capacity failed: {e}")))?; + let taken = std::mem::take(msg); + if await_ack { + let (result_tx, result_rx) = tokio::sync::oneshot::channel(); + permit.send(crate::asynchronous::SendingMessage::new_with_result( + taken, result_tx, + )); + result_rx + .await + .map_err(|_| Error::Others("writer task dropped result channel".to_string()))? + } else { + permit.send(crate::asynchronous::SendingMessage::new(taken)); + Ok(()) + } + } + + /// Sync-path size-check + enqueue (no transform, no lock needed when + /// security_extension is disabled — no stateful transforms exist). + #[cfg(feature = "sync")] + pub fn send_response_sync( + &self, + buf: Vec, + _aad: &[u8], + tx: &std::sync::mpsc::Sender<(crate::proto::MessageHeader, Vec)>, + mh: crate::proto::MessageHeader, + ) -> Result<(), Error> { + check_oversize(buf.len(), true)?; + let mh = crate::proto::MessageHeader { + length: buf.len() as u32, + ..mh + }; + tx.send((mh, buf)) + .map_err(|e| Error::Others(format!("send to wire channel failed: {e}"))) + } + } +} // mod hooks + +#[cfg(all(test, feature = "security_extension"))] +mod tests { + use super::*; + #[cfg(all(feature = "async", feature = "security_extension"))] + use crate::asynchronous::transport::Socket; + use crate::proto::GenMessage; + #[cfg(feature = "async")] + use std::os::unix::io::RawFd; + + // ── ConnectionData + ConnectionDataExt ──────────────────────────────── + + #[test] + fn connection_data_empty() { + let data = ConnectionData::new(); + assert!(data.is_empty()); + assert_eq!(data.get_typed::("missing"), None); + assert_eq!(data.get_typed::("missing"), None); + } + + #[test] + fn connection_data_insert_and_get_typed() { + let mut data = ConnectionData::new(); + data.insert("name".into(), Box::new(String::from("alice"))); + data.insert("count".into(), Box::new(42u64)); + data.insert("flag".into(), Box::new(true)); + + assert_eq!( + data.get_typed::("name").map(String::as_str), + Some("alice") + ); + assert_eq!(data.get_typed::("count"), Some(&42u64)); + assert_eq!(data.get_typed::("flag"), Some(&true)); + } + + #[test] + fn connection_data_wrong_type_returns_none() { + let mut data = ConnectionData::new(); + data.insert("key".into(), Box::new(String::from("hello"))); + // Request as u64 — wrong type, should return None + assert_eq!(data.get_typed::("key"), None); + assert_eq!(data.get_typed::("key"), None); + } + + #[test] + fn connection_data_missing_key_returns_none() { + let mut data = ConnectionData::new(); + data.insert("a".into(), Box::new(1u32)); + assert_eq!(data.get_typed::("b"), None); + } + + // ── HookError ────────────────────────────────────────────────────── + + #[test] + fn accept_error_debug_rejected() { + let e = HookError::Rejected("bad cid".into()); + assert_eq!(format!("{:?}", e), "Rejected(bad cid)"); + } + + #[test] + fn accept_error_debug_timeout() { + let e = HookError::Timeout; + assert_eq!(format!("{:?}", e), "Timeout"); + } + + #[test] + fn accept_error_debug_io() { + let e = HookError::Io(std::io::Error::new(std::io::ErrorKind::BrokenPipe, "pipe")); + let dbg = format!("{:?}", e); + assert!(dbg.starts_with("Io(")); + } + + #[test] + fn accept_error_debug_other() { + let e = HookError::Other("misc".into()); + assert_eq!(format!("{:?}", e), "Other(misc)"); + } + + #[test] + fn accept_error_display_rejected() { + let e = HookError::Rejected("no auth".into()); + assert_eq!(format!("{}", e), "connection rejected: no auth"); + } + + #[test] + fn accept_error_display_timeout() { + let e = HookError::Timeout; + assert_eq!(format!("{}", e), "handshake timeout"); + } + + #[test] + fn accept_error_display_io() { + let e = HookError::Io(std::io::Error::new( + std::io::ErrorKind::ConnectionRefused, + "refused", + )); + let s = format!("{}", e); + assert!(s.starts_with("I/O error:")); + } + + #[test] + fn accept_error_display_other() { + let e = HookError::Other("something went wrong".into()); + assert_eq!(format!("{}", e), "something went wrong"); + } + + // ── HookOutput ───────────────────────────────────────────────────── + + #[test] + fn accept_output_empty_data_no_transform() { + let out = HookOutput { + data: ConnectionData::new(), + payload_transform: None, + }; + assert!(out.data.is_empty()); + assert!(out.payload_transform.is_none()); + } + + #[test] + fn accept_output_with_data_and_transform() { + let mut data = ConnectionData::new(); + data.insert("peer".into(), Box::new(1234u32)); + + let out = HookOutput { + data, + payload_transform: Some(Box::new(NoopTransform)), + }; + assert_eq!(out.data.get_typed::("peer"), Some(&1234u32)); + assert!(out.payload_transform.is_some()); + } + + // ── Helper: NoopTransform ──────────────────────────────────────────── + + #[derive(Debug)] + struct NoopTransform; + + impl PayloadTransform for NoopTransform { + fn transform_inbound(&self, data: Vec, _aad: &[u8]) -> Result, String> { + Ok(data) + } + fn transform_outbound(&self, data: Vec, _aad: &[u8]) -> Result, String> { + Ok(data) + } + } + + // ── ConnectionContext construction ─────────────────────────────────── + + #[test] + fn context_default_is_empty() { + let ctx = ConnectionContext::default(); + assert!(ctx.data.is_empty()); + assert!(ctx.payload_transform.is_none()); + } + + #[test] + fn context_new_none_equals_default() { + let ctx = ConnectionContext::new(None); + assert!(ctx.data.is_empty()); + assert!(ctx.payload_transform.is_none()); + } + + #[test] + fn context_new_with_data_only() { + let mut data = ConnectionData::new(); + data.insert("role".into(), Box::new(String::from("server"))); + let output = HookOutput { + data, + payload_transform: None, + }; + let ctx = ConnectionContext::new(Some(output)); + assert_eq!( + ctx.data.get_typed::("role").map(String::as_str), + Some("server") + ); + assert!(ctx.payload_transform.is_none()); + } + + #[test] + fn context_new_with_data_and_transform() { + let mut data = ConnectionData::new(); + data.insert("cid".into(), Box::new(42u32)); + let output = HookOutput { + data, + payload_transform: Some(Box::new(NoopTransform)), + }; + let ctx = ConnectionContext::new(Some(output)); + assert_eq!(ctx.data.get_typed::("cid"), Some(&42u32)); + assert!(ctx.payload_transform.is_some()); + } + + #[test] + fn context_debug_format() { + let ctx = ConnectionContext::default(); + let dbg = format!("{:?}", ctx); + assert!(dbg.contains("ConnectionContext")); + } + + // ── ConnectionContext transform pass-through (no transform) ────────── + + #[test] + fn transform_inbound_no_transform_returns_original() { + let ctx = ConnectionContext::default(); + let payload = vec![1, 2, 3, 4]; + let result = ctx.transform_inbound(payload.clone(), &[]).unwrap(); + assert_eq!(result, payload); + } + + #[test] + fn transform_outbound_no_transform_returns_original() { + let ctx = ConnectionContext::default(); + let payload = vec![5, 6, 7, 8]; + let result = ctx.transform_outbound(payload.clone(), &[]).unwrap(); + assert_eq!(result, payload); + } + + #[test] + fn inbound_no_transform_is_noop() { + let ctx = ConnectionContext::default(); + let mut msg = GenMessage { + payload: vec![10, 20, 30], + ..Default::default() + }; + ctx.inbound(&mut msg, false).unwrap(); + assert_eq!(msg.payload, vec![10, 20, 30]); + } + + #[test] + fn outbound_no_transform_is_noop() { + let ctx = ConnectionContext::default(); + let mut msg = GenMessage { + payload: vec![10, 20, 30], + ..Default::default() + }; + ctx.outbound(&mut msg, false).unwrap(); + assert_eq!(msg.payload, vec![10, 20, 30]); + } + + // ── ConnectionContext with NoopTransform ───────────────────────────── + + #[test] + fn transform_inbound_noop_preserves_data() { + let output = HookOutput { + data: ConnectionData::new(), + payload_transform: Some(Box::new(NoopTransform)), + }; + let ctx = ConnectionContext::new(Some(output)); + let payload = vec![0xDE, 0xAD, 0xBE, 0xEF]; + let result = ctx.transform_inbound(payload.clone(), &[]).unwrap(); + assert_eq!(result, payload); + } + + #[test] + fn transform_outbound_noop_preserves_data() { + let output = HookOutput { + data: ConnectionData::new(), + payload_transform: Some(Box::new(NoopTransform)), + }; + let ctx = ConnectionContext::new(Some(output)); + let payload = vec![0xCA, 0xFE]; + let result = ctx.transform_outbound(payload.clone(), &[]).unwrap(); + assert_eq!(result, payload); + } + + // ── XOR-0xA5A5 mock encryption ────────────────────────────────────── + // + // Wire format: [algo_id: u16 BE][payload_len: u32 BE][encrypted_data...] + // Header = 6 bytes. algo_id = 0x0001 for XOR-0xA5A5. + // Encryption: XOR each u16 (big-endian) chunk with 0xA5A5. + // Odd-length payloads are padded with 0x00 before encryption; + // decryption truncates to original payload_len. + + const XOR_ALGO_ID: u16 = 0x0001; + const XOR_KEY: u16 = 0xA5A5; + const XOR_HEADER_LEN: usize = 6; // 2 (algo_id) + 4 (payload_len) + + #[derive(Debug)] + struct XorPayloadTransform; + + impl PayloadTransform for XorPayloadTransform { + /// Encrypt: prepend header, XOR payload with 0xA5A5 in u16 chunks. + fn transform_outbound(&self, data: Vec, _aad: &[u8]) -> Result, String> { + let payload_len = data.len(); + // Pad to even length for u16 XOR + let mut padded = data; + if padded.len() % 2 != 0 { + padded.push(0x00); + } + // XOR encrypt in u16 big-endian chunks + let mut encrypted = Vec::with_capacity(XOR_HEADER_LEN + padded.len()); + // Header: algo_id (u16 BE) + payload_len (u32 BE) + encrypted.extend_from_slice(&XOR_ALGO_ID.to_be_bytes()); + encrypted.extend_from_slice(&(payload_len as u32).to_be_bytes()); + // Encrypted body + for chunk in padded.chunks(2) { + let val = u16::from_be_bytes([chunk[0], chunk[1]]); + let xored = val ^ XOR_KEY; + encrypted.extend_from_slice(&xored.to_be_bytes()); + } + Ok(encrypted) + } + + /// Decrypt: verify header, XOR encrypted data with 0xA5A5, truncate. + fn transform_inbound(&self, data: Vec, _aad: &[u8]) -> Result, String> { + if data.len() < XOR_HEADER_LEN { + return Err(format!( + "xor: packet too short ({} < {})", + data.len(), + XOR_HEADER_LEN + )); + } + let algo_id = u16::from_be_bytes([data[0], data[1]]); + if algo_id != XOR_ALGO_ID { + return Err(format!("xor: unknown algo 0x{:04X}", algo_id)); + } + let payload_len = u32::from_be_bytes([data[2], data[3], data[4], data[5]]) as usize; + let encrypted = &data[XOR_HEADER_LEN..]; + if encrypted.len() % 2 != 0 { + return Err("xor: encrypted data has odd length".into()); + } + // Decrypt: XOR with same key (symmetric) + let mut decrypted = Vec::with_capacity(encrypted.len()); + for chunk in encrypted.chunks(2) { + let val = u16::from_be_bytes([chunk[0], chunk[1]]); + let xored = val ^ XOR_KEY; + decrypted.extend_from_slice(&xored.to_be_bytes()); + } + // Truncate to original payload length + decrypted.truncate(payload_len); + Ok(decrypted) + } + } + + /// Helper: build a ConnectionContext with XorPayloadTransform. + fn xor_context() -> ConnectionContext { + let output = HookOutput { + data: ConnectionData::new(), + payload_transform: Some(Box::new(XorPayloadTransform)), + }; + ConnectionContext::new(Some(output)) + } + + // ── XOR: roundtrip tests ──────────────────────────────────────────── + + #[test] + fn xor_roundtrip_basic() { + let ctx = xor_context(); + let original = vec![0x01, 0x02, 0x03, 0x04, 0x05, 0x06]; + let encrypted = ctx.transform_outbound(original.clone(), &[]).unwrap(); + let decrypted = ctx.transform_inbound(encrypted, &[]).unwrap(); + assert_eq!(decrypted, original); + } + + #[test] + fn xor_roundtrip_empty_payload() { + let ctx = xor_context(); + let original = vec![]; + let encrypted = ctx.transform_outbound(original.clone(), &[]).unwrap(); + // Header only, no body + assert_eq!(encrypted.len(), XOR_HEADER_LEN); + let decrypted = ctx.transform_inbound(encrypted, &[]).unwrap(); + assert_eq!(decrypted, original); + } + + #[test] + fn xor_roundtrip_odd_length() { + let ctx = xor_context(); + let original = vec![0xAA, 0xBB, 0xCC]; // 3 bytes + let encrypted = ctx.transform_outbound(original.clone(), &[]).unwrap(); + // Header (6) + 2 u16 chunks (4 bytes, padded) + assert_eq!(encrypted.len(), XOR_HEADER_LEN + 4); + let decrypted = ctx.transform_inbound(encrypted, &[]).unwrap(); + assert_eq!(decrypted, original); + } + + #[test] + fn xor_roundtrip_single_byte() { + let ctx = xor_context(); + let original = vec![0xFF]; + let encrypted = ctx.transform_outbound(original.clone(), &[]).unwrap(); + let decrypted = ctx.transform_inbound(encrypted, &[]).unwrap(); + assert_eq!(decrypted, original); + } + + #[test] + fn xor_roundtrip_large_payload() { + let ctx = xor_context(); + let original: Vec = (0..1024).map(|i| (i % 256) as u8).collect(); + let encrypted = ctx.transform_outbound(original.clone(), &[]).unwrap(); + let decrypted = ctx.transform_inbound(encrypted, &[]).unwrap(); + assert_eq!(decrypted, original); + } + + // ── XOR: header verification ───────────────────────────────────────── + + #[test] + fn xor_header_algo_id() { + let xform = XorPayloadTransform; + let encrypted = xform.transform_outbound(vec![0x00, 0x01], &[]).unwrap(); + let algo_id = u16::from_be_bytes([encrypted[0], encrypted[1]]); + assert_eq!(algo_id, XOR_ALGO_ID); + } + + #[test] + fn xor_header_payload_length() { + let xform = XorPayloadTransform; + let payload = vec![0x10, 0x20, 0x30, 0x40, 0x50]; // 5 bytes + let encrypted = xform.transform_outbound(payload, &[]).unwrap(); + let recorded_len = + u32::from_be_bytes([encrypted[2], encrypted[3], encrypted[4], encrypted[5]]) as usize; + assert_eq!(recorded_len, 5); + } + + #[test] + fn xor_encrypted_data_differs_from_plaintext() { + let xform = XorPayloadTransform; + let original = vec![0x00, 0x00, 0x00, 0x00]; // all zeros + let encrypted = xform.transform_outbound(original.clone(), &[]).unwrap(); + let body = &encrypted[XOR_HEADER_LEN..]; + // XOR(0x0000, 0xA5A5) = 0xA5A5, so body != plaintext + assert_ne!(body, &original[..]); + assert_eq!(body, &[0xA5, 0xA5, 0xA5, 0xA5]); + } + + // ── XOR: error handling ───────────────────────────────────────────── + + #[test] + fn xor_inbound_too_short() { + let xform = XorPayloadTransform; + let result = xform.transform_inbound(vec![0x00, 0x01], &[]); + assert!(result.is_err()); + assert!(result.unwrap_err().contains("too short")); + } + + #[test] + fn xor_inbound_wrong_algo() { + let xform = XorPayloadTransform; + // algo_id = 0x0099, payload_len = 0, 2-byte body (even length) + let data = vec![0x00, 0x99, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00]; + let result = xform.transform_inbound(data, &[]); + assert!(result.is_err()); + assert!(result.unwrap_err().contains("unknown algo")); + } + + #[test] + fn xor_inbound_odd_body_length() { + let xform = XorPayloadTransform; + // Valid header + 3-byte body (odd) + let data = vec![0x00, 0x01, 0x00, 0x00, 0x00, 0x02, 0xAB, 0xCD, 0xEF]; + let result = xform.transform_inbound(data, &[]); + assert!(result.is_err()); + assert!(result.unwrap_err().contains("odd length")); + } + + // ── XOR: GenMessage integration ───────────────────────────────────── + + #[test] + fn xor_outbound_msg() { + let ctx = xor_context(); + let mut msg = GenMessage { + payload: vec![0x01, 0x02, 0x03, 0x04], + ..Default::default() + }; + ctx.outbound(&mut msg, false).unwrap(); + // Encrypted: header (6) + body (4) + assert_eq!(msg.payload.len(), XOR_HEADER_LEN + 4); + // Verify algo_id in header + let algo = u16::from_be_bytes([msg.payload[0], msg.payload[1]]); + assert_eq!(algo, XOR_ALGO_ID); + } + + #[test] + fn xor_inbound_msg() { + let ctx = xor_context(); + let original = vec![0xDE, 0xAD, 0xBE, 0xEF]; + + // Simulate outbound (encrypt) + let mut msg = GenMessage { + payload: original.clone(), + ..Default::default() + }; + ctx.outbound(&mut msg, false).unwrap(); + + // Simulate inbound (decrypt) + ctx.inbound(&mut msg, false).unwrap(); + assert_eq!(msg.payload, original); + } + + #[test] + fn xor_msg_roundtrip_updates_header_length() { + let ctx = xor_context(); + let mut msg = GenMessage::default(); + msg.header.stream_id = 42; + msg.header.type_ = 1; + msg.header.flags = 0x80; + msg.payload = vec![0x11, 0x22, 0x33]; + + // Original payload is 3 bytes, XOR adds 6-byte header + 1 byte padding = 10 bytes + ctx.outbound(&mut msg, false).unwrap(); + assert_eq!(msg.header.length, 10); // Updated to transformed length + assert_eq!(msg.header.stream_id, 42); // Other fields unchanged + + ctx.inbound(&mut msg, false).unwrap(); + // After roundtrip: payload restored, header.length reflects original + assert_eq!(msg.header.length, 3); // Back to original payload length + assert_eq!(msg.payload, vec![0x11, 0x22, 0x33]); + } + + // ── Failing transform ─────────────────────────────────────────────── + + #[derive(Debug)] + struct FailTransform; + + impl PayloadTransform for FailTransform { + fn transform_inbound(&self, _data: Vec, _aad: &[u8]) -> Result, String> { + Err("inbound failed".into()) + } + fn transform_outbound(&self, _data: Vec, _aad: &[u8]) -> Result, String> { + Err("outbound failed".into()) + } + } + + #[test] + fn failing_transform_propagates_error() { + let output = HookOutput { + data: ConnectionData::new(), + payload_transform: Some(Box::new(FailTransform)), + }; + let ctx = ConnectionContext::new(Some(output)); + assert!(ctx.transform_inbound(vec![1], &[]).is_err()); + assert!(ctx.transform_outbound(vec![1], &[]).is_err()); + } + + #[test] + fn failing_transform_pipeline_propagates_error() { + let output = HookOutput { + data: ConnectionData::new(), + payload_transform: Some(Box::new(FailTransform)), + }; + let ctx = ConnectionContext::new(Some(output)); + let mut msg = GenMessage { + payload: vec![1, 2, 3], + ..Default::default() + }; + assert!(ctx.inbound(&mut msg, false).is_err()); + assert!(ctx.outbound(&mut msg, false).is_err()); + } + + // ── ConnectionContext Arc sharing ──────────────────────────────────── + + #[test] + fn context_data_is_arc_shared() { + let mut data = ConnectionData::new(); + data.insert("key".into(), Box::new(99u64)); + let output = HookOutput { + data, + payload_transform: None, + }; + let ctx = ConnectionContext::new(Some(output)); + // Clone the Arc — should point to same data + let data2 = ctx.data.clone(); + assert_eq!(Arc::strong_count(&ctx.data), 2); + assert_eq!(data2.get_typed::("key"), Some(&99u64)); + } + + #[test] + fn context_transform_is_arc_shared() { + let output = HookOutput { + data: ConnectionData::new(), + payload_transform: Some(Box::new(NoopTransform)), + }; + let ctx = ConnectionContext::new(Some(output)); + let t1 = ctx.payload_transform.clone().unwrap(); + let t2 = ctx.payload_transform.clone().unwrap(); + // Both Arc clones point to the same transform + assert!(Arc::ptr_eq(&t1, &t2)); + } + + // ── on_accept unit tests (async, requires Socket) ────────────────── + + #[cfg(feature = "async")] + #[derive(Debug)] + struct MockAcceptHook { + should_reject: bool, + } + + #[cfg(feature = "async")] + impl AcceptHook for MockAcceptHook { + fn on_accept(&self, _fd: RawFd) -> std::result::Result { + if self.should_reject { + Err(HookError::Rejected("mock rejection".into())) + } else { + Ok(HookOutput { + data: ConnectionData::new(), + payload_transform: None, + }) + } + } + } + + /// Helper: build a ServerExtensionConfig with an optional accept_hook. + #[cfg(feature = "async")] + fn cfg_with_hook(hook: Option>) -> ServerExtensionConfig { + ServerExtensionConfig { accept_hook: hook } + } + + #[cfg(feature = "async")] + #[tokio::test] + async fn on_accept_no_hook_returns_ok_none() { + let cfg = ServerExtensionConfig::default(); + // Create a real socket via UnixStream pair + let (client_stream, _server_stream) = tokio::net::UnixStream::pair().unwrap(); + let socket = Socket::from(client_stream); + let result = cfg.on_accept(&socket).await; + assert!(result.is_ok(), "on_accept should not fail"); + assert!(result.unwrap().is_none(), "No hook → Ok(None)"); + } + + #[cfg(feature = "async")] + #[tokio::test] + async fn on_accept_socket_without_raw_fd_returns_err() { + let hook: Arc = Arc::new(MockAcceptHook { + should_reject: false, + }); + let cfg = cfg_with_hook(Some(hook)); + // Socket::new() does NOT capture raw_fd + let (client_stream, _) = tokio::net::UnixStream::pair().unwrap(); + let socket = Socket::new(client_stream); + assert_eq!(socket.as_raw_fd(), None); + let err = match cfg.on_accept(&socket).await { + Ok(_) => panic!("hook configured but no fd → should fail"), + Err(e) => e, + }; + let err_str = format!("{}", err); + assert!( + err_str.contains("accept hook requires a raw fd"), + "error should tell caller to use Server::bind() or Socket::from: {}", + err_str + ); + } + + #[cfg(feature = "async")] + #[tokio::test] + async fn on_accept_hook_succeeds_returns_ok_some() { + let hook: Arc = Arc::new(MockAcceptHook { + should_reject: false, + }); + let cfg = cfg_with_hook(Some(hook)); + // Socket::from() DOES capture raw_fd + let (client_stream, _) = tokio::net::UnixStream::pair().unwrap(); + let socket = Socket::from(client_stream); + assert!(socket.as_raw_fd().is_some()); + let result = cfg.on_accept(&socket).await; + assert!(result.is_ok()); + assert!( + result.unwrap().is_some(), + "Hook succeeded → Ok(Some(output))" + ); + } + + #[cfg(feature = "async")] + #[tokio::test] + async fn on_accept_hook_rejected_returns_err() { + let hook: Arc = Arc::new(MockAcceptHook { + should_reject: true, + }); + let cfg = cfg_with_hook(Some(hook)); + let (client_stream, _) = tokio::net::UnixStream::pair().unwrap(); + let socket = Socket::from(client_stream); + let result = cfg.on_accept(&socket).await; + assert!(result.is_err(), "Hook rejected → Err(HookError)"); + } + + // ── XOR: additional edge cases ────────────────────────────────────── + + #[test] + fn xor_roundtrip_all_zeros() { + let ctx = xor_context(); + let original = vec![0x00; 16]; + let encrypted = ctx.transform_outbound(original.clone(), &[]).unwrap(); + // All zeros XOR 0xA5A5 → all 0xA5A5 + let body = &encrypted[XOR_HEADER_LEN..]; + assert!(body.iter().all(|&b| b == 0xA5)); + let decrypted = ctx.transform_inbound(encrypted, &[]).unwrap(); + assert_eq!(decrypted, original); + } + + #[test] + fn xor_roundtrip_all_ones() { + let ctx = xor_context(); + let original = vec![0xFF; 20]; + let encrypted = ctx.transform_outbound(original.clone(), &[]).unwrap(); + let decrypted = ctx.transform_inbound(encrypted, &[]).unwrap(); + assert_eq!(decrypted, original); + } + + #[test] + fn xor_inbound_empty_header_exactly() { + let xform = XorPayloadTransform; + // Exactly XOR_HEADER_LEN bytes = valid header with 0-length body + let data = vec![0x00, 0x01, 0x00, 0x00, 0x00, 0x00]; + let result = xform.transform_inbound(data, &[]); + assert!(result.is_ok()); + assert!(result.unwrap().is_empty()); + } + + #[test] + fn xor_transform_inbound_idempotent_after_outbound() { + // Encrypt, then decrypt twice — second decrypt should fail (header mismatch) + let xform = XorPayloadTransform; + let original = vec![0x01, 0x02, 0x03, 0x04]; + let encrypted = xform.transform_outbound(original.clone(), &[]).unwrap(); + let decrypted = xform.transform_inbound(encrypted.clone(), &[]).unwrap(); + assert_eq!(decrypted, original); + // Decrypting already-decrypted data should fail (algo_id mismatch) + let result = xform.transform_inbound(decrypted, &[]); + assert!(result.is_err(), "Double-decrypt should fail"); + } +} diff --git a/src/sync/client.rs b/src/sync/client.rs index 50d9bd5d..cde48628 100644 --- a/src/sync/client.rs +++ b/src/sync/client.rs @@ -14,9 +14,6 @@ //! Sync client of ttrpc. -#[cfg(unix)] -use std::os::unix::io::RawFd; - use protobuf::Message; use std::collections::HashMap; use std::sync::mpsc; @@ -28,9 +25,16 @@ use crate::error::{Error, Result}; use crate::proto::{ check_oversize, Code, Codec, MessageHeader, Request, Response, MESSAGE_TYPE_RESPONSE, }; +use crate::security_extension::serialize_aad; use crate::sync::channel::{read_message, write_message}; use crate::sync::sys::ClientConnection; +#[cfg(feature = "security_extension")] +use std::os::unix::io::RawFd; +use crate::ConnectionContext; +#[cfg(feature = "security_extension")] +use crate::security_extension::{ConnectHook, HookOutput}; + #[cfg(windows)] use super::sys::PipeConnection; @@ -43,25 +47,72 @@ type ReciverMap = Arc>>>>>; pub struct Client { _connection: Arc, sender_tx: Sender, + _conn_ctx: Arc, } impl Client { pub fn connect(sockaddr: &str) -> Result { let conn = ClientConnection::client_connect(sockaddr)?; - Self::new_client(conn) + Self::new_client(conn, None) + } + + /// Create a sync client with a [`ConnectHook`] for security negotiation. + /// + /// The hook is invoked with the connection's raw file descriptor before + /// any ttrpc messages are exchanged. It can perform handshakes and return + /// a [`PayloadTransform`](crate::security_extension::PayloadTransform) for + /// connection-level encryption. + /// + /// The fd is captured by a `ClientConnection` + /// **before** the hook runs, so if the hook rejects, the connection's + /// `Drop` impl closes the fd and prevents leaks. + #[cfg(feature = "security_extension")] + pub fn with_hook(fd: RawFd, hook: H) -> Result { + // Take ownership of the fd BEFORE invoking the hook. ClientConnection + // has a Drop impl that closes `fd` (and its internal socket_pair), so + // if the hook rejects we just propagate the error and the Drop cleanup + // releases the fd — no leak, no double-close. + let conn = ClientConnection::new(fd) + .map_err(err_to_others_err!(e, "new ClientConnection"))?; + let output = hook.on_connect(fd).map_err(|e| { + Error::Others(format!( + "sync client connect hook failed (fd={}): {}", + fd, e + )) + })?; + Self::new_client(conn, Some(output)) + } + + /// Returns the per-connection metadata from the [`ConnectHook`]. + /// + /// This is the [`ConnectionData`](crate::security_extension::ConnectionData) + /// returned by the hook during connection establishment. Empty (default) + /// when no hook was configured. + #[cfg(feature = "security_extension")] + pub fn connection_data(&self) -> &crate::security_extension::ConnectionData { + &self._conn_ctx.data } #[cfg(unix)] /// Initialize a new [`Client`] from raw file descriptor. - pub fn new(fd: RawFd) -> Result { + pub fn new(fd: std::os::unix::io::RawFd) -> Result { let conn = ClientConnection::new(fd).map_err(err_to_others_err!(e, "new ClientConnection"))?; - Self::new_client(conn) + Self::new_client(conn, None) } - fn new_client(pipe_client: ClientConnection) -> Result { + fn new_client( + pipe_client: ClientConnection, + #[cfg(feature = "security_extension")] + hook_output: Option, + #[cfg(not(feature = "security_extension"))] _hook_output: Option<()>, + ) -> Result { + #[cfg(feature = "security_extension")] + let conn_ctx = Arc::new(ConnectionContext::new(hook_output)); + #[cfg(not(feature = "security_extension"))] + let conn_ctx = Arc::new(ConnectionContext::default()); let client = Arc::new(pipe_client); let weak_client = Arc::downgrade(&client); let (sender_tx, rx): (Sender, Receiver) = mpsc::channel(); @@ -72,6 +123,7 @@ impl Client { let sender_client = connection.clone(); //Sender + let sender_ctx = conn_ctx.clone(); thread::spawn(move || { let mut stream_id: u32 = 1; for (buf, recver_tx) in rx.iter() { @@ -85,6 +137,28 @@ impl Client { let mut mh = MessageHeader::new_request(0, buf.len() as u32); mh.set_stream_id(current_stream_id); + // ── outbound transform ── + let (mh, buf) = match sender_ctx.outbound_buf( + buf, + &serialize_aad(&mh), + false, + ) { + Ok(transformed) => { + mh.length = transformed.len() as u32; + (mh, transformed) + } + Err(e) => { + { + let mut map = receiver_map.lock().unwrap(); + map.remove(¤t_stream_id); + } + recver_tx + .send(Err(e)) + .unwrap_or_else(|_e| error!("The request has returned")); + continue; + } + }; + if let Err(e) = write_message(&sender_client, mh, buf) { //Remove current_stream_id and recver_tx to recver_map { @@ -105,6 +179,7 @@ impl Client { //ClientConnection's drop will be not call until the thread finished. It means if all the external references are finished, //this thread should be release. let receiver_client = weak_client.clone(); + let receiver_ctx = conn_ctx.clone(); thread::spawn(move || { loop { //The count of ClientConnection's Arc will be add one , and back to original value when this code ends. @@ -124,7 +199,15 @@ impl Client { } match read_message(&receiver_connection) { - Ok((mh, buf)) => { + Ok((mh, y)) => { + let buf = match y { + Ok(data) => receiver_ctx.inbound_buf( + data, + &serialize_aad(&mh), + false, + ), + Err(e) => Err(e), + }; trans_resp(recver_map_orig.clone(), mh, buf); } Err(x) => match x { @@ -155,6 +238,7 @@ impl Client { Ok(Client { _connection: client, sender_tx, + _conn_ctx: conn_ctx, }) } pub fn request(&self, req: Request) -> Result { diff --git a/src/sync/mod.rs b/src/sync/mod.rs index 53d0680c..66a0464f 100644 --- a/src/sync/mod.rs +++ b/src/sync/mod.rs @@ -17,5 +17,6 @@ pub use client::Client; pub use server::Server; #[doc(hidden)] +#[allow(deprecated)] pub use utils::response_to_channel; -pub use utils::{MethodHandler, TtrpcContext}; +pub use utils::{send_response, MethodHandler, TtrpcContext}; diff --git a/src/sync/server.rs b/src/sync/server.rs index 555224de..c8908eb9 100644 --- a/src/sync/server.rs +++ b/src/sync/server.rs @@ -27,13 +27,35 @@ use std::sync::{Arc, Mutex}; use std::thread; use std::thread::JoinHandle; -use super::utils::{response_error_to_channel, response_to_channel}; +use super::utils::{response_error_to_channel, send_response}; use crate::context; use crate::error::{get_status, Error, Result}; use crate::proto::{Code, MessageHeader, Request, Response, MESSAGE_TYPE_REQUEST}; use crate::sync::channel::{read_message, write_message}; use crate::sync::sys::{PipeConnection, PipeListener}; use crate::{MethodHandler, TtrpcContext}; +use crate::ConnectionContext; +use crate::security_extension::serialize_aad; + +#[cfg(feature = "security_extension")] +use crate::security_extension::{AcceptHook, HookError}; + +/// Invoke an [`AcceptHook`] for a newly accepted connection and wrap the +/// output in a [`ConnectionContext`]. +/// +/// Returns `Ok(ConnectionContext::default())` when `hook` is `None`. On hook +/// failure the error is forwarded to the caller so it can log, close the +/// rejected connection, and continue the accept loop. +#[cfg(feature = "security_extension")] +fn connection_context_from_hook( + fd: i32, + hook: Option<&Arc>, +) -> std::result::Result { + match hook { + Some(h) => h.on_accept(fd).map(|o| ConnectionContext::new(Some(o))), + None => Ok(ConnectionContext::default()), + } +} // poll_queue will create WAIT_THREAD_COUNT_DEFAULT threads in begin. // If wait thread count < WAIT_THREAD_COUNT_MIN, create number to WAIT_THREAD_COUNT_DEFAULT. @@ -60,6 +82,8 @@ pub struct Server { thread_count_min: usize, thread_count_max: usize, accept_retry_interval: Duration, + #[cfg(feature = "security_extension")] + accept_hook: Option>, } struct Connection { @@ -93,6 +117,7 @@ struct ThreadS<'a> { default: usize, min: usize, max: usize, + conn_ctx: &'a Arc, } #[allow(clippy::too_many_arguments)] @@ -107,6 +132,7 @@ fn start_method_handler_thread( cancel_rx: crossbeam::channel::Receiver<()>, min: usize, max: usize, + conn_ctx: Arc, ) { thread::spawn(move || { while !quit.load(Ordering::SeqCst) { @@ -142,7 +168,12 @@ fn start_method_handler_thread( buf = y; } Ok((mh, Err(e))) => { - if let Err(x) = response_error_to_channel(mh.stream_id, e, res_tx.clone()) { + if let Err(x) = response_error_to_channel( + mh.stream_id, + e, + res_tx.clone(), + conn_ctx.as_ref(), + ) { debug!("response_error_to_channel get error {:?}", x); quit_connection(quit, control_tx); break; @@ -168,7 +199,12 @@ fn start_method_handler_thread( let status = get_status(Code::INVALID_ARGUMENT, x.to_string()); let mut res = Response::new(); res.set_status(status); - if let Err(x) = response_to_channel(mh.stream_id, res, res_tx.clone()) { + if let Err(x) = send_response( + mh.stream_id, + res, + res_tx.clone(), + conn_ctx.as_ref(), + ) { debug!("response_to_channel get error {:?}", x); quit_connection(quit, control_tx); break; @@ -184,7 +220,12 @@ fn start_method_handler_thread( let status = get_status(Code::INVALID_ARGUMENT, format!("{path} does not exist")); let mut res = Response::new(); res.set_status(status); - if let Err(x) = response_to_channel(mh.stream_id, res, res_tx.clone()) { + if let Err(x) = send_response( + mh.stream_id, + res, + res_tx.clone(), + conn_ctx.as_ref(), + ) { info!("response_to_channel get error {:?}", x); quit_connection(quit, control_tx); break; @@ -198,6 +239,7 @@ fn start_method_handler_thread( res_tx: res_tx.clone(), metadata: context::from_pb(&req.metadata), timeout_nano: req.timeout_nano, + conn_ctx: conn_ctx.clone(), }; if let Err(x) = method.handler(ctx, req) { debug!("method handle {} get error {:?}", path, x); @@ -224,6 +266,7 @@ fn start_method_handler_threads(num: usize, ts: &ThreadS) { ts.cancel_rx.clone(), ts.min, ts.max, + ts.conn_ctx.clone(), ); } } @@ -248,6 +291,8 @@ impl Default for Server { thread_count_min: DEFAULT_WAIT_THREAD_COUNT_MIN, thread_count_max: DEFAULT_WAIT_THREAD_COUNT_MAX, accept_retry_interval: DEFAULT_ACCEPT_RETRY_INTERVAL, + #[cfg(feature = "security_extension")] + accept_hook: None, } } } @@ -314,6 +359,23 @@ impl Server { self } + /// Set a hook invoked on each accepted connection (Unix only). + /// + /// The hook receives the connection's raw file descriptor and can perform + /// authentication, peer identity inspection, or handshake negotiation. + /// It returns [`HookOutput`](crate::security_extension::HookOutput) with optional + /// per-connection data and a + /// [`PayloadTransform`](crate::security_extension::PayloadTransform) for encryption. + /// + /// If the hook returns an error, the connection is closed and the accept + /// loop continues to the next connection. + #[cfg(feature = "security_extension")] + pub fn set_accept_hook(mut self, hook: H) -> Self { + let hook: Arc = Arc::new(hook); + self.accept_hook = Some(hook); + self + } + pub fn start_listen(&mut self) -> Result<()> { let connections = self.connections.clone(); @@ -363,6 +425,9 @@ impl Server { } }; + #[cfg(feature = "security_extension")] + let accept_hook = self.accept_hook.clone(); + let handler = thread::Builder::new() .name("listener_loop".into()) .spawn(move || { @@ -395,11 +460,38 @@ impl Server { } }; + // ── Accept hook invocation ── + #[cfg(feature = "security_extension")] + let conn_ctx = match connection_context_from_hook( + pipe_connection.id(), + accept_hook.as_ref(), + ) { + Ok(ctx) => Arc::new(ctx), + Err(e) => { + let fd = pipe_connection.id(); + log::warn!("accept hook failed: {:?}", e); + // Close the rejected connection; PipeConnection has no + // Drop impl so the fd must be closed explicitly. + if let Err(ce) = pipe_connection.close() { + log::warn!( + "failed to close rejected connection (fd={}): {:?}", + fd, + ce + ); + } + continue; + } + }; + #[cfg(not(feature = "security_extension"))] + let conn_ctx = Arc::new(ConnectionContext::default()); + let methods = methods.clone(); let quit = Arc::new(AtomicBool::new(false)); let child_quit = quit.clone(); let reaper_tx_child = reaper_tx.clone(); let pipe_connection_child = pipe_connection.clone(); + let reader_ctx = conn_ctx.clone(); + let handler_ctx = conn_ctx.clone(); let (sync_tx, sync_rx) = channel(); @@ -438,7 +530,22 @@ impl Server { while !quit_reader.load(Ordering::SeqCst) { let msg = read_message(&pipe_reader); match msg { - Ok((x, y)) => { + Ok((mut x, y)) => { + // ── inbound transform ── + let y = match y { + Ok(data) => { + let res = reader_ctx.inbound_buf( + data, + &serialize_aad(&x), + true, + ); + if let Ok(ref buf) = res { + x.length = buf.len() as u32; + } + res + } + Err(e) => Err(e), + }; let res = workload_tx.send((x, y)); match res { Ok(_) => {} @@ -490,6 +597,7 @@ impl Server { default, min, max, + conn_ctx: &handler_ctx, }; start_method_handler_threads(ts.default, &ts); diff --git a/src/sync/utils.rs b/src/sync/utils.rs index 616b615e..f704893b 100644 --- a/src/sync/utils.rs +++ b/src/sync/utils.rs @@ -3,44 +3,103 @@ // SPDX-License-Identifier: Apache-2.0 // -use crate::error::{Error, Result}; +use crate::error::{get_status, Error, Result}; use crate::proto::{ - check_oversize, Codec, MessageHeader, Request, Response, MESSAGE_TYPE_RESPONSE, + Codec, Code, MessageHeader, Request, Response, MESSAGE_TYPE_RESPONSE, }; +use crate::ConnectionContext; +use crate::security_extension::serialize_aad; use std::collections::HashMap; +use std::sync::Arc; -/// Response message through a channel. -/// Eventually the message will sent to Client. -pub fn response_to_channel( +/// Send a [`Response`] through the channel with full outbound pipeline. +/// +/// Encodes the response, applies the connection transform (if any), and +/// sends the result through `tx`. Handles oversize responses safely: +/// +/// 1. **Pre-check**: verifies the raw payload fits within the transform-safe +/// limit (`ConnectionContext::max_raw_payload_len`). If it doesn't, builds +/// a small error response instead — avoiding state advancement of stateful +/// transforms (e.g., AEAD nonce counters) on data that would be rejected. +/// 2. **Transform**: applies `transform_outbound` and the post-transform +/// size guard exactly once. +/// +/// This single-pass design prevents the double-transform bug where a +/// stateful cipher would produce undecryptable output on a fallback message. +pub fn send_response( stream_id: u32, res: Response, tx: std::sync::mpsc::Sender<(MessageHeader, Vec)>, + ctx: &ConnectionContext, ) -> Result<()> { let mut buf = res.encode().map_err(err_to_others_err!(e, ""))?; - if let Err(e) = check_oversize(buf.len(), true) { - let resp: Response = e.into(); - buf = resp.encode().map_err(err_to_others_err!(e, ""))?; - }; + // Pre-check: ensure raw payload fits within the transform-safe limit. + let max_len = ctx.max_raw_payload_len(); + if buf.len() > max_len { + let err_msg = format!( + "response payload {} bytes exceeds safe limit {} bytes (after transform overhead)", + buf.len(), + max_len + ); + let err_resp: Response = + Error::RpcStatus(get_status(Code::INVALID_ARGUMENT, err_msg)).into(); + buf = err_resp.encode().map_err(err_to_others_err!(e, ""))?; + } + // Build the header early so we can serialize AAD for the transform. + // `length` is excluded from AAD, so the placeholder value doesn't matter. + let aad = serialize_aad(&MessageHeader::new_response(stream_id, 0)); let mh = MessageHeader { - length: buf.len() as u32, + length: 0, // Will be set by send_response_sync after transform. stream_id, type_: MESSAGE_TYPE_RESPONSE, flags: 0, }; - tx.send((mh, buf)).map_err(err_to_others_err!(e, ""))?; + // ── Serialize transform + enqueue ── + ctx.send_response_sync(buf, &aad, &tx, mh)?; Ok(()) } +/// Response message through a channel. Eventually the message will be sent +/// to Client. +/// +/// # Deprecated +/// +/// This helper does **not** apply the connection's payload transform. +/// On an encrypted connection, responses sent through this function +/// **leak as plaintext** and an encrypted client will fail to decode them. +/// +/// Use [`TtrpcContext::respond`] (from generated `request_handler!` macros or +/// hand-written handlers) or [`send_response`] with the connection's +/// [`ConnectionContext`] instead. +#[deprecated( + since = "0.9.0", + note = "Bypasses payload transform. Use TtrpcContext::respond() or send_response() instead." +)] +pub fn response_to_channel( + stream_id: u32, + res: Response, + tx: std::sync::mpsc::Sender<(MessageHeader, Vec)>, +) -> Result<()> { + let ctx = ConnectionContext::default(); + send_response(stream_id, res, tx, &ctx) +} + +/// Send an error as a transform-aware response through the channel. +/// +/// If the connection has a [`PayloadTransform`](crate::security_extension::PayloadTransform) +/// configured, the error response is encrypted before being sent so the +/// client can decode it on an encrypted connection. pub fn response_error_to_channel( stream_id: u32, e: Error, tx: std::sync::mpsc::Sender<(MessageHeader, Vec)>, + ctx: &ConnectionContext, ) -> Result<()> { - response_to_channel(stream_id, e.into(), tx) + send_response(stream_id, e.into(), tx, ctx) } /// Handle request in sync mode. @@ -74,7 +133,7 @@ macro_rules! request_handler { } }, } - ::ttrpc::response_to_channel($ctx.mh.stream_id, res, $ctx.res_tx)? + $ctx.respond($ctx.mh.stream_id, res)? }; } @@ -116,6 +175,21 @@ pub struct TtrpcContext { pub res_tx: std::sync::mpsc::Sender<(MessageHeader, Vec)>, pub metadata: HashMap>, pub timeout_nano: i64, + /// 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, +} + +impl TtrpcContext { + /// Encode, optionally transform, and send a response through this context's channel. + /// + /// This is the preferred way to send responses from within a [`MethodHandler`] + /// implementation. The response is encoded, the connection's + /// [`PayloadTransform`](crate::security_extension::PayloadTransform) is applied (if any), + /// an oversize check is performed, and the result is sent to the client. + pub fn respond(&self, stream_id: u32, res: Response) -> Result<()> { + send_response(stream_id, res, self.res_tx.clone(), self.conn_ctx.as_ref()) + } } /// Trait that implements handler which is a proxy to the desired method (sync). diff --git a/tests/common/mod.rs b/tests/common/mod.rs new file mode 100644 index 00000000..dd95697f --- /dev/null +++ b/tests/common/mod.rs @@ -0,0 +1,237 @@ +// Copyright 2026 Alibaba Cloud. All rights reserved. +// +// SPDX-License-Identifier: Apache-2.0 +// +//! Shared test utilities for hook integration tests. + +use std::sync::atomic::{AtomicU64, Ordering}; + +use ttrpc::security_extension::PayloadTransform; + +// ── XOR Transform constants ───────────────────────────────────────────────── + +pub const XOR_KEY: u16 = 0xA5A5; +pub const XOR_ALGO_ID: u16 = 0x0001; +/// algo_id(2) + aad_tag(2) + payload_len(4) +pub const XOR_HEADER_LEN: usize = 8; + +// ── XOR PayloadTransform ──────────────────────────────────────────────────── + +/// Symmetric XOR-based payload transform for testing. +/// +/// Wire format: `[algo_id:u16][aad_tag:u16][payload_len:u32][encrypted_data...]` +/// Encryption: XOR each u16 word with an AAD-derived effective key. +/// +/// The effective key is `XOR_KEY ^ aad_fold(aad)`, so any tampering with +/// the header fields (`stream_id`, `type_`, `flags`) is detected via the +/// `aad_tag` check on the inbound path. +#[derive(Debug)] +pub struct XorPayloadTransform; + +/// Fold the 6-byte AAD (`stream_id || type_ || flags`) into a u16 by +/// XOR-ing adjacent byte pairs. +pub fn aad_to_u16(aad: &[u8]) -> u16 { + let mut k: u16 = 0; + for (i, &b) in aad.iter().enumerate() { + if i % 2 == 0 { + k ^= (b as u16) << 8; + } else { + k ^= b as u16; + } + } + k +} + +fn effective_key(aad: &[u8]) -> u16 { + XOR_KEY ^ aad_to_u16(aad) +} + +impl PayloadTransform for XorPayloadTransform { + fn transform_outbound(&self, data: Vec, aad: &[u8]) -> Result, String> { + let key = effective_key(aad); + let payload_len = data.len(); + let mut padded = data; + if padded.len() % 2 != 0 { + padded.push(0x00); + } + let mut encrypted = Vec::with_capacity(XOR_HEADER_LEN + padded.len()); + encrypted.extend_from_slice(&XOR_ALGO_ID.to_be_bytes()); + encrypted.extend_from_slice(&aad_to_u16(aad).to_be_bytes()); // aad_tag + encrypted.extend_from_slice(&(payload_len as u32).to_be_bytes()); + for chunk in padded.chunks(2) { + let val = u16::from_be_bytes([chunk[0], chunk[1]]); + let xored = val ^ key; + encrypted.extend_from_slice(&xored.to_be_bytes()); + } + Ok(encrypted) + } + + fn transform_inbound(&self, data: Vec, aad: &[u8]) -> Result, String> { + if data.len() < XOR_HEADER_LEN { + return Err(format!( + "xor: packet too short ({} < {})", + data.len(), + XOR_HEADER_LEN + )); + } + let algo_id = u16::from_be_bytes([data[0], data[1]]); + if algo_id != XOR_ALGO_ID { + return Err(format!("xor: unknown algo 0x{:04X}", algo_id)); + } + // Verify AAD integrity tag + let stored_tag = u16::from_be_bytes([data[2], data[3]]); + let expected_tag = aad_to_u16(aad); + if stored_tag != expected_tag { + return Err(format!( + "xor: AAD mismatch (stored=0x{:04X}, expected=0x{:04X}) — header tampered", + stored_tag, expected_tag + )); + } + let payload_len = u32::from_be_bytes([data[4], data[5], data[6], data[7]]) as usize; + let encrypted = &data[XOR_HEADER_LEN..]; + if encrypted.len() % 2 != 0 { + return Err("xor: encrypted data has odd length".into()); + } + let key = effective_key(aad); + let mut decrypted = Vec::with_capacity(encrypted.len()); + for chunk in encrypted.chunks(2) { + let val = u16::from_be_bytes([chunk[0], chunk[1]]); + let xored = val ^ key; + decrypted.extend_from_slice(&xored.to_be_bytes()); + } + decrypted.truncate(payload_len); + Ok(decrypted) + } +} + +// ── Socket path helper ────────────────────────────────────────────────────── + +static SOCKET_COUNTER: AtomicU64 = AtomicU64::new(0); + +/// Generate a unique temporary Unix socket path and clean up any stale file. +pub fn temp_unix_socket_path() -> String { + let id = SOCKET_COUNTER.fetch_add(1, Ordering::SeqCst); + let path = format!("/tmp/ttrpc_test_{}_{}.sock", std::process::id(), id); + let _ = std::fs::remove_file(&path); + path +} + +/// Remove a socket file (ignores errors). +pub fn cleanup_socket_file(path: &str) { + let _ = std::fs::remove_file(path); +} + +// ── AAD integrity tests ────────────────────────────────────────────────────── + +/// Build a standard 6-byte AAD from `stream_id`, `type_`, `flags`. +pub fn make_aad(stream_id: u32, type_: u8, flags: u8) -> [u8; 6] { + let mut aad = [0u8; 6]; + aad[0..4].copy_from_slice(&stream_id.to_be_bytes()); + aad[4] = type_; + aad[5] = flags; + aad +} + +#[cfg(test)] +mod aad_tests { + use super::*; + + #[test] + fn roundtrip_with_matching_aad() { + let xform = XorPayloadTransform; + let aad = make_aad(1, 0x01, 0x00); // stream_id=1, REQUEST, no flags + let plaintext = b"hello world".to_vec(); + + let encrypted = xform.transform_outbound(plaintext.clone(), &aad).unwrap(); + let decrypted = xform.transform_inbound(encrypted, &aad).unwrap(); + assert_eq!(decrypted, plaintext); + } + + #[test] + fn detect_stream_id_tampering() { + let xform = XorPayloadTransform; + let aad_orig = make_aad(1, 0x01, 0x00); + let plaintext = b"secret data".to_vec(); + + let encrypted = xform.transform_outbound(plaintext, &aad_orig).unwrap(); + + // Attacker changes stream_id in the header + let aad_tampered = make_aad(99, 0x01, 0x00); + let err = xform + .transform_inbound(encrypted, &aad_tampered) + .unwrap_err(); + assert!( + err.contains("AAD mismatch"), + "expected AAD error, got: {}", + err + ); + } + + #[test] + fn detect_message_type_tampering() { + let xform = XorPayloadTransform; + let aad_orig = make_aad(3, 0x01, 0x00); // REQUEST + let plaintext = b"request payload".to_vec(); + + let encrypted = xform.transform_outbound(plaintext, &aad_orig).unwrap(); + + // Attacker changes type_ from REQUEST (0x01) to RESPONSE (0x02) + let aad_tampered = make_aad(3, 0x02, 0x00); + let err = xform + .transform_inbound(encrypted, &aad_tampered) + .unwrap_err(); + assert!( + err.contains("AAD mismatch"), + "expected AAD error, got: {}", + err + ); + } + + #[test] + fn detect_flags_tampering() { + let xform = XorPayloadTransform; + let aad_orig = make_aad(5, 0x03, 0x00); // DATA, no flags + let plaintext = b"streaming chunk".to_vec(); + + let encrypted = xform.transform_outbound(plaintext, &aad_orig).unwrap(); + + // Attacker sets flags to 0xFF + let aad_tampered = make_aad(5, 0x03, 0xFF); + let err = xform + .transform_inbound(encrypted, &aad_tampered) + .unwrap_err(); + assert!( + err.contains("AAD mismatch"), + "expected AAD error, got: {}", + err + ); + } + + #[test] + fn detect_close_frame_aad() { + // Close messages are DATA frames with FLAG_REMOTE_CLOSED | FLAG_NO_DATA. + // AAD = stream_id || type_(DATA=0x03) || flags(0x05). + let xform = XorPayloadTransform; + let aad_close = make_aad(7, 0x03, 0x05); // DATA type, close flags + let empty_payload = vec![]; // close has empty payload + + let encrypted = xform + .transform_outbound(empty_payload.clone(), &aad_close) + .unwrap(); + + // Legitimate close decrypts fine + let decrypted = xform + .transform_inbound(encrypted.clone(), &aad_close) + .unwrap(); + assert_eq!(decrypted, empty_payload); + + // But if stream_id is changed, the close is rejected + let aad_wrong = make_aad(99, 0x03, 0x05); + let err = xform.transform_inbound(encrypted, &aad_wrong).unwrap_err(); + assert!( + err.contains("AAD mismatch"), + "expected AAD error, got: {}", + err + ); + } +} diff --git a/tests/hook_integration_async_unix.rs b/tests/hook_integration_async_unix.rs new file mode 100644 index 00000000..f6c7eda1 --- /dev/null +++ b/tests/hook_integration_async_unix.rs @@ -0,0 +1,1239 @@ +// Copyright 2026 Alibaba Cloud. All rights reserved. +// +// SPDX-License-Identifier: Apache-2.0 +// +#![cfg(all(feature = "async", feature = "security_extension"))] +//! Integration tests for AcceptHook, ConnectHook, and PayloadTransform +//! exercising real client-server connections over Unix sockets. + +mod common; + +use std::collections::HashMap; +use std::os::unix::io::RawFd; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::Arc; +use std::time::Duration; + +use async_trait::async_trait; +use tokio::time::sleep; + +use ttrpc::asynchronous::{Client, MethodHandler, Server, Service, TtrpcContext}; +use ttrpc::proto::{Request, Response, Status}; +use ttrpc::security_extension::{ + AcceptHook, ConnectHook, ConnectionData, ConnectionDataExt, HookError, HookOutput, + PayloadTransform, +}; + +use common::{cleanup_socket_file, temp_unix_socket_path, XorPayloadTransform}; + +// ── Test constants ────────────────────────────────────────────────────────── + +const TEST_SERVICE: &str = "test.TestService"; +const TEST_METHOD: &str = "Echo"; + +// ── Test AcceptHook (server side) ────────────────────────────────────────── + +#[derive(Debug)] +struct TestAcceptHook { + call_count: Arc, + attach_data: bool, + attach_transform: bool, + reject: bool, +} + +impl TestAcceptHook { + fn new_rejecting() -> Self { + Self { + call_count: Arc::new(AtomicUsize::new(0)), + attach_data: false, + attach_transform: false, + reject: true, + } + } + + fn new_accepting(with_data: bool, with_transform: bool) -> Self { + Self { + call_count: Arc::new(AtomicUsize::new(0)), + attach_data: with_data, + attach_transform: with_transform, + reject: false, + } + } +} + +impl AcceptHook for TestAcceptHook { + fn on_accept(&self, _fd: RawFd) -> Result { + self.call_count.fetch_add(1, Ordering::SeqCst); + + if self.reject { + return Err(HookError::Rejected("test rejection".into())); + } + + let mut data = ConnectionData::new(); + if self.attach_data { + data.insert("peer_role".into(), Box::new(String::from("test_client"))); + data.insert("peer_cid".into(), Box::new(42u32)); + } + + let payload_transform = if self.attach_transform { + Some(Box::new(XorPayloadTransform) as Box) + } else { + None + }; + + Ok(HookOutput { + data, + payload_transform, + }) + } +} + +// ── Test ConnectHook (client side) ───────────────────────────────────────── + +#[derive(Debug)] +struct TestConnectHook { + call_count: Arc, + attach_data: bool, + attach_transform: bool, +} + +impl TestConnectHook { + fn new(with_data: bool, with_transform: bool) -> Self { + Self { + call_count: Arc::new(AtomicUsize::new(0)), + attach_data: with_data, + attach_transform: with_transform, + } + } +} + +impl ConnectHook for TestConnectHook { + fn on_connect(&self, _fd: RawFd) -> Result { + self.call_count.fetch_add(1, Ordering::SeqCst); + + let mut data = ConnectionData::new(); + if self.attach_data { + data.insert("client_role".into(), Box::new(String::from("test_client"))); + data.insert("client_id".into(), Box::new(99u64)); + } + + let payload_transform = if self.attach_transform { + Some(Box::new(XorPayloadTransform) as Box) + } else { + None + }; + + Ok(HookOutput { + data, + payload_transform, + }) + } +} + +// ── Echo MethodHandler ───────────────────────────────────────────────────── + +#[derive(Debug)] +struct EchoHandler; + +#[async_trait] +impl MethodHandler for EchoHandler { + async fn handler(&self, ctx: TtrpcContext, req: Request) -> ttrpc::Result { + // Echo the request payload back as the response payload + let mut resp = Response::new(); + resp.set_status(Status::default()); + // Copy request payload to response payload + resp.payload = req.payload; + + // Optionally attach connection data to response for verification + if let Some(role) = ctx.connection_data.get_typed::("peer_role") { + // Prepend role info to response for test verification + let prefix = format!("role:{}|", role); + let mut new_payload = prefix.into_bytes(); + new_payload.extend_from_slice(&resp.payload); + resp.payload = new_payload; + } + + Ok(resp) + } +} + +// ── SlowEcho MethodHandler (for timeout testing) ─────────────────────────── + +#[derive(Debug)] +struct SlowEchoHandler; + +#[async_trait] +impl MethodHandler for SlowEchoHandler { + async fn handler(&self, _ctx: TtrpcContext, req: Request) -> ttrpc::Result { + // Sleep just long enough to guarantee client timeout + sleep(Duration::from_secs(1)).await; + let mut resp = Response::new(); + resp.set_status(Status::default()); + resp.payload = req.payload; + Ok(resp) + } +} + +fn build_slow_echo_service() -> HashMap { + let mut methods: HashMap> = HashMap::new(); + methods.insert(TEST_METHOD.to_string(), Box::new(SlowEchoHandler)); + + let mut services = HashMap::new(); + services.insert( + TEST_SERVICE.to_string(), + Service { + methods, + streams: HashMap::new(), + }, + ); + services +} + +fn build_test_service() -> HashMap { + let mut methods: HashMap> = HashMap::new(); + // Register with just the method name, not the full path + methods.insert(TEST_METHOD.to_string(), Box::new(EchoHandler)); + + let mut services = HashMap::new(); + services.insert( + TEST_SERVICE.to_string(), + Service { + methods, + streams: HashMap::new(), + }, + ); + services +} + +fn build_echo_request(payload: &[u8]) -> Request { + let mut req = Request::new(); + req.service = TEST_SERVICE.to_string(); + req.method = TEST_METHOD.to_string(); + req.payload = payload.to_vec(); + req.timeout_nano = 5_000_000_000; // 5 seconds + req +} + +/// Brief yield to allow the server's background accept loop to start. +/// +/// We intentionally avoid a "connect-and-retry" readiness probe because +/// any connection to the server would trigger the accept hook, altering +/// the expected hook-call counts that tests assert on. +async fn wait_for_server_ready() { + tokio::task::yield_now().await; + sleep(Duration::from_millis(30)).await; +} + +// ── Tests ────────────────────────────────────────────────────────────────── + +#[tokio::test] +async fn test_server_accept_hook_called_on_connection() { + let sock_path = temp_unix_socket_path(); + let hook = TestAcceptHook::new_accepting(true, false); + let hook_count = hook.call_count.clone(); + + let mut server = Server::new() + .bind(&format!("unix://{}", sock_path)) + .unwrap() + .set_accept_hook(hook) + .register_service(build_test_service()); + + server.start().await.unwrap(); + wait_for_server_ready().await; + + // Connect client (no hook on client side) + let client = Client::connect(&format!("unix://{}", sock_path)) + .await + .unwrap(); + + // Send a request to trigger connection establishment. + // The accept hook fires during connection acceptance, which completes + // before the request is processed, so no extra sleep is needed. + let req = build_echo_request(b"hello"); + let _resp = client.request(req).await.unwrap(); + + // Verify hook was called exactly once + assert_eq!( + hook_count.load(Ordering::SeqCst), + 1, + "AcceptHook should be called once" + ); + + // Cleanup + server.shutdown().await.unwrap(); + cleanup_socket_file(&sock_path); +} + +#[tokio::test] +async fn test_server_accept_hook_rejects_connection() { + let sock_path = temp_unix_socket_path(); + let hook = TestAcceptHook::new_rejecting(); + let hook_count = hook.call_count.clone(); + + let mut server = Server::new() + .bind(&format!("unix://{}", sock_path)) + .unwrap() + .set_accept_hook(hook) + .register_service(build_test_service()); + + server.start().await.unwrap(); + wait_for_server_ready().await; + + // Connect client + let client = Client::connect(&format!("unix://{}", sock_path)) + .await + .unwrap(); + + // Request should fail because server rejects the connection + let req = build_echo_request(b"hello"); + let result = client.request(req).await; + assert!( + result.is_err(), + "Expected request to fail when server rejects connection" + ); + + // Wait for server to process the rejection + sleep(Duration::from_millis(200)).await; + + // Verify hook was called + assert!( + hook_count.load(Ordering::SeqCst) >= 1, + "AcceptHook should be called at least once" + ); + + // Cleanup + server.shutdown().await.unwrap(); + cleanup_socket_file(&sock_path); +} + +#[tokio::test] +async fn test_connection_data_propagated_to_handler() { + let sock_path = temp_unix_socket_path(); + let hook = TestAcceptHook::new_accepting(true, false); + + let mut server = Server::new() + .bind(&format!("unix://{}", sock_path)) + .unwrap() + .set_accept_hook(hook) + .register_service(build_test_service()); + + server.start().await.unwrap(); + wait_for_server_ready().await; + + let client = Client::connect(&format!("unix://{}", sock_path)) + .await + .unwrap(); + + let req = build_echo_request(b"test_data"); + let resp = client.request(req).await.unwrap(); + + // EchoHandler prepends "role:test_client|" if peer_role is in connection_data + let resp_str = String::from_utf8_lossy(&resp.payload); + assert!( + resp_str.starts_with("role:test_client|"), + "Expected connection data to be propagated to handler, got: {}", + resp_str + ); + + server.shutdown().await.unwrap(); + cleanup_socket_file(&sock_path); +} + +#[tokio::test] +async fn test_xor_transform_on_wire() { + let sock_path = temp_unix_socket_path(); + + // Server with XOR transform + let server_hook = TestAcceptHook::new_accepting(false, true); + let mut server = Server::new() + .bind(&format!("unix://{}", sock_path)) + .unwrap() + .set_accept_hook(server_hook) + .register_service(build_test_service()); + + server.start().await.unwrap(); + wait_for_server_ready().await; + + // Client with XOR transform (symmetric) + let client_hook = TestConnectHook::new(false, true); + let socket = ttrpc::asynchronous::transport::Socket::connect(&format!("unix://{}", sock_path)) + .await + .unwrap(); + let client = Client::with_hook(socket, client_hook).unwrap(); + + // 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(); + + // Response should be the echo of the original payload (after decrypt→encrypt→decrypt) + assert_eq!( + &resp.payload, original_payload, + "XOR roundtrip failed: expected original payload back" + ); + + server.shutdown().await.unwrap(); + cleanup_socket_file(&sock_path); +} + +#[tokio::test] +async fn test_client_connect_hook_called() { + let sock_path = temp_unix_socket_path(); + + let mut server = Server::new() + .bind(&format!("unix://{}", sock_path)) + .unwrap() + .register_service(build_test_service()); + + server.start().await.unwrap(); + wait_for_server_ready().await; + + let client_hook = TestConnectHook::new(true, false); + let hook_count = client_hook.call_count.clone(); + + let socket = ttrpc::asynchronous::transport::Socket::connect(&format!("unix://{}", sock_path)) + .await + .unwrap(); + let client = Client::with_hook(socket, client_hook).unwrap(); + + let req = build_echo_request(b"hook_test"); + let _resp = client.request(req).await.unwrap(); + + assert_eq!( + hook_count.load(Ordering::SeqCst), + 1, + "ConnectHook should be called exactly once" + ); + + server.shutdown().await.unwrap(); + cleanup_socket_file(&sock_path); +} + +#[tokio::test] +async fn test_symmetric_hooks_with_data_and_transform() { + let sock_path = temp_unix_socket_path(); + + // Server: attach data + XOR transform + let server_hook = TestAcceptHook::new_accepting(true, true); + let server_count = server_hook.call_count.clone(); + + let mut server = Server::new() + .bind(&format!("unix://{}", sock_path)) + .unwrap() + .set_accept_hook(server_hook) + .register_service(build_test_service()); + + server.start().await.unwrap(); + wait_for_server_ready().await; + + // Client: attach data + XOR transform + let client_hook = TestConnectHook::new(true, true); + let client_count = client_hook.call_count.clone(); + + let socket = ttrpc::asynchronous::transport::Socket::connect(&format!("unix://{}", sock_path)) + .await + .unwrap(); + let client = Client::with_hook(socket, client_hook).unwrap(); + + // Multiple requests to verify transform persists across messages + for i in 0..5 { + let payload = format!("message_{}", i); + let req = build_echo_request(payload.as_bytes()); + let resp = client.request(req).await.unwrap(); + + // EchoHandler prepends role info from connection_data + let resp_str = String::from_utf8_lossy(&resp.payload); + assert!( + resp_str.starts_with("role:test_client|"), + "Request {}: expected role prefix, got: {}", + i, + resp_str + ); + } + + // Both hooks should have been called exactly once + assert_eq!(server_count.load(Ordering::SeqCst), 1); + assert_eq!(client_count.load(Ordering::SeqCst), 1); + + server.shutdown().await.unwrap(); + cleanup_socket_file(&sock_path); +} + +#[tokio::test] +async fn test_no_hook_plaintext_passthrough() { + let sock_path = temp_unix_socket_path(); + + // Server without hook + let mut server = Server::new() + .bind(&format!("unix://{}", sock_path)) + .unwrap() + .register_service(build_test_service()); + + server.start().await.unwrap(); + wait_for_server_ready().await; + + // Client without hook + let client = Client::connect(&format!("unix://{}", sock_path)) + .await + .unwrap(); + + let original = b"plaintext_message"; + let req = build_echo_request(original); + let resp = client.request(req).await.unwrap(); + + // No transform, no connection data — pure echo + assert_eq!( + &resp.payload, original, + "Plaintext passthrough should return exact payload" + ); + + server.shutdown().await.unwrap(); + cleanup_socket_file(&sock_path); +} + +// ── Streaming tests ───────────────────────────────────────────────────────── + +use ttrpc::asynchronous::StreamHandler; +use ttrpc::r#async::StreamInner; + +const TEST_STREAM_METHOD: &str = "DuplexEcho"; + +#[derive(Debug)] +struct DuplexEchoHandler; + +#[async_trait] +impl StreamHandler for DuplexEchoHandler { + async fn handler( + &self, + ctx: TtrpcContext, + mut stream: StreamInner, + ) -> ttrpc::Result> { + use ttrpc::proto::Codec; + // Echo loop: receive Request, send back Response with same payload + loop { + match stream.recv().await { + Ok(data) => { + // Decode as Request to extract payload + let req = Request::decode(&data).unwrap_or_else(|_| { + let mut r = Request::new(); + r.payload = data.clone(); + r + }); + // Build Response with same payload and encode it + let mut resp = Response::new(); + resp.set_status(Status::default()); + resp.payload = req.payload; + let encoded = resp + .encode() + .map_err(|e| ttrpc::Error::Others(format!("encode resp failed: {}", e)))?; + stream.send(encoded).await?; + } + Err(ttrpc::Error::Eof) => break, + Err(ttrpc::Error::RemoteClosed) => break, + Err(e) => return Err(e), + } + } + // Send final response + let mut resp = Response::new(); + resp.set_status(Status::default()); + if let Some(role) = ctx.connection_data.get_typed::("peer_role") { + resp.payload = format!("stream_done:{}", role).into_bytes(); + } else { + resp.payload = b"stream_done".to_vec(); + } + Ok(Some(resp)) + } +} + +fn build_test_service_with_stream() -> HashMap { + let mut methods: HashMap> = HashMap::new(); + methods.insert(TEST_METHOD.to_string(), Box::new(EchoHandler)); + + let mut streams: HashMap> = HashMap::new(); + streams.insert(TEST_STREAM_METHOD.to_string(), Arc::new(DuplexEchoHandler)); + + let mut services = HashMap::new(); + services.insert(TEST_SERVICE.to_string(), Service { methods, streams }); + services +} + +fn build_stream_request() -> Request { + let mut req = Request::new(); + req.service = TEST_SERVICE.to_string(); + req.method = TEST_STREAM_METHOD.to_string(); + req.timeout_nano = 5_000_000_000; + req +} + +#[tokio::test] +async fn test_streaming_with_xor_transform() { + let sock_path = temp_unix_socket_path(); + + // Server with XOR transform + stream handler + let server_hook = TestAcceptHook::new_accepting(false, true); + let mut server = Server::new() + .bind(&format!("unix://{}", sock_path)) + .unwrap() + .set_accept_hook(server_hook) + .register_service(build_test_service_with_stream()); + + server.start().await.unwrap(); + wait_for_server_ready().await; + + // Client with XOR transform + let client_hook = TestConnectHook::new(false, true); + let socket = ttrpc::asynchronous::transport::Socket::connect(&format!("unix://{}", sock_path)) + .await + .unwrap(); + let client = Client::with_hook(socket, client_hook).unwrap(); + + // Open a duplex stream + let inner = client + .new_stream(build_stream_request(), true, true) + .await + .unwrap(); + let mut stream = ttrpc::r#async::ClientStream::::new(inner); + + // Send 3 messages and verify echo + for i in 0u32..3 { + let mut msg = Request::new(); + msg.payload = format!("stream_msg_{}", i).into_bytes(); + stream.send(&msg).await.unwrap(); + + let echoed = stream.recv().await.unwrap(); + assert_eq!(echoed.payload, msg.payload, "Stream echo {} failed", i); + } + + // Close the client side + stream.close_send().await.unwrap(); + + // Wait for server to process close and send final response + sleep(Duration::from_millis(200)).await; + + server.shutdown().await.unwrap(); + cleanup_socket_file(&sock_path); +} + +#[tokio::test] +async fn test_streaming_without_transform_plaintext() { + let sock_path = temp_unix_socket_path(); + + // Server without hook (no transform) + let mut server = Server::new() + .bind(&format!("unix://{}", sock_path)) + .unwrap() + .register_service(build_test_service_with_stream()); + + server.start().await.unwrap(); + wait_for_server_ready().await; + + let client = Client::connect(&format!("unix://{}", sock_path)) + .await + .unwrap(); + + let inner = client + .new_stream(build_stream_request(), true, true) + .await + .unwrap(); + let mut stream = ttrpc::r#async::ClientStream::::new(inner); + + let mut msg = Request::new(); + msg.payload = b"plain_stream_data".to_vec(); + stream.send(&msg).await.unwrap(); + + let echoed = stream.recv().await.unwrap(); + assert_eq!(echoed.payload, b"plain_stream_data"); + + stream.close_send().await.unwrap(); + sleep(Duration::from_millis(200)).await; + + server.shutdown().await.unwrap(); + cleanup_socket_file(&sock_path); +} + +// ── ConnectHook error path ────────────────────────────────────────────────── + +#[tokio::test] +async fn test_client_connect_hook_rejected_fails_connection() { + let sock_path = temp_unix_socket_path(); + + let mut server = Server::new() + .bind(&format!("unix://{}", sock_path)) + .unwrap() + .register_service(build_test_service()); + + server.start().await.unwrap(); + wait_for_server_ready().await; + + // ConnectHook that returns Err — client creation must fail. + #[derive(Debug)] + struct RejectingConnectHook; + impl ConnectHook for RejectingConnectHook { + fn on_connect(&self, _fd: RawFd) -> Result { + Err(HookError::Rejected("client self-reject".into())) + } + } + + let socket = ttrpc::asynchronous::transport::Socket::connect(&format!("unix://{}", sock_path)) + .await + .unwrap(); + let err = match Client::with_hook(socket, RejectingConnectHook) { + Ok(_) => panic!("Expected Client::with_hook to fail when connect hook rejects"), + Err(e) => e, + }; + let err_str = format!("{}", err); + assert!( + err_str.contains("client self-reject") || err_str.contains("connect hook rejected"), + "Expected hook rejection error, got: {}", + err_str + ); + + server.shutdown().await.unwrap(); + cleanup_socket_file(&sock_path); +} + +// ── Multiple concurrent connections ───────────────────────────────────────── + +#[tokio::test] +async fn test_multiple_concurrent_connections_with_transform() { + let sock_path = temp_unix_socket_path(); + + let server_hook = TestAcceptHook::new_accepting(true, true); + let server_count = server_hook.call_count.clone(); + + let mut server = Server::new() + .bind(&format!("unix://{}", sock_path)) + .unwrap() + .set_accept_hook(server_hook) + .register_service(build_test_service()); + + server.start().await.unwrap(); + wait_for_server_ready().await; + + // Create 3 concurrent clients, each with XOR transform + let mut handles = vec![]; + for i in 0..3 { + let path = sock_path.clone(); + handles.push(tokio::spawn(async move { + let client_hook = TestConnectHook::new(true, true); + let socket = + ttrpc::asynchronous::transport::Socket::connect(&format!("unix://{}", path)) + .await + .unwrap(); + let client = Client::with_hook(socket, client_hook).unwrap(); + + let payload = format!("concurrent_client_{}", i); + let req = build_echo_request(payload.as_bytes()); + let resp = client.request(req).await.unwrap(); + + // Verify response contains role prefix (from server's connection data) + let resp_str = String::from_utf8_lossy(&resp.payload); + assert!( + resp_str.starts_with("role:test_client|"), + "Client {}: expected role prefix, got: {}", + i, + resp_str + ); + })); + } + + for handle in handles { + handle.await.unwrap(); + } + + // Each connection should have triggered AcceptHook once + assert_eq!( + server_count.load(Ordering::SeqCst), + 3, + "AcceptHook should be called once per connection" + ); + + server.shutdown().await.unwrap(); + cleanup_socket_file(&sock_path); +} + +// ── Coverage gap tests ──────────────────────────────────────────────────────── + +// Test 1: streaming_client=false with XOR transform +// +// When streaming_client=false with a non-empty request payload, the server's +// handle_stream() creates a synthetic DATA message from the REQUEST payload +// that was already decrypted by handle_request() (Injection Point 2/10). +// The fix: use StreamMsg::PreDecoded so StreamReceiver::recv() +// (Injection Point 10/10) passes it through without re-applying +// transform_inbound. Works for any PayloadTransform type. +#[tokio::test] +async fn test_streaming_client_false_with_xor_transform() { + let sock_path = temp_unix_socket_path(); + + let server_hook = TestAcceptHook::new_accepting(false, true); + let mut server = Server::new() + .bind(&format!("unix://{}", sock_path)) + .unwrap() + .set_accept_hook(server_hook) + .register_service(build_test_service_with_stream()); + + server.start().await.unwrap(); + wait_for_server_ready().await; + + let client_hook = TestConnectHook::new(false, true); + let socket = ttrpc::asynchronous::transport::Socket::connect(&format!("unix://{}", sock_path)) + .await + .unwrap(); + let client = Client::with_hook(socket, client_hook).unwrap(); + + // streaming_client = false, non-empty payload → triggers handle_stream + // faked DATA with re-encryption → StreamReceiver::recv() decrypts correctly + let mut req = build_stream_request(); + req.payload = b"initial_payload".to_vec(); + + let inner = client.new_stream(req, false, true).await.unwrap(); + let mut stream = ttrpc::r#async::ClientStream::::new(inner); + + // Server handler (DuplexEchoHandler) receives the initial payload as + // the first DATA message, echoes it back, then enters the recv loop. + // The echo should contain the original payload bytes — proving that + // the double-transform was correctly compensated. + let echoed = stream.recv().await.unwrap(); + assert_eq!( + echoed.payload, b"initial_payload", + "streaming_client=false initial payload should survive transform fix" + ); + + // With streaming_client=false, the client cannot send or close_send. + // Just drop the stream to signal we're done. + drop(stream); + sleep(Duration::from_millis(200)).await; + + server.shutdown().await.unwrap(); + cleanup_socket_file(&sock_path); +} + +// Test 2: Stream close semantics — close_send + post-close send error +#[tokio::test] +async fn test_stream_close_send_then_send_fails() { + let sock_path = temp_unix_socket_path(); + + let mut server = Server::new() + .bind(&format!("unix://{}", sock_path)) + .unwrap() + .register_service(build_test_service_with_stream()); + + server.start().await.unwrap(); + wait_for_server_ready().await; + + let client = Client::connect(&format!("unix://{}", sock_path)) + .await + .unwrap(); + + let inner = client + .new_stream(build_stream_request(), true, true) + .await + .unwrap(); + let stream = ttrpc::r#async::ClientStream::::new(inner); + + // Send one message successfully + let mut msg = Request::new(); + msg.payload = b"before_close".to_vec(); + stream.send(&msg).await.unwrap(); + + // Close the send side + stream.close_send().await.unwrap(); + + // Sending after close should fail with LocalClosed + let result = stream.send(&msg).await; + assert!(result.is_err(), "send after close_send should fail"); + let err_msg = format!("{}", result.unwrap_err()); + assert!( + err_msg.contains("LocalClosed") || err_msg.contains("closed"), + "Expected LocalClosed error, got: {}", + err_msg + ); + + sleep(Duration::from_millis(100)).await; + server.shutdown().await.unwrap(); + cleanup_socket_file(&sock_path); +} + +// Test 3: Server shutdown during active streaming +// +// Verifies that server.shutdown() completes cleanly even when a stream +// is still open on the client side. The post-shutdown stream operations +// are wrapped in a timeout because the client may not detect the closed +// connection immediately. +#[tokio::test] +async fn test_server_shutdown_during_active_stream() { + let sock_path = temp_unix_socket_path(); + + let mut server = Server::new() + .bind(&format!("unix://{}", sock_path)) + .unwrap() + .register_service(build_test_service_with_stream()); + + server.start().await.unwrap(); + wait_for_server_ready().await; + + let client = Client::connect(&format!("unix://{}", sock_path)) + .await + .unwrap(); + + let inner = client + .new_stream(build_stream_request(), true, true) + .await + .unwrap(); + let mut stream = ttrpc::r#async::ClientStream::::new(inner); + + // Send a message and receive echo — verifies stream works before shutdown + let mut msg = Request::new(); + msg.payload = b"before_shutdown".to_vec(); + stream.send(&msg).await.unwrap(); + let echoed = stream.recv().await.unwrap(); + assert_eq!(echoed.payload, b"before_shutdown"); + + // Shutdown server while stream is still open — must complete without deadlock + server.shutdown().await.unwrap(); + cleanup_socket_file(&sock_path); + + // Post-shutdown: recv() should detect the closed connection (with timeout) + sleep(Duration::from_millis(200)).await; + let result = tokio::time::timeout(Duration::from_secs(3), stream.recv()).await; + match result { + Ok(Ok(_)) => {} // Unexpected but not a failure — server may have sent a final message + Ok(Err(_)) => {} // Expected: recv detects closed connection + Err(_) => {} // Timeout: client hasn't detected shutdown yet — acceptable + } +} + +// Test 4: ConnectHook receives valid raw_fd +#[tokio::test] +async fn test_connect_hook_receives_valid_raw_fd() { + let sock_path = temp_unix_socket_path(); + + let mut server = Server::new() + .bind(&format!("unix://{}", sock_path)) + .unwrap() + .register_service(build_test_service()); + + server.start().await.unwrap(); + wait_for_server_ready().await; + + // ConnectHook that captures the fd it receives + let captured_fd = Arc::new(std::sync::atomic::AtomicI32::new(-1)); + let captured_fd_clone = captured_fd.clone(); + + #[derive(Debug)] + struct FdCapturingHook { + captured: Arc, + } + impl ConnectHook for FdCapturingHook { + fn on_connect(&self, fd: RawFd) -> Result { + self.captured.store(fd, Ordering::SeqCst); + Ok(HookOutput { + data: ConnectionData::new(), + payload_transform: None, + }) + } + } + + let socket = ttrpc::asynchronous::transport::Socket::connect(&format!("unix://{}", sock_path)) + .await + .unwrap(); + let client = Client::with_hook( + socket, + FdCapturingHook { + captured: captured_fd_clone, + }, + ) + .unwrap(); + + // Trigger connection + let req = build_echo_request(b"fd_check"); + let _resp = client.request(req).await.unwrap(); + + let fd = captured_fd.load(Ordering::SeqCst); + assert!( + fd >= 0, + "ConnectHook should receive a valid raw fd (>= 0), got: {}", + fd + ); + + server.shutdown().await.unwrap(); + cleanup_socket_file(&sock_path); +} + +// Test 5: Server-initiated stream close — client receives Eof +// +// Handler that sends 2 echo responses then returns (server closes the stream). +#[derive(Debug)] +struct LimitedEchoHandler; + +#[async_trait] +impl StreamHandler for LimitedEchoHandler { + async fn handler( + &self, + _ctx: TtrpcContext, + mut stream: StreamInner, + ) -> ttrpc::Result> { + // Receive one message, echo it as DATA, then return final response + if let Ok(data) = stream.recv().await { + // Echo the raw data back as a DATA message + stream.send(data).await?; + } + // Return final response — triggers server stream close + let mut resp = Response::new(); + resp.set_status(Status::default()); + resp.payload = b"limited_done".to_vec(); + Ok(Some(resp)) + } +} + +fn build_limited_echo_service() -> HashMap { + let mut streams: HashMap> = HashMap::new(); + streams.insert("LimitedEcho".to_string(), Arc::new(LimitedEchoHandler)); + + let mut services = HashMap::new(); + services.insert( + TEST_SERVICE.to_string(), + Service { + methods: HashMap::new(), + streams, + }, + ); + services +} + +#[tokio::test] +async fn test_server_initiated_stream_close_client_gets_final_response() { + let sock_path = temp_unix_socket_path(); + + let mut server = Server::new() + .bind(&format!("unix://{}", sock_path)) + .unwrap() + .register_service(build_limited_echo_service()); + + server.start().await.unwrap(); + wait_for_server_ready().await; + + let client = Client::connect(&format!("unix://{}", sock_path)) + .await + .unwrap(); + + // Build stream request for LimitedEcho method + let mut req = Request::new(); + req.service = TEST_SERVICE.to_string(); + req.method = "LimitedEcho".to_string(); + req.timeout_nano = 5_000_000_000; + + let inner = client.new_stream(req, true, true).await.unwrap(); + let mut stream = ttrpc::r#async::ClientStream::::new(inner); + + // Send one message, server echoes it as DATA, then returns final response + let mut msg = Request::new(); + msg.payload = b"limited_test".to_vec(); + stream.send(&msg).await.unwrap(); + + // First recv: gets the echo DATA message (raw bytes, decoded as Response). + // We don't check the echo content because it's raw StreamInner::send() bytes, + // but we do assert success so a failed echo doesn't silently cascade. + let _echoed = stream.recv().await.expect("echo recv should succeed"); + + // Second recv: gets the final RESPONSE from the handler's return value + // Use timeout to prevent hanging + let result = tokio::time::timeout(Duration::from_secs(5), stream.recv()).await; + match result { + Ok(Ok(final_resp)) => { + assert_eq!( + final_resp.payload, b"limited_done", + "Expected final response payload 'limited_done'" + ); + } + 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 + ); + } + Err(_) => { + panic!("Timed out waiting for server's final response or stream close"); + } + } + + server.shutdown().await.unwrap(); + cleanup_socket_file(&sock_path); +} + +// Test 6: Unary request timeout (server-side DEADLINE_EXCEEDED + client-side timeout) +// +// Covers the timeout code path in server.rs handle_method() (tokio::time::timeout +// around the handler) and client.rs request() (tokio::time::timeout on the response). +#[tokio::test] +async fn test_unary_request_timeout() { + let sock_path = temp_unix_socket_path(); + + // Server with SlowEchoHandler (sleeps 1s) + let mut server = Server::new() + .bind(&format!("unix://{}", sock_path)) + .unwrap() + .register_service(build_slow_echo_service()); + + server.start().await.unwrap(); + wait_for_server_ready().await; + + let client = Client::connect(&format!("unix://{}", sock_path)) + .await + .unwrap(); + + // Build request with very short timeout (100ms) — handler sleeps 1s + let mut req = Request::new(); + req.service = TEST_SERVICE.to_string(); + req.method = TEST_METHOD.to_string(); + req.payload = b"timeout_test".to_vec(); + req.timeout_nano = 100_000_000; // 100ms + + let result = client.request(req).await; + assert!(result.is_err(), "Expected timeout error"); + let err = result.unwrap_err(); + let err_str = format!("{:?}", err); + assert!( + err_str.contains("timeout") + || err_str.contains("Timeout") + || err_str.contains("DEADLINE_EXCEEDED"), + "Expected timeout-related error, got: {}", + err_str + ); + + server.shutdown().await.unwrap(); + cleanup_socket_file(&sock_path); +} + +// Test 7: Streaming with streaming_server=false — client rejects DATA messages +// +// When streaming_server=false, the client's StreamReceiver has receivable=false. +// If the server sends DATA messages, recv() returns an error because the client +// does not expect streaming data from the server. +#[tokio::test] +async fn test_streaming_server_false_rejects_data() { + let sock_path = temp_unix_socket_path(); + + let server_hook = TestAcceptHook::new_accepting(false, true); + let mut server = Server::new() + .bind(&format!("unix://{}", sock_path)) + .unwrap() + .set_accept_hook(server_hook) + .register_service(build_test_service_with_stream()); + + server.start().await.unwrap(); + wait_for_server_ready().await; + + let client_hook = TestConnectHook::new(false, true); + let socket = ttrpc::asynchronous::transport::Socket::connect(&format!("unix://{}", sock_path)) + .await + .unwrap(); + let client = Client::with_hook(socket, client_hook).unwrap(); + + // streaming_client=true, streaming_server=false + let inner = client + .new_stream(build_stream_request(), true, false) + .await + .unwrap(); + let mut stream = ttrpc::r#async::ClientStream::::new(inner); + + // Send a message — server's DuplexEchoHandler will echo it back as DATA + let mut msg = Request::new(); + msg.payload = b"test_data".to_vec(); + stream.send(&msg).await.unwrap(); + + // recv() should fail because streaming_server=false (receivable=false) + let result = tokio::time::timeout(Duration::from_secs(3), stream.recv()).await; + match result { + Ok(Err(e)) => { + let err_str = format!("{}", e); + assert!( + err_str.contains("non-streaming server") + || err_str.contains("Eof") + || err_str.contains("RemoteClosed") + || err_str.contains("closed"), + "Expected non-streaming-server error or close, got: {}", + err_str + ); + } + Ok(Ok(_)) => { + panic!("recv() should have failed for streaming_server=false, but succeeded"); + } + Err(_) => { + panic!("recv() timed out — expected non-streaming-server error"); + } + } + + drop(stream); + sleep(Duration::from_millis(200)).await; + server.shutdown().await.unwrap(); + cleanup_socket_file(&sock_path); +} + +// Test 8: Multiple concurrent streams on one connection +#[tokio::test] +async fn test_multiple_concurrent_streams_on_one_connection() { + let sock_path = temp_unix_socket_path(); + + let server_hook = TestAcceptHook::new_accepting(false, true); + let mut server = Server::new() + .bind(&format!("unix://{}", sock_path)) + .unwrap() + .set_accept_hook(server_hook) + .register_service(build_test_service_with_stream()); + + server.start().await.unwrap(); + wait_for_server_ready().await; + + // Single client with XOR transform + let client_hook = TestConnectHook::new(false, true); + let socket = ttrpc::asynchronous::transport::Socket::connect(&format!("unix://{}", sock_path)) + .await + .unwrap(); + let client = Client::with_hook(socket, client_hook).unwrap(); + + // Open 2 independent streams on the same connection + let inner1 = client + .new_stream(build_stream_request(), true, true) + .await + .unwrap(); + let mut stream1 = ttrpc::r#async::ClientStream::::new(inner1); + + let inner2 = client + .new_stream(build_stream_request(), true, true) + .await + .unwrap(); + let mut stream2 = ttrpc::r#async::ClientStream::::new(inner2); + + // Interleave sends/receives across both streams + let mut msg1 = Request::new(); + msg1.payload = b"stream1_data".to_vec(); + stream1.send(&msg1).await.unwrap(); + + let mut msg2 = Request::new(); + msg2.payload = b"stream2_data".to_vec(); + stream2.send(&msg2).await.unwrap(); + + // Receive in reverse order to verify independence + let echoed2 = stream2.recv().await.unwrap(); + assert_eq!(echoed2.payload, b"stream2_data", "Stream 2 echo mismatch"); + + let echoed1 = stream1.recv().await.unwrap(); + assert_eq!(echoed1.payload, b"stream1_data", "Stream 1 echo mismatch"); + + // Close both streams + stream1.close_send().await.unwrap(); + stream2.close_send().await.unwrap(); + + sleep(Duration::from_millis(200)).await; + server.shutdown().await.unwrap(); + cleanup_socket_file(&sock_path); +} diff --git a/tests/hook_integration_sync_unix.rs b/tests/hook_integration_sync_unix.rs new file mode 100644 index 00000000..3d41d6f8 --- /dev/null +++ b/tests/hook_integration_sync_unix.rs @@ -0,0 +1,349 @@ +// Copyright 2026 Alibaba Cloud. All rights reserved. +// +// SPDX-License-Identifier: Apache-2.0 +// +#![cfg(all(feature = "sync", feature = "security_extension"))] +//! Integration tests for sync AcceptHook, ConnectHook, and PayloadTransform. + +mod common; + +use std::collections::HashMap; +use std::os::unix::io::RawFd; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::Arc; +use std::thread; +use std::time::Duration; + +use ttrpc::proto::{Request, Response}; +use ttrpc::security_extension::{ + AcceptHook, ConnectHook, ConnectionData, ConnectionDataExt, HookError, HookOutput, +}; +use ttrpc::sync::{Client, MethodHandler, Server, TtrpcContext}; +use ttrpc::{get_status, Code}; + +use common::{cleanup_socket_file, temp_unix_socket_path, XorPayloadTransform}; + +// ── Test constants ────────────────────────────────────────────────────────── + +const TEST_SERVICE: &str = "test.SyncTestService"; +const TEST_METHOD: &str = "Echo"; + +// ── Test hooks ────────────────────────────────────────────────────────────── + +#[derive(Debug)] +struct CountingAcceptHook { + call_count: Arc, + reject: bool, +} + +impl AcceptHook for CountingAcceptHook { + fn on_accept(&self, _fd: RawFd) -> Result { + self.call_count.fetch_add(1, Ordering::SeqCst); + if self.reject { + return Err(HookError::Rejected("rejected by hook".into())); + } + let mut data = ConnectionData::new(); + data.insert("peer_role".into(), Box::new(String::from("sync-client"))); + Ok(HookOutput { + data, + payload_transform: Some(Box::new(XorPayloadTransform)), + }) + } +} + +#[derive(Debug)] +struct CountingConnectHook { + call_count: Arc, +} + +impl ConnectHook for CountingConnectHook { + fn on_connect(&self, _fd: RawFd) -> Result { + self.call_count.fetch_add(1, Ordering::SeqCst); + Ok(HookOutput { + data: ConnectionData::new(), + payload_transform: Some(Box::new(XorPayloadTransform)), + }) + } +} + +#[derive(Debug)] +struct RejectingConnectHook; + +impl ConnectHook for RejectingConnectHook { + fn on_connect(&self, _fd: RawFd) -> Result { + Err(HookError::Rejected("client rejected".into())) + } +} + +// ── Sync MethodHandler ────────────────────────────────────────────────────── + +#[derive(Debug)] +struct SyncEchoHandler; + +impl MethodHandler for SyncEchoHandler { + fn handler(&self, ctx: TtrpcContext, req: Request) -> ttrpc::Result<()> { + let mut resp = Response::new(); + resp.set_status(get_status(Code::OK, "".to_string())); + // Echo the request payload back, optionally prefixed with connection data + let mut payload = Vec::new(); + if let Some(role) = ctx.conn_ctx.data.get_typed::("peer_role") { + payload.extend_from_slice(format!("role:{}|", role).as_bytes()); + } + payload.extend_from_slice(&req.payload); + resp.payload = payload; + ctx.respond(ctx.mh.stream_id, resp)?; + Ok(()) + } +} + +fn build_sync_echo_service() -> HashMap> { + let mut methods: HashMap> = HashMap::new(); + let path = format!("/{}/{}", TEST_SERVICE, TEST_METHOD); + methods.insert(path, Box::new(SyncEchoHandler)); + methods +} + +// ── Helper: wait for server to be ready ───────────────────────────────────── + +fn wait_for_server_ready() { + thread::sleep(Duration::from_millis(50)); +} + +// ── Tests ─────────────────────────────────────────────────────────────────── + +#[test] +fn test_sync_no_hook_plaintext_passthrough() { + let sock_path = temp_unix_socket_path(); + let path = format!("unix://{}", sock_path); + + let mut server = Server::new() + .bind(&path) + .unwrap() + .register_service(build_sync_echo_service()); + server.start().unwrap(); + wait_for_server_ready(); + + let client = Client::connect(&path).unwrap(); + + let mut req = Request::new(); + req.set_service(TEST_SERVICE.to_string()); + req.set_method(TEST_METHOD.to_string()); + req.payload = b"hello".to_vec(); + + let resp = client.request(req).unwrap(); + assert_eq!(resp.payload, b"hello"); + + server.shutdown(); + cleanup_socket_file(&sock_path); +} + +#[test] +fn test_sync_accept_hook_called_on_connection() { + let sock_path = temp_unix_socket_path(); + let path = format!("unix://{}", sock_path); + + let hook_count = Arc::new(AtomicUsize::new(0)); + let hook = CountingAcceptHook { + call_count: hook_count.clone(), + reject: false, + }; + + let mut server = Server::new() + .bind(&path) + .unwrap() + .register_service(build_sync_echo_service()) + .set_accept_hook(hook); + server.start().unwrap(); + wait_for_server_ready(); + + // Connect (triggers accept hook) + let hook_count_client = Arc::new(AtomicUsize::new(0)); + let connect_hook = CountingConnectHook { + call_count: hook_count_client.clone(), + }; + let fd = create_connected_fd(&sock_path); + let client = Client::with_hook(fd, connect_hook).unwrap(); + + // Send a request to ensure the server processes the accept + let mut req = Request::new(); + req.set_service(TEST_SERVICE.to_string()); + req.set_method(TEST_METHOD.to_string()); + req.payload = b"test".to_vec(); + + let resp = client.request(req).unwrap(); + + // Hook should have been called once on accept + assert_eq!(hook_count.load(Ordering::SeqCst), 1); + assert_eq!(hook_count_client.load(Ordering::SeqCst), 1); + + // Response should have connection data prefix + let resp_str = String::from_utf8_lossy(&resp.payload); + assert!(resp_str.starts_with("role:sync-client|")); + + server.shutdown(); + cleanup_socket_file(&sock_path); +} + +#[test] +fn test_sync_accept_hook_rejects_connection() { + let sock_path = temp_unix_socket_path(); + let path = format!("unix://{}", sock_path); + + let hook_count = Arc::new(AtomicUsize::new(0)); + let hook = CountingAcceptHook { + call_count: hook_count.clone(), + reject: true, + }; + + let mut server = Server::new() + .bind(&path) + .unwrap() + .register_service(build_sync_echo_service()) + .set_accept_hook(hook); + server.start().unwrap(); + wait_for_server_ready(); + + // Try to connect — server will reject via hook + let fd = create_connected_fd(&sock_path); + let client = Client::new(fd).unwrap(); + + // Send a request to trigger the accept loop processing the connection + let mut req = Request::new(); + req.set_service(TEST_SERVICE.to_string()); + req.set_method(TEST_METHOD.to_string()); + req.payload = b"fail".to_vec(); + req.timeout_nano = 2_000_000_000; // 2 second timeout + + let result = client.request(req); + + // Hook was called + assert_eq!(hook_count.load(Ordering::SeqCst), 1); + + // Request should fail (server closed the connection after hook rejection) + assert!(result.is_err()); + + server.shutdown(); + cleanup_socket_file(&sock_path); +} + +#[test] +fn test_sync_connect_hook_rejected_fails() { + let sock_path = temp_unix_socket_path(); + let path = format!("unix://{}", sock_path); + + let mut server = Server::new() + .bind(&path) + .unwrap() + .register_service(build_sync_echo_service()); + server.start().unwrap(); + wait_for_server_ready(); + + let fd = create_connected_fd(&sock_path); + let result = Client::with_hook(fd, RejectingConnectHook); + assert!(result.is_err()); + if let Err(e) = result { + let err_msg = format!("{}", e); + assert!(err_msg.contains("rejected")); + } + + server.shutdown(); + cleanup_socket_file(&sock_path); +} + +#[test] +fn test_sync_xor_transform_roundtrip() { + let sock_path = temp_unix_socket_path(); + let path = format!("unix://{}", sock_path); + + let hook_count = Arc::new(AtomicUsize::new(0)); + let hook = CountingAcceptHook { + call_count: hook_count.clone(), + reject: false, + }; + + let mut server = Server::new() + .bind(&path) + .unwrap() + .register_service(build_sync_echo_service()) + .set_accept_hook(hook); + server.start().unwrap(); + wait_for_server_ready(); + + let client_hook_count = Arc::new(AtomicUsize::new(0)); + let connect_hook = CountingConnectHook { + call_count: client_hook_count.clone(), + }; + let fd = create_connected_fd(&sock_path); + let client = Client::with_hook(fd, connect_hook).unwrap(); + + // Send various payload sizes + for payload in &[ + b"a".to_vec(), + b"ab".to_vec(), + b"abc".to_vec(), + vec![0u8; 1024], + ] { + let mut req = Request::new(); + req.set_service(TEST_SERVICE.to_string()); + req.set_method(TEST_METHOD.to_string()); + req.payload = payload.clone(); + + let resp = client.request(req).unwrap(); + // Response has "role:sync-client|" prefix + original payload + let prefix = b"role:sync-client|"; + assert!(resp.payload.starts_with(prefix)); + let original = &resp.payload[prefix.len()..]; + assert_eq!(original, payload.as_slice()); + } + + server.shutdown(); + cleanup_socket_file(&sock_path); +} + +#[test] +fn test_sync_connection_data_propagated() { + let sock_path = temp_unix_socket_path(); + let path = format!("unix://{}", sock_path); + + let hook = CountingAcceptHook { + call_count: Arc::new(AtomicUsize::new(0)), + reject: false, + }; + + let mut server = Server::new() + .bind(&path) + .unwrap() + .register_service(build_sync_echo_service()) + .set_accept_hook(hook); + server.start().unwrap(); + wait_for_server_ready(); + + let connect_hook = CountingConnectHook { + call_count: Arc::new(AtomicUsize::new(0)), + }; + let fd = create_connected_fd(&sock_path); + let client = Client::with_hook(fd, connect_hook).unwrap(); + + let mut req = Request::new(); + req.set_service(TEST_SERVICE.to_string()); + req.set_method(TEST_METHOD.to_string()); + req.payload = b"data".to_vec(); + + let resp = client.request(req).unwrap(); + // Verify the connection data was accessible in the handler + let resp_str = String::from_utf8_lossy(&resp.payload); + assert!(resp_str.starts_with("role:sync-client|")); + + server.shutdown(); + cleanup_socket_file(&sock_path); +} + +// ── Utility: create a connected fd ────────────────────────────────────────── + +fn create_connected_fd(path: &str) -> RawFd { + use std::os::unix::io::IntoRawFd; + use std::os::unix::net::UnixStream; + + let stream = UnixStream::connect(path).expect("connect"); + stream.into_raw_fd() +} diff --git a/ttrpc-codegen/Makefile b/ttrpc-codegen/Makefile index bb69e69a..291e6c00 100644 --- a/ttrpc-codegen/Makefile +++ b/ttrpc-codegen/Makefile @@ -1 +1,2 @@ +FEATURES = --all-features include ../Makefile