dynomite/net/dnode_client.rs
1//! Inbound peer-connection driver for the dnode peer plane.
2//!
3//! The local node is the receiver. The driver:
4//!
5//! 1. Reads bytes off the transport into a contiguous buffer.
6//! 2. Drives the dnode header parser ([`crate::proto::dnode::DnodeParser`])
7//! over the buffer until a full `Dmsg` header has been observed.
8//! 3. If the header marks the payload as encrypted, decrypts it
9//! using the per-connection AES key bound during the handshake
10//! via [`crate::crypto::Crypto`]. When the header indicates a
11//! plaintext payload (the peer-plane was negotiated unsecured),
12//! the bytes pass through unchanged.
13//! 4. Drives the datastore parser over the (decrypted) payload to
14//! reconstruct a [`Msg`].
15//! 5. Hands the parsed [`Msg`] to the supplied
16//! [`ClientHandler`]'s dispatcher and routes the dispatcher's
17//! response back through the per-connection responder channel.
18
19use tokio::io::{AsyncReadExt, AsyncWriteExt};
20use tokio::sync::mpsc;
21
22use crate::msg::Msg;
23use crate::msg::MsgParseResult;
24use crate::msg::MsgType;
25use crate::net::client::ClientHandler;
26use crate::net::conn::Conn;
27use crate::net::dispatcher::OutboundEnvelope;
28use crate::net::NetError;
29use crate::proto::dnode::{DmsgType, DnodeParser, ParseStep};
30
31/// Type alias for the dnode client handler bundle.
32pub type DnodeClientHandler = ClientHandler;
33
34/// Drive a DNODE_PEER_CLIENT FSM until the peer closes.
35///
36/// `rx` receives responses produced by the cluster dispatcher; the
37/// driver writes the response bytes back through the same
38/// transport.
39///
40/// # Errors
41/// Surfaces transport- and DNODE-level errors.
42pub async fn dnode_client_loop(
43 mut conn: Conn,
44 handler: ClientHandler,
45 mut rx: mpsc::Receiver<OutboundEnvelope>,
46) -> Result<(), NetError> {
47 let mut read_buf = vec![0u8; 4096];
48 let mut accumulated = Vec::<u8>::new();
49 let mut parser = DnodeParser::new();
50
51 loop {
52 if conn.is_eof() && conn.imsg_q().is_empty() && conn.omsg_q().is_empty() {
53 conn.set_done();
54 return Ok(());
55 }
56
57 tokio::select! {
58 res = async {
59 if let Some(t) = conn.transport_mut() {
60 t.read(&mut read_buf).await
61 } else {
62 Ok(0)
63 }
64 } => {
65 let n = res?;
66 if n == 0 {
67 conn.set_eof();
68 continue;
69 }
70 conn.record_recv(n);
71 accumulated.extend_from_slice(&read_buf[..n]);
72 drive_dnode_parser(&mut conn, &handler, &mut accumulated, &mut parser).await?;
73 }
74 Some(env) = rx.recv() => {
75 // Forward dispatcher-produced responses back to
76 // the peer over this same transport. The peer-
77 // originator's `DnodeServerConn` parses incoming
78 // bytes with `DnodeParser`, so the response must
79 // be dnode-framed (header + payload). Without
80 // this header the originator's parser hangs in
81 // `NeedMore`, the dispatcher's `responder` mpsc
82 // never gets the reply, and the originating
83 // client times out.
84 let bytes: Vec<u8> = env
85 .rsp
86 .mbufs()
87 .iter()
88 .flat_map(|b| b.readable().to_vec())
89 .collect();
90 if !bytes.is_empty() {
91 let mut header_buf = conn.mbuf_pool().get();
92 crate::proto::dnode::dmsg_write(
93 &mut header_buf,
94 env.req_id,
95 crate::proto::dnode::DmsgType::Res,
96 0,
97 true,
98 None,
99 u32::try_from(bytes.len()).unwrap_or(u32::MAX),
100 )
101 .map_err(|e| NetError::Dnode(format!("{e:?}")))?;
102 let header_len = header_buf.readable().len();
103 if let Some(t) = conn.transport_mut() {
104 t.write_all(header_buf.readable()).await?;
105 t.write_all(&bytes).await?;
106 conn.record_send(header_len + bytes.len());
107 }
108 }
109 conn.outstanding_mut().remove(&env.req_id);
110 if let Some(front) = conn.omsg_q_mut().front() {
111 if front.id() == env.req_id {
112 let _ = conn.omsg_q_mut().pop_front();
113 }
114 }
115 }
116 }
117 }
118}
119
120/// Frame `reply` bytes back to a peer as a dnode `Res` carrying
121/// `reply_id`, over the connection's transport. Used by the read-
122/// coordination path: a replica answers a `DtFetch` query with its
123/// local CRDT state so the coordinator can merge the replica set.
124async fn write_replica_reply(
125 conn: &mut Conn,
126 reply_id: crate::core::types::MsgId,
127 reply: &[u8],
128) -> Result<(), NetError> {
129 let mut header_buf = conn.mbuf_pool().get();
130 crate::proto::dnode::dmsg_write(
131 &mut header_buf,
132 reply_id,
133 crate::proto::dnode::DmsgType::Res,
134 0,
135 true,
136 None,
137 u32::try_from(reply.len()).unwrap_or(u32::MAX),
138 )
139 .map_err(|e| NetError::Dnode(format!("{e:?}")))?;
140 let header_len = header_buf.readable().len();
141 if let Some(t) = conn.transport_mut() {
142 t.write_all(header_buf.readable()).await?;
143 t.write_all(reply).await?;
144 conn.record_send(header_len + reply.len());
145 }
146 Ok(())
147}
148
149async fn drive_dnode_parser(
150 conn: &mut Conn,
151 handler: &ClientHandler,
152 accumulated: &mut Vec<u8>,
153 parser: &mut DnodeParser,
154) -> Result<(), NetError> {
155 loop {
156 if accumulated.is_empty() {
157 return Ok(());
158 }
159 let step = parser.step(accumulated.as_slice());
160 match step {
161 ParseStep::NeedMore { .. } => return Ok(()),
162 ParseStep::Error { consumed } => {
163 return Err(NetError::Dnode(format!(
164 "dnode header parse error after {consumed} bytes"
165 )));
166 }
167 ParseStep::HeaderDone { consumed } => {
168 let header_end = consumed;
169 let dmsg = parser.take_dmsg();
170 let plen = dmsg.plen as usize;
171 let total = header_end + plen;
172 if accumulated.len() < total {
173 // Wait for more bytes for the payload; rewind
174 // by stashing what we have. The parser was
175 // moved to PostDone but we need it to retry
176 // header parsing on the next chunk.
177 parser.reset();
178 return Ok(());
179 }
180 let payload = accumulated[header_end..total].to_vec();
181 accumulated.drain(0..total);
182 parser.reset();
183
184 // Gossip-class frames are control plane: feed the
185 // sender's identity into the gossip handler's
186 // failure detector and skip the datastore parse
187 // path. Without this fork the datastore parser
188 // sees an opaque ASCII pname (e.g. `127.0.0.1:8101`)
189 // and rejects it with a parse error, causing the
190 // dnode_client_loop to tear the connection down.
191 if is_gossip_ty(dmsg.ty) {
192 handle_gossip_frame(handler, dmsg.ty, &payload);
193 continue;
194 }
195
196 // Dyniak cross-node object-replica frames are
197 // applied to the LOCAL datastore via the attached
198 // replica sink and are never re-forwarded, so a
199 // replica write fans out exactly once. Without a
200 // sink wired (a non-dyniak pool), the frame is
201 // dropped: no other node treats RiakReplica as a
202 // data-plane request.
203 if matches!(dmsg.ty, DmsgType::RiakReplica) {
204 if let Some(sink) = handler.replica_sink() {
205 // A replica op may be a fire-and-forget write
206 // (apply, no reply) or a read-coordination
207 // query (apply_query returns Some(reply)). Try
208 // the query path first; if it yields reply
209 // bytes, frame them back to the requester as a
210 // `Res` carrying the same dmsg id so the
211 // coordinator's outstanding request resolves.
212 let reply_id = dmsg.id;
213 if let Some(reply) = sink.apply_query(&payload).await {
214 write_replica_reply(conn, reply_id, &reply).await?;
215 } else {
216 sink.apply(&payload).await;
217 }
218 }
219 continue;
220 }
221
222 // Decrypt if the dnode header indicates the payload
223 // is encrypted and we have an AES key.
224 let decoded = if dmsg.is_encrypted() {
225 let Some(key) = conn.aes_key() else {
226 // No key has been negotiated yet; the
227 // peer-plane handshake should have run
228 // first. Surface a single opaque parse
229 // error and let the driver close the
230 // connection.
231 return Err(NetError::Dnode(
232 "dnode payload marked encrypted but no aes key bound".into(),
233 ));
234 };
235 decrypt_dnode_payload(key, &payload)?
236 } else {
237 payload
238 };
239
240 // Feed the decoded payload through the datastore
241 // parser to reconstruct a Msg.
242 let mut msg = Msg::new(dmsg.id, MsgType::Unknown, true);
243 let dmsg_ty = dmsg.ty;
244 msg.set_dmsg(dmsg);
245 let parse_result = match handler.data_store() {
246 crate::conf::DataStore::Valkey | crate::conf::DataStore::Dyniak => {
247 crate::proto::redis::redis_parse_req(&mut msg, &decoded)
248 }
249 crate::conf::DataStore::Memcache => {
250 crate::proto::memcache::memcache_parse_req(&mut msg, &decoded)
251 }
252 };
253 if matches!(dmsg_ty, DmsgType::ReqForward) {
254 // A `ReqForward` is the wire signal that this
255 // request was already routed by an upstream
256 // dispatcher (e.g. a quorum coalescer issuing
257 // a read-repair write back to a divergent
258 // replica). The receiver must NOT re-fan it
259 // out: tag the parsed request as
260 // `LocalNodeOnly` so the cluster dispatcher
261 // hands it straight to its local datastore.
262 msg.set_routing(crate::msg::MsgRouting::LocalNodeOnly);
263 }
264 match parse_result {
265 MsgParseResult::Ok | MsgParseResult::Noop => {
266 let pool = conn.mbuf_pool().clone();
267 let mut buf = pool.get();
268 buf.recv(&decoded);
269 msg.mbufs_mut().push_back(buf);
270 msg.recompute_mlen();
271 conn.outstanding_mut().insert(msg.id(), msg.id());
272 conn.enqueue_out(Msg::new(msg.id(), msg.ty(), true))?;
273 // Hand the parsed peer request to the
274 // configured dispatcher. The dispatcher
275 // either takes ownership and replies
276 // asynchronously through `responder`, or it
277 // returns an inline / error response that
278 // we forward immediately, or it asks the
279 // FSM to drop the request.
280 let outcome = handler
281 .dispatcher()
282 .dispatch(msg, handler.response_tx().clone());
283 match outcome {
284 crate::net::dispatcher::DispatchOutcome::Pending
285 | crate::net::dispatcher::DispatchOutcome::Drop => {}
286 crate::net::dispatcher::DispatchOutcome::Inline(rsp)
287 | crate::net::dispatcher::DispatchOutcome::Error(rsp) => {
288 let env = OutboundEnvelope {
289 req_id: rsp.id(),
290 rsp,
291 span: tracing::Span::current(),
292 source_peer_idx: None,
293 };
294 let _ = handler.response_tx().send(env).await;
295 }
296 }
297 }
298 MsgParseResult::Again => return Ok(()),
299 other => {
300 return Err(NetError::Parse(format!("dnode payload parse: {other:?}")));
301 }
302 }
303 }
304 }
305 }
306}
307
308/// Decrypt a dnode peer-plane payload using the per-connection AES
309/// key.
310///
311/// AES-128-CBC with PKCS#7 padding, IV from the trailing 16 bytes
312/// of the 32-byte key buffer. Returns a single opaque
313/// [`NetError::Dnode`] on failure regardless of whether the
314/// underlying error was bad padding, a length mismatch, or a
315/// key/iv mismatch, so peers cannot distinguish the cases, which
316/// avoids a padding-oracle surface.
317fn decrypt_dnode_payload(
318 key: &[u8; crate::crypto::AES_KEYLEN],
319 payload: &[u8],
320) -> Result<Vec<u8>, NetError> {
321 crate::crypto::Crypto::aes_decrypt(payload, key)
322 .map_err(|_| NetError::Dnode("dnode payload decrypt failed".into()))
323}
324
325/// True for any dnode message type that belongs to the gossip
326/// control plane. The data plane (`Req`, `ReqForward`, `Res`,
327/// `CryptoHandshake`, `Unknown`, `Debug`, `ParseError`) returns
328/// `false`.
329fn is_gossip_ty(ty: DmsgType) -> bool {
330 matches!(
331 ty,
332 DmsgType::GossipSyn
333 | DmsgType::GossipSynReply
334 | DmsgType::GossipAck
335 | DmsgType::GossipDigestSyn
336 | DmsgType::GossipDigestAck
337 | DmsgType::GossipDigestAck2
338 | DmsgType::GossipShutdown
339 )
340}
341
342/// Process a gossip control-plane frame. The payload is the
343/// sender peer's `host:port` (ASCII). Heartbeat-class frames
344/// feed the failure detector; `GossipShutdown` immediately
345/// transitions the sender to [`crate::cluster::peer::PeerState::Down`].
346///
347/// Frames received before the run loop has attached a gossip
348/// handler are silently dropped; this matches the reference
349/// engine's behaviour of ignoring stray gossip while the
350/// failure detector is still being constructed.
351fn handle_gossip_frame(handler: &ClientHandler, ty: DmsgType, payload: &[u8]) {
352 let Some(gossip) = handler.gossip() else {
353 return;
354 };
355 let Ok(pname) = std::str::from_utf8(payload) else {
356 return;
357 };
358 let pname = pname.trim();
359 if pname.is_empty() {
360 return;
361 }
362 let now = std::time::Instant::now();
363 match ty {
364 DmsgType::GossipShutdown => {
365 gossip.mark_down_pname(pname);
366 }
367 _ => {
368 gossip.record_heartbeat_pname(pname, now);
369 }
370 }
371}
372
373#[cfg(test)]
374mod tests {
375 use super::*;
376 use crate::io::reactor::{ConnRole, TcpTransport};
377 use tokio::net::{TcpListener, TcpStream};
378
379 #[tokio::test]
380 async fn build_and_drop() {
381 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
382 let addr = listener.local_addr().unwrap();
383 let _accept = tokio::spawn(async move {
384 let (s, _) = listener.accept().await.unwrap();
385 drop(s);
386 });
387 let s = TcpStream::connect(addr).await.unwrap();
388 let _conn = Conn::new(
389 Box::new(TcpTransport::new(s, ConnRole::DnodePeerClient)),
390 ConnRole::DnodePeerClient,
391 );
392 }
393}