iroh_relay/server/
streams.rs

1//! Streams used in the server-side implementation of iroh relays.
2
3use std::{
4    pin::Pin,
5    sync::{Arc, atomic::AtomicBool},
6    task::{Context, Poll},
7};
8
9use n0_error::{ensure, stack_error};
10use n0_future::{FutureExt, Sink, Stream, ready, time};
11use tokio::{
12    io::{AsyncRead, AsyncWrite},
13    sync::watch,
14};
15use tracing::{instrument, warn};
16
17use super::{ClientRateLimit, Metrics};
18use crate::{
19    ExportKeyingMaterial, KeyCache, MAX_PACKET_SIZE,
20    protos::{
21        relay::{ClientToRelayMsg, Error as ProtoError, RelayToClientMsg},
22        streams::{StreamError, WsBytesFramed},
23    },
24};
25
26/// The relay's connection to a client.
27///
28/// This implements
29/// - a [`Stream`] of [`ClientToRelayMsg`]s that are received from the client,
30/// - a [`Sink`] of [`RelayToClientMsg`]s that can be sent to the client.
31///
32/// Generic over the inner stream type to support different WebSocket implementations.
33#[derive(Debug)]
34pub struct RelayedStream<S> {
35    pub(crate) inner: S,
36    pub(crate) key_cache: KeyCache,
37}
38
39impl<S> RelayedStream<S> {
40    /// Creates a new RelayedStream from an inner stream and key cache.
41    ///
42    /// This is the primary constructor for external integrations using custom
43    /// WebSocket implementations.
44    pub fn new(inner: S, key_cache: KeyCache) -> Self {
45        Self { inner, key_cache }
46    }
47}
48
49/// Type alias for the standard server-side relay stream
50#[allow(dead_code)]
51pub(super) type ServerRelayedStream = RelayedStream<WsBytesFramed<RateLimited<MaybeTlsStream>>>;
52
53#[cfg(test)]
54impl ServerRelayedStream {
55    pub(crate) fn test(stream: tokio::io::DuplexStream) -> Self {
56        let stream = MaybeTlsStream::Test(stream);
57        let stream = RateLimited::unlimited(stream, Arc::new(Metrics::default()));
58        Self {
59            inner: WsBytesFramed {
60                io: tokio_websockets::ServerBuilder::new()
61                    .limits(Self::limits())
62                    .serve(stream),
63            },
64            key_cache: KeyCache::test(),
65        }
66    }
67
68    pub(crate) fn test_limited(
69        stream: tokio::io::DuplexStream,
70        max_burst_bytes: u32,
71        bytes_per_second: u32,
72    ) -> Result<(Self, watch::Receiver<u64>), InvalidBucketConfig> {
73        let stream = MaybeTlsStream::Test(stream);
74        let stream = RateLimited::new(
75            stream,
76            max_burst_bytes,
77            bytes_per_second,
78            Arc::new(Metrics::default()),
79        )?;
80        let limited_watcher = stream.limited_watcher();
81        let stream = Self {
82            inner: WsBytesFramed {
83                io: tokio_websockets::ServerBuilder::new()
84                    .limits(Self::limits())
85                    .serve(stream),
86            },
87            key_cache: KeyCache::test(),
88        };
89        Ok((stream, limited_watcher))
90    }
91
92    fn limits() -> tokio_websockets::Limits {
93        tokio_websockets::Limits::default()
94            .max_payload_len(Some(crate::protos::relay::MAX_FRAME_SIZE))
95    }
96}
97
98/// Relay send errors
99#[stack_error(derive, add_meta)]
100#[non_exhaustive]
101pub enum SendError {
102    /// Error from the underlying WebSocket stream
103    #[error(transparent)]
104    StreamError {
105        #[error(from, std_err)]
106        /// The underlying stream error
107        source: StreamError,
108    },
109    /// Packet size exceeds the maximum allowed size
110    #[error("Packet exceeds max packet size")]
111    ExceedsMaxPacketSize {
112        /// The size of the packet that was too large
113        size: usize,
114    },
115    /// Attempted to send an empty packet
116    #[error("Attempted to send empty packet")]
117    EmptyPacket {},
118}
119
120impl<S> Sink<RelayToClientMsg> for RelayedStream<S>
121where
122    S: Sink<bytes::Bytes, Error = StreamError> + Unpin,
123{
124    type Error = SendError;
125
126    fn poll_ready(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
127        Pin::new(&mut self.inner).poll_ready(cx).map_err(Into::into)
128    }
129
130    fn start_send(mut self: Pin<&mut Self>, item: RelayToClientMsg) -> Result<(), Self::Error> {
131        let size = item.encoded_len();
132        ensure!(
133            size <= MAX_PACKET_SIZE,
134            SendError::ExceedsMaxPacketSize { size }
135        );
136        if let RelayToClientMsg::Datagrams { datagrams, .. } = &item {
137            ensure!(!datagrams.contents.is_empty(), SendError::EmptyPacket);
138        }
139
140        Pin::new(&mut self.inner)
141            .start_send(item.to_bytes().freeze())
142            .map_err(Into::into)
143    }
144
145    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
146        Pin::new(&mut self.inner).poll_flush(cx).map_err(Into::into)
147    }
148
149    fn poll_close(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
150        Pin::new(&mut self.inner).poll_close(cx).map_err(Into::into)
151    }
152}
153
154/// Relay receive errors
155#[stack_error(derive, add_meta, from_sources)]
156#[non_exhaustive]
157pub enum RecvError {
158    /// Error decoding the relay protocol message
159    #[error(transparent)]
160    Proto {
161        /// The protocol decoding error
162        source: ProtoError,
163    },
164    /// Error from the underlying WebSocket stream
165    #[error(transparent)]
166    StreamError {
167        #[error(std_err)]
168        /// The underlying stream error
169        source: StreamError,
170    },
171}
172
173impl<S> Stream for RelayedStream<S>
174where
175    S: Stream<Item = Result<bytes::Bytes, StreamError>> + Unpin,
176{
177    type Item = Result<ClientToRelayMsg, RecvError>;
178
179    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
180        Poll::Ready(match ready!(Pin::new(&mut self.inner).poll_next(cx)) {
181            Some(Ok(msg)) => {
182                Some(ClientToRelayMsg::from_bytes(msg, &self.key_cache).map_err(Into::into))
183            }
184            Some(Err(e)) => Some(Err(e.into())),
185            None => None,
186        })
187    }
188}
189
190/// The main underlying IO stream type used for the relay server.
191///
192/// Allows choosing whether or not the underlying [`tokio::net::TcpStream`] is served over TLS.
193#[derive(Debug)]
194#[allow(clippy::large_enum_variant)]
195#[non_exhaustive]
196pub enum MaybeTlsStream {
197    /// A plain non-Tls [`tokio::net::TcpStream`]
198    Plain(tokio::net::TcpStream),
199    /// A Tls wrapped [`tokio::net::TcpStream`]
200    Tls(tokio_rustls::server::TlsStream<tokio::net::TcpStream>),
201    /// An in-memory bidirectional pipe.
202    #[cfg(test)]
203    Test(tokio::io::DuplexStream),
204}
205
206impl MaybeTlsStream {
207    /// Tries to disable the nagle algorithm on the TCP stream.
208    ///
209    /// This sets the NO_DELAY option on the TCP stream, which turns off the
210    /// nagle algorithm for coalecing writes together.
211    ///
212    /// If this fails, this will print a warning the first time it fails.
213    pub(super) fn disable_nagle(&self) {
214        let stream = match self {
215            #[cfg(test)]
216            Self::Test(_) => return,
217            Self::Plain(stream) => stream,
218            Self::Tls(tls_stream) => tls_stream.get_ref().0,
219        };
220
221        if stream.set_nodelay(true).is_err() {
222            use std::sync::atomic::Ordering::Relaxed;
223
224            static FAILED_NO_DELAY: AtomicBool = AtomicBool::new(false);
225            if !FAILED_NO_DELAY.swap(true, Relaxed) {
226                warn!(
227                    "Failed to set TCP socket to NO_DELAY (turning off Nagle failed). This will impair relay performance."
228                );
229            }
230        }
231    }
232}
233
234impl ExportKeyingMaterial for MaybeTlsStream {
235    fn export_keying_material<T: AsMut<[u8]>>(
236        &self,
237        output: T,
238        label: &[u8],
239        context: Option<&[u8]>,
240    ) -> Option<T> {
241        let Self::Tls(tls) = self else {
242            return None;
243        };
244
245        tls.get_ref()
246            .1
247            .export_keying_material(output, label, context)
248            .ok()
249    }
250}
251
252impl AsyncRead for MaybeTlsStream {
253    fn poll_read(
254        mut self: Pin<&mut Self>,
255        cx: &mut Context<'_>,
256        buf: &mut tokio::io::ReadBuf<'_>,
257    ) -> Poll<std::io::Result<()>> {
258        match &mut *self {
259            MaybeTlsStream::Plain(s) => Pin::new(s).poll_read(cx, buf),
260            MaybeTlsStream::Tls(s) => Pin::new(s).poll_read(cx, buf),
261            #[cfg(test)]
262            MaybeTlsStream::Test(s) => Pin::new(s).poll_read(cx, buf),
263        }
264    }
265}
266
267impl AsyncWrite for MaybeTlsStream {
268    fn poll_flush(
269        mut self: Pin<&mut Self>,
270        cx: &mut Context<'_>,
271    ) -> Poll<std::result::Result<(), std::io::Error>> {
272        match &mut *self {
273            MaybeTlsStream::Plain(s) => Pin::new(s).poll_flush(cx),
274            MaybeTlsStream::Tls(s) => Pin::new(s).poll_flush(cx),
275            #[cfg(test)]
276            MaybeTlsStream::Test(s) => Pin::new(s).poll_flush(cx),
277        }
278    }
279
280    fn poll_shutdown(
281        mut self: Pin<&mut Self>,
282        cx: &mut Context<'_>,
283    ) -> Poll<std::result::Result<(), std::io::Error>> {
284        match &mut *self {
285            MaybeTlsStream::Plain(s) => Pin::new(s).poll_shutdown(cx),
286            MaybeTlsStream::Tls(s) => Pin::new(s).poll_shutdown(cx),
287            #[cfg(test)]
288            MaybeTlsStream::Test(s) => Pin::new(s).poll_shutdown(cx),
289        }
290    }
291
292    fn poll_write(
293        mut self: Pin<&mut Self>,
294        cx: &mut Context<'_>,
295        buf: &[u8],
296    ) -> Poll<std::result::Result<usize, std::io::Error>> {
297        match &mut *self {
298            MaybeTlsStream::Plain(s) => Pin::new(s).poll_write(cx, buf),
299            MaybeTlsStream::Tls(s) => Pin::new(s).poll_write(cx, buf),
300            #[cfg(test)]
301            MaybeTlsStream::Test(s) => Pin::new(s).poll_write(cx, buf),
302        }
303    }
304
305    fn poll_write_vectored(
306        mut self: Pin<&mut Self>,
307        cx: &mut Context<'_>,
308        bufs: &[std::io::IoSlice<'_>],
309    ) -> Poll<std::result::Result<usize, std::io::Error>> {
310        match &mut *self {
311            MaybeTlsStream::Plain(s) => Pin::new(s).poll_write_vectored(cx, bufs),
312            MaybeTlsStream::Tls(s) => Pin::new(s).poll_write_vectored(cx, bufs),
313            #[cfg(test)]
314            MaybeTlsStream::Test(s) => Pin::new(s).poll_write_vectored(cx, bufs),
315        }
316    }
317
318    fn is_write_vectored(&self) -> bool {
319        match self {
320            MaybeTlsStream::Plain(s) => s.is_write_vectored(),
321            MaybeTlsStream::Tls(s) => s.is_write_vectored(),
322            #[cfg(test)]
323            MaybeTlsStream::Test(s) => s.is_write_vectored(),
324        }
325    }
326}
327
328/// Rate limiter for reading from a [`RelayedStream`].
329///
330/// The writes to the sink are not rate limited.
331///
332/// This potentially buffers one frame if the rate limiter does not allow this frame.
333/// While the frame is buffered the underlying stream is no longer polled.
334#[derive(Debug)]
335pub(crate) struct RateLimited<S> {
336    inner: S,
337    bucket: Option<Bucket>,
338    bucket_refilled: Option<Pin<Box<time::Sleep>>>,
339    /// Counts how often reads from this stream have been rate-limited.
340    ///
341    /// Subscribe via [`RateLimited::limited_watcher`] to be notified about it, e.g. to
342    /// inform the client that it is being throttled.
343    limited_tx: watch::Sender<u64>,
344    /// Watcher for the per-client rate limit, if set.
345    ///
346    /// If set by constructing via [`RateLimited::from_watcher`], we check for
347    /// updates on each poll and reconfigure [`Self::bucket`] on changes.
348    ///
349    /// Uses `tokio::sync::watch` and not `n0_watcher` because in tokio's watch channel,
350    /// checking if a value changed is a cheap atomic version compare, whereas `n0_watcher`
351    /// acquires a `ReadLock` for each change check. This matters because we need to check this
352    /// on each read poll.
353    rate_limit_watcher: Option<watch::Receiver<Option<ClientRateLimit>>>,
354    metrics: Arc<Metrics>,
355}
356
357/// A token bucket used for rate-limiting.
358///
359/// The bucket holds up to `max` tokens, refilled at a fixed rate of
360/// `bytes_per_second` over `refill_period` intervals. [`Bucket::consume`] drains
361/// tokens and reports when the next attempt is allowed.
362///
363/// Exposed so that embedders mounting the relay protocol on a custom HTTP
364/// server (where the stream is already framed and the internal
365/// `RateLimited` wrapper cannot be used) can apply the same rate-limit
366/// semantics at the frame layer.
367#[derive(Debug)]
368pub struct Bucket {
369    // The current bucket fill
370    fill: i64,
371    // The maximum bucket fill
372    max: i64,
373    // The bucket's last fill time
374    last_fill: time::Instant,
375    // Interval length of one refill
376    refill_period: time::Duration,
377    // How much we re-fill per refill period
378    refill: i64,
379}
380
381#[allow(missing_docs)]
382#[stack_error(derive, add_meta)]
383pub struct InvalidBucketConfig {
384    max: i64,
385    bytes_per_second: i64,
386    refill_period: time::Duration,
387}
388
389impl Bucket {
390    /// Creates a new bucket starting full at `max` tokens, refilled at
391    /// `bytes_per_second` over `refill_period` intervals.
392    ///
393    /// # Errors
394    ///
395    /// Returns [`InvalidBucketConfig`] when `max`, `bytes_per_second`, or
396    /// `refill_period` are non-positive, or when the configuration would refill
397    /// less than one token per period.
398    pub fn new(
399        max: i64,
400        bytes_per_second: i64,
401        refill_period: time::Duration,
402    ) -> Result<Self, InvalidBucketConfig> {
403        // milliseconds is the tokio timer resolution
404        let refill = bytes_per_second.saturating_mul(refill_period.as_millis() as i64) / 1000;
405        ensure!(
406            max > 0 && bytes_per_second > 0 && refill_period.as_millis() as u32 > 0 && refill > 0,
407            InvalidBucketConfig {
408                max,
409                bytes_per_second,
410                refill_period
411            }
412        );
413        Ok(Self {
414            fill: max,
415            max,
416            last_fill: time::Instant::now(),
417            refill_period,
418            refill,
419        })
420    }
421
422    fn from_config(cfg: Option<ClientRateLimit>) -> Result<Option<Self>, InvalidBucketConfig> {
423        match cfg {
424            Some(cfg) => {
425                let bytes_per_second = u32::from(cfg.bytes_per_second);
426                let max_burst_bytes = cfg.max_burst_bytes.map_or(bytes_per_second / 10, u32::from);
427                Ok(Some(Bucket::new(
428                    max_burst_bytes as i64,
429                    bytes_per_second as i64,
430                    time::Duration::from_millis(100),
431                )?))
432            }
433            None => Ok(None),
434        }
435    }
436
437    fn update_state(&mut self) {
438        let now = time::Instant::now();
439        // div safety: self.refill_period.as_millis() is checked to be non-null in constructor
440        let refill_periods = now.saturating_duration_since(self.last_fill).as_millis() as u32
441            / self.refill_period.as_millis() as u32;
442        if refill_periods == 0 {
443            // Nothing to do - we won't refill yet
444            return;
445        }
446
447        self.fill = self
448            .fill
449            .saturating_add(refill_periods as i64 * self.refill);
450        self.fill = std::cmp::min(self.fill, self.max);
451        self.last_fill += self.refill_period * refill_periods;
452    }
453
454    /// Attempts to consume `bytes` tokens from the bucket.
455    ///
456    /// Returns `Ok(())` if enough tokens were available.
457    ///
458    /// # Errors
459    ///
460    /// Returns `Err(deadline)` with the [`time::Instant`] at which enough
461    /// tokens will have refilled to allow the request.
462    pub fn consume(&mut self, bytes: usize) -> Result<(), time::Instant> {
463        let bytes = i64::try_from(bytes).unwrap_or(i64::MAX);
464        self.update_state();
465
466        self.fill = self.fill.saturating_sub(bytes);
467
468        if self.fill > 0 {
469            return Ok(());
470        }
471
472        let missing = self.fill.saturating_neg();
473
474        let periods_needed = (missing / self.refill) + 1;
475        let periods_needed = u32::try_from(periods_needed).unwrap_or(u32::MAX);
476
477        Err(self.last_fill + periods_needed * self.refill_period)
478    }
479}
480
481impl<S> RateLimited<S> {
482    /// Wraps `io` with a live rate limit driven by `rate_limit_watcher`.
483    ///
484    /// The bucket is initialized from the watcher's current value and reconfigured
485    /// whenever it changes.
486    pub(crate) fn from_watcher(
487        io: S,
488        mut rate_limit_watcher: watch::Receiver<Option<ClientRateLimit>>,
489        metrics: Arc<Metrics>,
490    ) -> Result<Self, InvalidBucketConfig> {
491        let initial_cfg = *rate_limit_watcher.borrow_and_update();
492        let bucket = Bucket::from_config(initial_cfg)?;
493        Ok(Self {
494            inner: io,
495            bucket,
496            bucket_refilled: None,
497            limited_tx: watch::Sender::new(0),
498            rate_limit_watcher: Some(rate_limit_watcher),
499            metrics,
500        })
501    }
502
503    #[cfg(test)]
504    pub(crate) fn new(
505        inner: S,
506        max_burst_bytes: u32,
507        bytes_per_second: u32,
508        metrics: Arc<Metrics>,
509    ) -> Result<Self, InvalidBucketConfig> {
510        Ok(Self {
511            inner,
512            bucket: Some(Bucket::new(
513                max_burst_bytes as i64,
514                bytes_per_second as i64,
515                time::Duration::from_millis(100),
516            )?),
517            bucket_refilled: None,
518            limited_tx: watch::Sender::new(0),
519            rate_limit_watcher: None,
520            metrics,
521        })
522    }
523
524    #[cfg(test)]
525    pub(crate) fn unlimited(inner: S, metrics: Arc<Metrics>) -> Self {
526        Self {
527            inner,
528            bucket: None,
529            bucket_refilled: None,
530            limited_tx: watch::Sender::new(0),
531            rate_limit_watcher: None,
532            metrics,
533        }
534    }
535
536    /// Returns a watcher that counts how often reads from this stream have been rate-limited.
537    pub(crate) fn limited_watcher(&self) -> watch::Receiver<u64> {
538        self.limited_tx.subscribe()
539    }
540
541    /// Records metrics about being rate-limited.
542    fn record_rate_limited(&mut self, bytes: usize) {
543        // TODO: add a label for the frame type.
544        self.metrics.bytes_rx_ratelimited_total.inc_by(bytes as u64);
545        self.limited_tx.send_modify(|count| {
546            if *count == 0 {
547                self.metrics.conns_rx_ratelimited_total.inc();
548            }
549            *count += 1;
550        });
551    }
552}
553
554impl<S: ExportKeyingMaterial> ExportKeyingMaterial for RateLimited<S> {
555    fn export_keying_material<T: AsMut<[u8]>>(
556        &self,
557        output: T,
558        label: &[u8],
559        context: Option<&[u8]>,
560    ) -> Option<T> {
561        self.inner.export_keying_material(output, label, context)
562    }
563}
564
565impl<S: AsyncRead + Unpin> AsyncRead for RateLimited<S> {
566    #[instrument(name = "rate_limited_poll_read", skip_all)]
567    fn poll_read(
568        mut self: Pin<&mut Self>,
569        cx: &mut std::task::Context<'_>,
570        buf: &mut tokio::io::ReadBuf<'_>,
571    ) -> Poll<std::io::Result<()>> {
572        let this = &mut *self;
573
574        // Pick up a live rate change before consulting the bucket. `has_changed`
575        // is a cheap atomic version compare, so we only take the watch lock via
576        // `borrow_and_update` when the value actually changed.
577        if let Some(rate_limit_watcher) = this.rate_limit_watcher.as_mut()
578            && rate_limit_watcher.has_changed().unwrap_or(false)
579        {
580            let cfg = rate_limit_watcher.borrow_and_update();
581            match Bucket::from_config(*cfg) {
582                Ok(bucket) => {
583                    this.bucket = bucket;
584                    // The limit changed, so drop any pending refill wait: the new
585                    // bucket starts full and should be consulted immediately.
586                    this.bucket_refilled = None;
587                }
588                Err(err) => warn!(%err, "ignoring invalid live rate-limit update"),
589            }
590        }
591
592        let Some(bucket) = &mut this.bucket else {
593            // If there is no rate-limiter, then directly poll the inner.
594            return Pin::new(&mut this.inner).poll_read(cx, buf);
595        };
596
597        // If we're currently limited, wait until we've got some bucket space again
598        if let Some(bucket_refilled) = &mut this.bucket_refilled {
599            ready!(bucket_refilled.poll(cx));
600            this.bucket_refilled = None;
601        }
602
603        // We're not currently limited, let's read
604
605        // Poll inner for a new item.
606        let bytes_before = buf.remaining();
607        ready!(Pin::new(&mut this.inner).poll_read(cx, buf))?;
608        let bytes_read = bytes_before - buf.remaining();
609
610        // Record how much we've read, rate limit accordingly, if need be.
611        if let Err(refill_time) = bucket.consume(bytes_read) {
612            this.record_rate_limited(bytes_read);
613            this.bucket_refilled = Some(Box::pin(time::sleep_until(refill_time)));
614        }
615
616        Poll::Ready(Ok(()))
617    }
618}
619
620impl<S: AsyncWrite + Unpin> AsyncWrite for RateLimited<S> {
621    fn poll_write(
622        mut self: Pin<&mut Self>,
623        cx: &mut std::task::Context<'_>,
624        buf: &[u8],
625    ) -> Poll<Result<usize, std::io::Error>> {
626        Pin::new(&mut self.inner).poll_write(cx, buf)
627    }
628
629    fn poll_flush(
630        mut self: Pin<&mut Self>,
631        cx: &mut std::task::Context<'_>,
632    ) -> Poll<Result<(), std::io::Error>> {
633        Pin::new(&mut self.inner).poll_flush(cx)
634    }
635
636    fn poll_shutdown(
637        mut self: Pin<&mut Self>,
638        cx: &mut std::task::Context<'_>,
639    ) -> Poll<Result<(), std::io::Error>> {
640        Pin::new(&mut self.inner).poll_shutdown(cx)
641    }
642}
643
644#[cfg(test)]
645mod tests {
646    use std::sync::Arc;
647
648    use n0_error::{Result, StdResultExt};
649    use n0_future::time;
650    use n0_tracing_test::traced_test;
651    use tokio::io::{AsyncReadExt, AsyncWriteExt};
652
653    use super::Bucket;
654    use crate::server::{Metrics, streams::RateLimited};
655
656    #[tokio::test(start_paused = true)]
657    #[traced_test]
658    async fn test_ratelimiter() -> Result {
659        let (read, mut write) = tokio::io::duplex(4096);
660
661        let send_total = 10 * 1024 * 1024; // 10MiB
662        let send_data = vec![42u8; send_total];
663
664        let bytes_per_second = 12_345;
665
666        let mut rate_limited = RateLimited::new(
667            read,
668            bytes_per_second / 10,
669            bytes_per_second,
670            Arc::new(Metrics::default()),
671        )?;
672
673        let before = time::Instant::now();
674        n0_future::future::try_zip(
675            async {
676                let mut remaining = send_total;
677                let mut buf = [0u8; 4096];
678                while remaining > 0 {
679                    remaining -= rate_limited.read(&mut buf).await?;
680                }
681                Ok(())
682            },
683            async {
684                write.write_all(&send_data).await?;
685                write.flush().await
686            },
687        )
688        .await
689        .anyerr()?;
690
691        let duration = time::Instant::now().duration_since(before);
692        assert_ne!(duration.as_millis(), 0);
693
694        let actual_bytes_per_second = send_total as f64 / duration.as_secs_f64();
695        println!("{actual_bytes_per_second}");
696        assert_eq!(actual_bytes_per_second.round() as u32, bytes_per_second);
697
698        Ok(())
699    }
700
701    /// A rate limit installed live through the watch channel takes effect on an
702    /// existing stream without dropping it: the connection first reads a chunk
703    /// unlimited at full speed, then a limit is pushed mid-connection and the
704    /// throughput of the remainder matches the limit.
705    #[tokio::test(start_paused = true)]
706    #[traced_test]
707    async fn test_ratelimiter_live_update() -> Result {
708        use crate::server::ClientRateLimit;
709
710        let (read, mut write) = tokio::io::duplex(4096);
711
712        let bytes_per_second = 12_345u32;
713
714        // Read this much at full speed before the limit is installed, then read
715        // exactly 10MiB under the live limit and time only that second phase.
716        let unlimited_total = 1024 * 1024; // 1MiB
717        let limited_total = 10 * 1024 * 1024; // 10MiB
718        let send_data = vec![42u8; unlimited_total + limited_total];
719
720        // Start unlimited; the limit is installed mid-connection below.
721        let (rate_tx, rate_rx) = tokio::sync::watch::channel(None);
722        let mut rate_limited =
723            RateLimited::from_watcher(read, rate_rx, Arc::new(Metrics::default()))?;
724        let mut limit = ClientRateLimit::new(bytes_per_second.try_into().anyerr()?);
725        limit.max_burst_bytes = Some((bytes_per_second / 10).try_into().anyerr()?);
726
727        let mut unlimited_duration = time::Duration::ZERO;
728        let mut limited_duration = time::Duration::ZERO;
729        n0_future::future::try_zip(
730            async {
731                let mut buf = [0u8; 4096];
732
733                // Phase 1: no limit installed yet, pull bytes at full speed.
734                let before = time::Instant::now();
735                let mut remaining = unlimited_total;
736                while remaining > 0 {
737                    let len = remaining.min(buf.len());
738                    let n = rate_limited.read(&mut buf[..len]).await?;
739                    remaining -= n;
740                }
741                unlimited_duration = time::Instant::now().duration_since(before);
742
743                // Install the limit live, mid-connection, without dropping the stream.
744                rate_tx.send_replace(Some(limit));
745
746                // Phase 2: the limit is now in effect, time the remainder.
747                let before = time::Instant::now();
748                let mut remaining = limited_total;
749                while remaining > 0 {
750                    let len = remaining.min(buf.len());
751                    let n = rate_limited.read(&mut buf[..len]).await?;
752                    remaining -= n;
753                }
754                limited_duration = time::Instant::now().duration_since(before);
755                Ok(())
756            },
757            async {
758                write.write_all(&send_data).await?;
759                write.flush().await
760            },
761        )
762        .await
763        .anyerr()?;
764
765        // The unlimited phase is not throttled, so no (virtual) time passes.
766        assert_eq!(
767            unlimited_duration.as_millis(),
768            0,
769            "reading before the limit is installed runs at full speed",
770        );
771
772        let actual_bytes_per_second = limited_total as f64 / limited_duration.as_secs_f64();
773        assert_eq!(
774            actual_bytes_per_second.round() as u32,
775            bytes_per_second,
776            "the live-installed limit governs throughput once applied mid-connection",
777        );
778
779        Ok(())
780    }
781
782    #[tokio::test(start_paused = true)]
783    async fn test_bucket_high_refill() -> Result {
784        let bytes_per_second = i64::MAX;
785        let mut bucket = Bucket::new(i64::MAX, bytes_per_second, time::Duration::from_millis(100))?;
786        for _ in 0..100 {
787            time::sleep(time::Duration::from_millis(100)).await;
788            assert!(bucket.consume(1_000_000).is_ok());
789        }
790
791        Ok(())
792    }
793
794    #[tokio::test(start_paused = true)]
795    async fn smoke_test_bucket_high_consume() -> Result {
796        let bytes_per_second = 123_456;
797        let mut bucket = Bucket::new(
798            bytes_per_second / 10,
799            bytes_per_second,
800            time::Duration::from_millis(100),
801        )?;
802        for _ in 0..100 {
803            let Err(until) = bucket.consume(usize::MAX) else {
804                panic!("i64::MAX shouldn't be within limits");
805            };
806            time::sleep_until(until).await;
807        }
808
809        Ok(())
810    }
811}
iroh_relay/server/streams.rs

iroh_relay/server/
streams.rs
