From 67c46d022ea71e230d0dd23ace229adbfc55e1d5 Mon Sep 17 00:00:00 2001
From: Bulat Fatykh <68301540+gitoutofthere@users.noreply.github.com>
Date: Wed, 30 Sep 2026 19:54:21 +0300
Subject: [PATCH] Validate a master client's address before it holds a
 connection slot (#14)

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 crates/iw4l-master/src/main.rs | 56 ++++++++++++++++++++++++++++++++--
 1 file changed, 54 insertions(+), 2 deletions(-)

diff --git a/crates/iw4l-master/src/main.rs b/crates/iw4l-master/src/main.rs
index f0c2f7c..706c113 100644
--- a/crates/iw4l-master/src/main.rs
+++ b/crates/iw4l-master/src/main.rs
@@ -3,7 +3,7 @@
 use std::collections::HashMap;
 use std::fs::File;
 use std::io::{BufReader, Write};
-use std::net::SocketAddr;
+use std::net::{IpAddr, SocketAddr};
 use std::path::{Path, PathBuf};
 use std::sync::Arc;
 use std::sync::atomic::{AtomicU64, Ordering};
@@ -35,12 +35,48 @@ const CLI_DEADLINE: Duration = Duration::from_secs(8);
 const BOOTSTRAP_FORWARD: Duration = Duration::from_secs(8);
 const HELLO_DEADLINE: Duration = Duration::from_secs(8);
 const MAX_CONNECTIONS: usize = 256;
+/// A household or a LAN party shares one address.
+const MAX_CONNECTIONS_PER_ADDRESS: usize = 32;
 
 const CLOSE_SERVICE_ERROR: u32 = 1;
 const CLOSE_SERVICE_DONE: u32 = 0;
 
 const MAX_CLOSE_REASON_BYTES: usize = 120;
 
+struct AddressSlot {
+    counts: Arc<std::sync::Mutex<HashMap<IpAddr, usize>>>,
+    ip: IpAddr,
+}
+
+impl AddressSlot {
+    fn claim(counts: &Arc<std::sync::Mutex<HashMap<IpAddr, usize>>>, ip: IpAddr) -> Option<Self> {
+        let mut held = counts.lock().expect("address counts poisoned");
+        let count = held.entry(ip).or_insert(0);
+        if *count >= MAX_CONNECTIONS_PER_ADDRESS {
+            return None;
+        }
+        *count += 1;
+        Some(Self {
+            counts: Arc::clone(counts),
+            ip,
+        })
+    }
+}
+
+impl Drop for AddressSlot {
+    fn drop(&mut self) {
+        let Ok(mut held) = self.counts.lock() else {
+            return;
+        };
+        if let Some(count) = held.get_mut(&self.ip) {
+            *count -= 1;
+            if *count == 0 {
+                held.remove(&self.ip);
+            }
+        }
+    }
+}
+
 enum ConnTask {
     Writer(Result<()>),
     Bootstrap,
@@ -472,7 +508,23 @@ async fn serve(bind: SocketAddr, cert_path: &Path, key_path: &Path) -> Result<()
         ALPN
     )?;
     let connection_slots = Arc::new(Semaphore::new(MAX_CONNECTIONS));
+    let address_counts = Arc::new(std::sync::Mutex::new(HashMap::new()));
     while let Some(incoming) = endpoint.accept().await {
+        // No slot before the address is proven: a forged Initial would hold
+        // one until its handshake timed out.
+        if !incoming.remote_address_validated() && incoming.may_retry() {
+            let _ = incoming.retry();
+            continue;
+        }
+        let remote = incoming.remote_address();
+        let Some(address_slot) = AddressSlot::claim(&address_counts, remote.ip()) else {
+            let _ = writeln!(
+                std::io::stderr(),
+                "connection refused: {remote} at per-address cap {MAX_CONNECTIONS_PER_ADDRESS}"
+            );
+            incoming.refuse();
+            continue;
+        };
         let Ok(permit) = connection_slots.clone().try_acquire_owned() else {
             let _ = writeln!(
                 std::io::stderr(),
@@ -483,7 +535,6 @@ async fn serve(bind: SocketAddr, cert_path: &Path, key_path: &Path) -> Result<()
         };
         let state = Arc::clone(&state);
         let connection_id = next_connection_id.fetch_add(1, Ordering::Relaxed);
-        let remote = incoming.remote_address();
         let started = Instant::now();
         let _ = writeln!(
             std::io::stderr(),
@@ -491,6 +542,7 @@ async fn serve(bind: SocketAddr, cert_path: &Path, key_path: &Path) -> Result<()
         );
         tokio::spawn(async move {
             let _permit = permit;
+            let _address_slot = address_slot;
             match incoming.await {
                 Ok(connection) => {
                     let _ = writeln!(