cliquenet/net/peer.rs
1#[cfg(test)]
2mod tests;
3
4use std::{
5 cmp::min, collections::VecDeque, future::ready, io, mem, num::NonZeroUsize, sync::Arc,
6 time::Duration,
7};
8
9use bon::bon;
10use bytes::{Bytes, BytesMut};
11use snow::TransportState;
12use tokio::{
13 select,
14 sync::{OwnedSemaphorePermit, Semaphore, mpsc::UnboundedSender},
15 time::{Instant, MissedTickBehavior, interval},
16};
17use tokio_util::sync::CancellationToken;
18use tracing::{info, trace, warn};
19
20use crate::{
21 Config, Metrics, PublicKey,
22 connection::Connection,
23 delay::DelayQueue,
24 error::{Empty, NetworkError},
25 msg::{
26 Ack, FrameType, Header, MAX_NOISE_MESSAGE_SIZE, MAX_PAYLOAD_SIZE, RESERVED_TAG_SIZE,
27 Trailer,
28 },
29 net::{PeerMessage, RetryPolicy},
30 queue::Queue,
31 time::Countdown,
32};
33
34type NoiseBuf = Box<[u8; MAX_NOISE_MESSAGE_SIZE]>;
35type Result<T> = std::result::Result<T, NetworkError>;
36
37/// A peer sends and receives messages over a connection with a remote.
38///
39/// Peers are initialised with a [`Connection`], which can be replaced
40/// later, should it be necessary.
41///
42/// Messages sent are expected to be acknowledged, i.e. the remote needs to
43/// send back an ACK frame. Otherwise the message is resent after some time.
44pub struct Peer {
45 /// Network configuration.
46 conf: Arc<Config>,
47
48 /// A budget limits how many message a peer can deliver to the application.
49 budget: Budget,
50
51 /// Messages the application wants to be sent to the remote.
52 msgs: Queue<(RetryPolicy, Bytes)>,
53
54 /// Messages waiting to be retried if no ACK has been received.
55 retry: DelayQueue,
56
57 /// The channel over which to deliver inbound messages to the application.
58 tx: UnboundedSender<PeerMessage>,
59
60 /// A healthcheck countdown.
61 ///
62 /// When dropped to 0 the connection should be replaced.
63 countdown: Countdown,
64
65 /// The true max. message size.
66 ///
67 /// It accounts for the additional `Trailer` bytes.
68 max_message_size: usize,
69
70 metrics: Arc<dyn Metrics>,
71}
72
73/// A budget limits how many messages a peer can delivery to the application.
74pub struct Budget(Arc<Semaphore>);
75
76impl Budget {
77 pub fn new(amount: NonZeroUsize) -> Self {
78 Self(Arc::new(Semaphore::new(amount.get())))
79 }
80
81 fn remaining(&self) -> usize {
82 self.0.available_permits()
83 }
84}
85
86#[bon]
87impl Peer {
88 #[builder]
89 pub fn new(
90 config: Arc<Config>,
91 budget: NonZeroUsize,
92 messages: Queue<(RetryPolicy, Bytes)>,
93 retry: DelayQueue,
94 inbound: UnboundedSender<PeerMessage>,
95 metrics: Arc<dyn Metrics>,
96 ) -> Self {
97 Self {
98 max_message_size: config
99 .max_message_size
100 .get()
101 .saturating_add(Trailer::MAX_SIZE),
102 conf: config.clone(),
103 budget: Budget::new(budget),
104 tx: inbound,
105 msgs: messages,
106 retry,
107 countdown: Countdown::new(),
108 metrics,
109 }
110 }
111
112 /// Start I/O with a connected peer.
113 ///
114 /// This method continues until an error occurs, after which callers may
115 /// want to reconnect and resume peer operation with a new `Connection`.
116 pub async fn start(
117 &mut self,
118 mut conn: Connection,
119 cancel: CancellationToken,
120 ) -> Result<Empty> {
121 /// Messages are broken into frames and each frame has a header.
122 ///
123 /// The `ReadState` tracks what we have received from the remote.
124 enum ReadState<'a> {
125 Header {
126 off: usize,
127 buf: [u8; Header::SIZE],
128 },
129 Frame {
130 hdr: Header,
131 typ: FrameType,
132 off: usize,
133 buf: &'a mut Vec<u8>,
134 },
135 }
136
137 /// This state tracks what we have sent to the remote. We currently
138 /// support either data or ACK frames. The offset and length values
139 /// are relative to the write buffer (see below).
140 enum WriteState {
141 /// No write operation is in progress.
142 Idle,
143 /// An ACK frame is sent.
144 Ack { off: usize, len: usize },
145 /// A data frame is sent.
146 Data { off: usize, len: usize },
147 }
148
149 impl WriteState {
150 fn is_idle(&self) -> bool {
151 matches!(self, Self::Idle)
152 }
153
154 /// Encrypt a single data frame with Noise.
155 fn data_frame(
156 data: &[u8],
157 is_partial: bool,
158 state: &mut TransportState,
159 buf: &mut NoiseBuf,
160 ) -> Result<Self> {
161 let n = state.write_message(data, &mut buf[Header::SIZE..])?;
162 let h = if is_partial {
163 Header::data(n as u16).partial()
164 } else {
165 Header::data(n as u16)
166 };
167 buf[..Header::SIZE].copy_from_slice(&h.to_bytes());
168 Ok(Self::Data {
169 off: 0,
170 len: n + Header::SIZE,
171 })
172 }
173
174 /// Encrypt a single ACK frame with Noise.
175 fn ack_frame(a: Ack, state: &mut TransportState, buf: &mut NoiseBuf) -> Result<Self> {
176 let n = state.write_message(&a.0, &mut buf[Header::SIZE..])?;
177 let h = Header::ack(n as u16);
178 buf[..Header::SIZE].copy_from_slice(&h.to_bytes());
179 Ok(Self::Ack {
180 off: 0,
181 len: n + Header::SIZE,
182 })
183 }
184 }
185
186 // Early check if we already got cancelled which can happen with many
187 // simultaneous connects where we need to drop connections.
188 if cancel.is_cancelled() {
189 return Err(NetworkError::PeerInterrupt);
190 }
191
192 let mut now = Instant::now();
193
194 // Reset messages scheduled for another send.
195 self.retry.reset(now);
196
197 // Write buffer (Noise packages are limited to 64KiB).
198 let mut wbuf = NoiseBuf::new([0; _]);
199
200 // Ack buffer.
201 let mut abuf = [0; size_of::<Ack>() + RESERVED_TAG_SIZE];
202
203 // A frame buffer for reading before it is decrypted.
204 let mut fbuf = Vec::new();
205
206 // An incoming message that is assembled from its frames.
207 let mut ibound_msg = BytesMut::new();
208
209 // An outgoing message. The integer tracks the chunk size as we need
210 // to break the message into frames that fit into a noise package.
211 let mut obound_msg: Option<(RetryPolicy, Bytes, usize)> = None;
212
213 // Pending outbound ACK messages. This is appended when we received a
214 // message and picked up and interleaved when sending frames or else
215 // at the start of the loop (see below).
216 let mut obound_acks: VecDeque<Ack> = VecDeque::new();
217
218 // Track write and read states:
219 let mut wstate = WriteState::Idle;
220 let mut rstate = ReadState::Header {
221 off: 0,
222 buf: [0; _],
223 };
224
225 // Each read needs a permit taken from our budget in order to deliver
226 // to the application.
227 let mut read_permit: Option<OwnedSemaphorePermit> = None;
228
229 // Measure time.
230 let mut clock = interval(Duration::from_secs(1));
231 clock.set_missed_tick_behavior(MissedTickBehavior::Skip);
232
233 // Ensure any previously fired countdown is reset.
234 self.countdown.stop();
235
236 loop {
237 trace!(
238 name = %self.conf.name,
239 peer = %conn.key,
240 addr = %conn.addr,
241 writing = %!wstate.is_idle(),
242 acks = %obound_acks.len(),
243 retries = %self.retry.len(),
244 can_read = %read_permit.is_some(),
245 "entering event loop"
246 );
247
248 select! {
249 // Wait for the next message if no write is in progress.
250 //
251 // We limit writing if we expect too many ACKs that the remote
252 // has not sent yet. This is to prevent an attack were a
253 // malicious peer never sends ACKs, which would cause our
254 // delay queue to grow unbounded.
255 m = self.msgs.next(), if wstate.is_idle() && self.retry.len() < self.conf.peer_budget.get() => {
256 trace!(name = %self.conf.name, peer = %conn.key, "next outbound message");
257 let (slot, id, (policy, bytes)) = m;
258 if policy.is_retry() {
259 self.retry.add(slot, id, bytes.clone(), policy, Instant::now());
260 }
261 self.msgs.remove(slot, id);
262 let chunk = min(bytes.len(), MAX_PAYLOAD_SIZE);
263 wstate = WriteState::data_frame(
264 &bytes[..chunk],
265 chunk < bytes.len(),
266 &mut conn.state,
267 &mut wbuf
268 )?;
269 obound_msg = Some((policy, bytes, chunk))
270 }
271
272 // Pick up an ACK and send it if possible.
273 () = ready(()), if wstate.is_idle() && !obound_acks.is_empty() => {
274 trace!(name = %self.conf.name, peer = %conn.key, "next outbound ack");
275 let ack = obound_acks.pop_front().expect("obound_acks is not empty");
276 wstate = WriteState::ack_frame(ack, &mut conn.state, &mut wbuf)?
277 }
278
279 // If requested, interrupt all I/O processing.
280 //
281 // Once a peer has been interrupted, its connection needs to be
282 // replaced before calling start again.
283 _ = cancel.cancelled() => {
284 info!(name = %self.conf.name, peer = %conn.key, "peer interrupt");
285 return Err(NetworkError::PeerInterrupt)
286 }
287
288 // Update time and metrics.
289 t = clock.tick() => {
290 now = t;
291 self.update_metrics(&conn.key);
292 }
293
294 // If messages should be re-sent and we can do so, send it:
295 () = ready(()), if wstate.is_idle() && self.retry.is_due(now) => {
296 trace!(name = %self.conf.name, peer = %conn.key, "resending message");
297 let Some((bytes, policy)) = self.retry.due(now) else {
298 continue
299 };
300 let chunk = min(bytes.len(), MAX_PAYLOAD_SIZE);
301 wstate = WriteState::data_frame(
302 &bytes[..chunk],
303 chunk < bytes.len(),
304 &mut conn.state,
305 &mut wbuf
306 )?;
307 obound_msg = Some((policy, bytes, chunk))
308 }
309
310 // Continue an ongoing write operation.
311 r = conn.stream.writable(), if !wstate.is_idle() => {
312 trace!(name = %self.conf.name, peer = %conn.key, "continue writing");
313 if let Err(e) = r {
314 return Err(e.into())
315 }
316 match &mut wstate {
317 WriteState::Ack { off, len } => {
318 match conn.stream.try_write(&wbuf[*off..*len]) {
319 Ok(n) => {
320 *off += n;
321 if *off < *len {
322 continue
323 }
324 if let Some((policy, bytes, chunk)) = &mut obound_msg {
325 if *chunk < bytes.len() {
326 let end = min(*chunk + MAX_PAYLOAD_SIZE, bytes.len());
327 wstate = WriteState::data_frame(
328 &bytes[*chunk..end],
329 end < bytes.len(),
330 &mut conn.state,
331 &mut wbuf
332 )?;
333 *chunk = end;
334 } else {
335 if policy.is_retry() {
336 self.countdown.start(self.conf.receive_timeout)
337 }
338 obound_msg = None;
339 wstate = WriteState::Idle;
340 }
341 } else {
342 wstate = WriteState::Idle;
343 }
344 }
345 Err(e) => if e.kind() != io::ErrorKind::WouldBlock {
346 return Err(e.into())
347 }
348 }
349 }
350 WriteState::Data { off, len } => {
351 match conn.stream.try_write(&wbuf[*off..*len]) {
352 Ok(n) => {
353 *off += n;
354 if *off < *len {
355 continue
356 }
357 if let Some(ack) = obound_acks.pop_front() {
358 wstate = WriteState::ack_frame(ack, &mut conn.state, &mut wbuf)?
359 } else if let Some((policy, bytes, chunk)) = &mut obound_msg {
360 if *chunk < bytes.len() {
361 let end = min(*chunk + MAX_PAYLOAD_SIZE, bytes.len());
362 wstate = WriteState::data_frame(
363 &bytes[*chunk..end],
364 end < bytes.len(),
365 &mut conn.state,
366 &mut wbuf
367 )?;
368 *chunk = end;
369 } else {
370 if policy.is_retry() {
371 self.countdown.start(self.conf.receive_timeout)
372 }
373 obound_msg = None;
374 wstate = WriteState::Idle;
375 }
376 } else {
377 wstate = WriteState::Idle;
378 }
379 }
380 Err(e) => if e.kind() != io::ErrorKind::WouldBlock {
381 return Err(e.into())
382 }
383 }
384 },
385 WriteState::Idle => { /* unreachable!() */ }
386 }
387 }
388
389 // Wait for the healthcheck countdown to finish.
390 //
391 // The countdown is started after writing a message and reset
392 // when we received a frame.
393 () = &mut self.countdown => {
394 trace!(name = %self.conf.name, peer = %conn.key, "read timeout");
395 return Err(NetworkError::Timeout)
396 }
397
398 // Await the next read permit.
399 //
400 // If our budget is used up, we need to wait for capacity to become
401 // available before we can continue to read from the socket (and
402 // eventually deliver the message to the application).
403 p = self.budget.0.clone().acquire_owned(), if read_permit.is_none() => {
404 trace!(name = %self.conf.name, peer = %conn.key, "next read permit");
405 read_permit = Some(p.map_err(|_| NetworkError::BudgetClosed)?);
406 }
407
408 // Continue reading from the socket if possible.
409 //
410 // NB that we require the ACKs that we have appended before to
411 // picked up. This should be very fast as writing interleaves
412 // ACKs in between frames. We do this to exercise backpressure
413 // because if the remote does not or can not read what we write
414 // but keeps sending us data we would accumulate ACKs without
415 // bound.
416 r = conn.stream.readable(), if read_permit.is_some() => {
417 trace!(name = %self.conf.name, peer = %conn.key, "continue reading");
418 if let Err(e) = r {
419 return Err(e.into())
420 }
421 if obound_acks.len() > self.conf.peer_budget.get() {
422 return Err(NetworkError::TooManyPendingAcks(conn.key))
423 }
424 match &mut rstate {
425 ReadState::Header { off, buf } => {
426 match conn.stream.try_read(&mut buf[*off..]) {
427 Ok(0) => {
428 let e = io::ErrorKind::UnexpectedEof.into();
429 return Err(NetworkError::Io(e))
430 }
431 Ok(n) => {
432 self.countdown.stop();
433 *off += n;
434 if *off < buf.len() {
435 continue
436 }
437 let hdr = Header::unvalidated(u32::from_be_bytes(*buf));
438 let typ = match hdr.frame_type() {
439 Ok(FrameType::Data) => FrameType::Data,
440 Ok(FrameType::Ack) => {
441 if hdr.is_partial() {
442 warn!(
443 name = %self.conf.name,
444 node = %self.conf.keypair.public_key(),
445 peer = %conn.key,
446 addr = %conn.addr,
447 "ACK header marked as partial"
448 );
449 return Err(NetworkError::InvalidAck)
450 }
451 if hdr.len() as usize > abuf.len() {
452 warn!(
453 name = %self.conf.name,
454 node = %self.conf.keypair.public_key(),
455 peer = %conn.key,
456 addr = %conn.addr,
457 len = %hdr.len(),
458 "ACK header length too large"
459 );
460 return Err(NetworkError::InvalidAck)
461 }
462 FrameType::Ack
463 }
464 Err(typ) => {
465 return Err(NetworkError::UnknownFrameType(typ))
466 }
467 };
468 fbuf.resize(hdr.len().into(), 0);
469 rstate = ReadState::Frame { hdr, typ, off: 0, buf: &mut fbuf }
470 }
471 Err(e) => if e.kind() != io::ErrorKind::WouldBlock {
472 return Err(e.into())
473 }
474 }
475 }
476 ReadState::Frame { hdr, typ, off, buf } => {
477 match conn.stream.try_read(&mut buf[*off..]) {
478 Ok(0) => {
479 let e = io::ErrorKind::UnexpectedEof.into();
480 return Err(NetworkError::Io(e))
481 }
482 Ok(n) => {
483 self.countdown.stop();
484 *off += n;
485 if *off < buf.len() {
486 continue
487 }
488 match typ {
489 FrameType::Data => {
490 let n = buf.len();
491 let i = ibound_msg.len();
492 ibound_msg.resize(i + n, 0);
493
494 let n = conn.state.read_message(buf, &mut ibound_msg[i..])?;
495 ibound_msg.truncate(i + n);
496
497 if ibound_msg.len() > self.max_message_size {
498 return Err(NetworkError::MessageTooLarge)
499 }
500
501 if !hdr.is_partial() { // message complete
502 let mut msg = mem::take(&mut ibound_msg).freeze();
503 let Some(t) = Trailer::from_bytes(&mut msg) else {
504 warn!(
505 name = %self.conf.name,
506 node = %self.conf.keypair.public_key(),
507 peer = %conn.key,
508 addr = %conn.addr,
509 "invalid trailer"
510 );
511 return Err(NetworkError::InvalidTrailer);
512 };
513 match t {
514 Trailer::Std { slot, id } =>
515 obound_acks.push_back(Ack::from((slot, id))),
516 Trailer::NoAck { slot: _ } => (),
517 Trailer::Unknown => ()
518 }
519 let p = read_permit.take();
520 debug_assert!(p.is_some());
521 if self.tx.send((conn.key, msg, p)).is_err() {
522 return Err(NetworkError::ChannelClosed)
523 }
524 trace!(
525 name = %self.conf.name,
526 node = %self.conf.keypair.public_key(),
527 peer = %conn.key,
528 addr = %conn.addr,
529 "message delivered"
530 );
531 }
532 rstate = ReadState::Header { off: 0, buf: [0; _] };
533 }
534 FrameType::Ack => {
535 let n = conn.state.read_message(buf, &mut abuf)?;
536 let Ok(a) = Ack::try_from(&abuf[..n]) else {
537 return Err(NetworkError::InvalidAck)
538 };
539 let (s, i) = a.into();
540 self.retry.remove(s, i);
541 rstate = ReadState::Header { off: 0, buf: [0; _] };
542 }
543 }
544 }
545 Err(e) => if e.kind() != io::ErrorKind::WouldBlock {
546 return Err(e.into())
547 }
548 }
549 }
550 }
551 }
552 }
553 }
554 }
555
556 fn update_metrics(&self, key: &PublicKey) {
557 self.metrics.set(key, "outbound_messages", self.msgs.len());
558 self.metrics.set(key, "retrying_messages", self.retry.len());
559 self.metrics
560 .set(key, "remaining_budget", self.budget.remaining());
561 }
562}