strangers messages small

1 files changed, 20 insertions, 5 deletions

diff --git a/src/transport/libp2p/mod.rs b/src/transport/libp2p/mod.rs
index 056b1b2..2e670bf 100644
--- a/src/transport/libp2p/mod.rs
+++ b/src/transport/libp2p/mod.rs
@@ -238,6 +238,8 @@ pub(crate) struct Libp2pTransport {
     book: Arc<Mutex<Book>>,
     incoming: tokio::sync::Mutex<Option<libp2p_stream::IncomingStreams>>,
     stopped: watch::Receiver<bool>,
+    /// Whose messages may be large (`transport::Admission`).
+    admission: Arc<std::sync::OnceLock<crate::transport::Admission>>,
 }
 
 impl Libp2pTransport {
@@ -308,6 +310,7 @@ impl Libp2pTransport {
             book,
             incoming: tokio::sync::Mutex::new(incoming),
             stopped,
+            admission: Arc::default(),
         })
     }
 
@@ -339,15 +342,16 @@ impl Libp2pTransport {
     }
 }
 
-/// One message from an inbound stream.
-async fn read_message(mut stream: ::libp2p::Stream) -> Result<Vec<u8>> {
+/// One message from an inbound stream, of at most `limit` bytes (`MAX_MESSAGE`, or less
+/// from a peer not admitted: `transport::Admission`).
+async fn read_message(mut stream: ::libp2p::Stream, limit: usize) -> Result<Vec<u8>> {
     let mut length = [0u8; 4];
     tokio::time::timeout(STREAM_TIMEOUT, stream.read_exact(&mut length))
         .await
         .map_err(|_| anyhow!("no length within {STREAM_TIMEOUT:?}"))??;
     let length = u32::from_be_bytes(length) as usize;
-    if length > MAX_MESSAGE {
-        bail!("a message of {length} bytes; at most {MAX_MESSAGE} are read");
+    if length > limit {
+        bail!("a message of {length} bytes; at most {limit} are read from this peer");
     }
     let mut body = vec![0u8; length];
     tokio::time::timeout(stream_timeout(length), stream.read_exact(&mut body))
@@ -635,6 +639,10 @@ impl Transport for Libp2pTransport {
         Ok(())
     }
 
+    fn set_admission(&self, admission: crate::transport::Admission) {
+        let _ = self.admission.set(admission);
+    }
+
     /// Reads inbound streams (`READS` at once) and handles their messages one at a
     /// time, in the order they were read; PINGs and PONGs at once (`is_probe`).
     async fn start_consumer(&self, handler: MessageHandler) -> Result<()> {
@@ -645,6 +653,7 @@ impl Transport for Libp2pTransport {
         let (sender, mut received) = mpsc::channel::<(PeerId, ParsedMessage, usize)>(READS);
         let reads = Arc::new(Semaphore::new(READS));
         let probes = Arc::clone(&handler);
+        let admission = Arc::clone(&self.admission);
         tokio::spawn(async move {
             while let Some((peer, stream)) = incoming.next().await {
                 let Ok(permit) = Arc::clone(&reads).acquire_owned().await else {
@@ -652,8 +661,14 @@ impl Transport for Libp2pTransport {
                 };
                 let sender = sender.clone();
                 let probes = Arc::clone(&probes);
+                let limit = crate::transport::Admission::limit(
+                    admission.get(),
+                    SCHEME,
+                    &peer.to_string(),
+                    MAX_MESSAGE,
+                );
                 tokio::spawn(async move {
-                    match read_message(stream).await {
+                    match read_message(stream, limit).await {
                         Ok(body) => match wire::decode(&body) {
                             Ok(message) if is_probe(&message) => {
                                 hand_on_probe(