strangers messages small

1 files changed, 64 insertions, 0 deletions

diff --git a/src/transport/mod.rs b/src/transport/mod.rs
index e81adfb..00e07e5 100644
--- a/src/transport/mod.rs
+++ b/src/transport/mod.rs
@@ -871,6 +871,41 @@ pub(crate) fn hand_on_probe(
     });
 }
 
+/// Most bytes a message from a peer that isn't established may take (`Admission`):
+/// before the handshake a peer can only say HELLO (or answer one: a HELLO_REPLY lists at
+/// most 64 peers with 8 addresses each, some tens of KiB), WHO or ME, or PING or PONG.
+/// Larger frames from it are refused before anything is allocated for them.
+pub(crate) const MAX_UNADMITTED: usize = 256 << 10;
+
+/// Whether a transport address (`scheme:identity`) is a node that proved itself a
+/// member of our swarm (its handshake, `swarm::peer`): only its messages may be large;
+/// a stranger's are refused above `MAX_UNADMITTED`, unread. Set by the identity layer
+/// (`identity::IdentityTransport`), which knows which node is at an address; without
+/// it (tests on a bare transport), everyone is.
+#[derive(Clone)]
+pub(crate) struct Admission(std::sync::Arc<dyn Fn(&str) -> bool + Send + Sync>);
+
+impl Admission {
+    pub(crate) fn new(admits: impl Fn(&str) -> bool + Send + Sync + 'static) -> Self {
+        Self(std::sync::Arc::new(admits))
+    }
+
+    /// The most bytes a message from `scheme:identity` may take.
+    pub(crate) fn limit(
+        admission: Option<&Self>,
+        scheme: &str,
+        identity: &str,
+        most: usize,
+    ) -> usize {
+        match admission {
+            Some(admission) if !(admission.0)(&format!("{scheme}:{identity}")) => {
+                most.min(MAX_UNADMITTED)
+            }
+            _ => most,
+        }
+    }
+}
+
 #[async_trait]
 pub(crate) trait Transport: Send + Sync {
     /// This node's identity on the transport (for Mycelium: its public key).
@@ -1038,6 +1073,18 @@ pub(crate) trait Transport: Send + Sync {
         stats::TransportStats::default()
     }
 
+    /// Who may send large messages (`Admission`), for a transport that reads messages
+    /// itself (it frames them). By default ignored.
+    fn set_admission(&self, admission: Admission) {
+        let _ = admission;
+    }
+
+    /// `node` proved itself a member of our swarm (its handshake): its messages may be
+    /// large from now on (`Admission`). By default nothing.
+    fn admit(&self, node: &str) {
+        let _ = node;
+    }
+
     /// Resolves a seed (anything that leads to a node, e.g. an address) to that node's
     /// identity, as `local_id` gives it on that node.
     async fn resolve_seed(&self, seed: &str) -> Result<String>;
@@ -1121,6 +1168,23 @@ mod tests {
         assert_eq!(seen.len(), heads.len());
     }
 
+    /// A stranger's messages are limited, a member's aren't; without an identity layer
+    /// (no admission), nobody's.
+    #[test]
+    fn strangers_messages_are_small() {
+        let admission = Admission::new(|address| address == "s:member");
+        let most = 64 << 20;
+        assert_eq!(
+            Admission::limit(Some(&admission), "s", "member", most),
+            most
+        );
+        assert_eq!(
+            Admission::limit(Some(&admission), "s", "stranger", most),
+            MAX_UNADMITTED
+        );
+        assert_eq!(Admission::limit(None, "s", "stranger", most), most);
+    }
+
     /// A tag names its epoch, not one of another generation or another ID, and resolves
     /// to the epoch it names among those known.
     #[test]