Replace bach’s UDP simulation backend with a first-pass turmoil-net adapter

bach/Cargo.toml

@@ -13,7 +13,7 @@ default = ["net"]
 coop = []
 full = ["coop", "metrics", "net", "net-monitor", "tracing"]
 metrics = ["dep:metrics"]
-net = ["dep:bytes", "dep:s2n-quic-core", "dep:siphasher"]
+net = ["dep:bytes", "dep:s2n-quic-core", "dep:siphasher", "dep:turmoil-net"]
 net-monitor = []
 tokio-compat = ["tokio/time"]
 tracing = ["dep:tracing"]
@@ -36,6 +36,7 @@ siphasher = { version = "1", default-features = false, optional = true }
 slotmap = "1"
 tokio = { version = "1", default-features = false, features = ["sync"] }
 tracing = { version = "0.1", optional = true }
+turmoil-net = { version = "0.1.0", optional = true }
 
 [dev-dependencies]
 bolero.workspace = true

bach/src/environment/net/registry.rs

@@ -2,17 +2,20 @@ use crate::{
     environment::net::{
         ip,
         monitor::List as Monitors,
-        port,
-        queue::{self, Dispatch},
-        socket::{self, reservation},
+        queue,
+        socket::{self, udp::Socket as UdpSocket},
     },
     group::Group,
     net::{monitor::Monitor, IpAddr},
     scope::define,
 };
+use core::{future::Future, pin::pin, task::Context};
 use std::{collections::HashMap, io};
 
-use super::{ip::transport, pcap};
+use super::ip::{header, transport, Packet, Transport};
+use bytes::Bytes;
+use turmoil_net::shim::tokio::net::UdpSocket as TurmoilUdpSocket;
+use turmoil_net::{EnterGuard, HostId, Net};
 
 define!(scope, Box<Registry>);
 
@@ -28,12 +31,12 @@ pub(crate) fn with_registry<F: FnOnce(&mut Registry) -> io::Result<R>, R>(f: F)
 
 pub struct Registry {
     hostnames: HashMap<String, (Group, IpAddr)>,
-    senders: Dispatch,
-    groups: HashMap<Group, GroupState>,
+    group_ids: HashMap<Group, HostId>,
     ips: ip::Allocator,
-    pcaps: pcap::Registry,
-    queue_alloc: Box<dyn queue::Allocator>,
     monitors: Monitors,
+    #[allow(dead_code)]
+    queue_alloc: Box<dyn queue::Allocator>,
+    guard: Option<EnterGuard>,
 }
 
 impl Default for Registry {
@@ -47,12 +50,11 @@ impl Registry {
         let monitors = Monitors::default();
         Self {
             hostnames: HashMap::new(),
-            senders: Dispatch::new(monitors.clone()),
-            groups: HashMap::new(),
+            group_ids: HashMap::new(),
             ips: ip::Allocator::default(),
-            pcaps: Default::default(),
             queue_alloc: queue,
             monitors,
+            guard: None,
         }
     }
 
@@ -61,7 +63,10 @@ impl Registry {
     }
 
     pub fn set_pcap_dir<P: Into<std::path::PathBuf>>(&mut self, pcap: P) -> io::Result<()> {
-        self.pcaps.set_dir(pcap)
+        // The first-pass turmoil-net backend does not emit PCAP files yet, but keep the
+        // builder method available so existing callers continue to compile.
+        let _ = pcap.into();
+        Ok(())
     }
 
     pub fn set_subnet(&mut self, subnet: IpAddr) {
@@ -80,23 +85,10 @@ impl Registry {
             }
         }
 
-        if let Some((owner, ip)) = self.hostnames.get(name).cloned() {
-            // the owner would have already resolved itself in the pcap
-            if owner == *group {
+        if let Some((owner, ip)) = self.hostnames.get(name).copied() {
+            if owner == *group || self.guard.is_none() {
                 return Ok(ip);
             }
-
-            let group_name = group.name();
-
-            // inject a DNS packet in the pcap
-            let first_time = self.pcaps.dns(group, name, &ip);
-
-            // if this is the first time `group` has queried `name`, then do a reverse query
-            // on the owner so the pcaps come through correctly on the other side
-            if first_time {
-                let _ = self.resolve_host(&owner, &group_name);
-            }
-
             return Ok(ip);
         }
 
@@ -111,13 +103,7 @@ impl Registry {
             ));
         }
 
-        let ip = self.ips.allocate();
-        self.hostnames.insert(group_name, (*group, ip));
-        self.groups.insert(*group, GroupState::default());
-
-        self.pcaps.dns(group, name, &ip);
-
-        Ok(ip)
+        self.ensure_group_host(group)
     }
 
     pub fn register_monitor<M: Monitor>(&mut self, monitor: M) {
@@ -129,7 +115,9 @@ impl Registry {
         group: &Group,
         options: &socket::Options,
     ) -> std::io::Result<Box<dyn socket::Socket>> {
-        let group_ip = self.resolve_host(group, &group.name())?;
+        let group_ip = self.ensure_group_host(group)?;
+        self.prepare()?;
+        self.set_current_group(group)?;
 
         let mut local_addr = options.local_addr;
 
@@ -142,40 +130,161 @@ impl Registry {
             ));
         }
 
-        let state = self.groups.get_mut(group).unwrap();
+        self.monitors
+            .on_socket_opened(&local_addr, transport::Kind::Udp)?;
 
-        let reservation = if local_addr.port() == 0 {
-            let res = state.udp.ephemeral()?;
-            local_addr.set_port(res.port());
-            res
-        } else {
-            state.udp.reserve(local_addr.port(), options.reuse_port)?
+        let socket = bind_udp_socket(local_addr)?;
+        let socket = UdpSocket::new(socket, self.monitors.clone())?;
+        Ok(Box::new(socket))
+    }
+
+    pub fn set_current_group(&mut self, group: &Group) -> io::Result<()> {
+        self.prepare()?;
+        let Some(host_id) = self.group_ids.get(group).copied() else {
+            return Err(io::Error::new(
+                io::ErrorKind::NotFound,
+                format!("group `{group}` is not registered in turmoil-net"),
+            ));
         };
+        turmoil_net::set_current(host_id);
+        Ok(())
+    }
 
-        self.monitors
-            .on_socket_opened(&local_addr, transport::Kind::Udp)?;
+    pub fn drain_packets(&self) -> Vec<turmoil_net::Packet> {
+        let Some(guard) = self.guard.as_ref() else {
+            return Vec::new();
+        };
 
-        let queue::PacketQueue {
-            local_sender: sender,
-            local_receiver: receiver,
-            remote_sender,
-        } = self.queue_alloc.for_udp(
-            group,
-            local_addr,
-            &self.senders,
-            &self.monitors,
-            &mut self.pcaps,
-        );
-
-        let reservation = (reservation, self.senders.reserve(local_addr, remote_sender));
-        let socket = socket::udp::Socket::new(sender, receiver, local_addr, self.monitors.clone());
-        let socket = reservation::Socket::new(socket, reservation);
+        let mut packets = Vec::new();
+        guard.egress_all(&mut packets);
+        packets
+    }
 
-        Ok(Box::new(socket))
+    pub fn deliver(&self, packet: turmoil_net::Packet) {
+        if let Some(guard) = self.guard.as_ref() {
+            guard.deliver(packet);
+        }
     }
+
+    pub fn monitors(&self) -> Monitors {
+        self.monitors.clone()
+    }
+
+    fn prepare(&mut self) -> io::Result<()> {
+        if self.guard.is_some() {
+            return Ok(());
+        }
+
+        let mut groups = crate::group::list();
+        groups.sort_by_key(Group::id);
+
+        for group in groups {
+            let _ = self.ensure_group_host(&group)?;
+        }
+
+        let mut hosts = self
+            .hostnames
+            .values()
+            .copied()
+            .collect::<Vec<(Group, IpAddr)>>();
+        hosts.sort_by_key(|(group, _)| group.id());
+
+        let mut net = Net::new();
+        let mut group_ids = HashMap::with_capacity(hosts.len());
+
+        for (group, ip) in hosts {
+            let host_id = net.add_host(ip);
+            group_ids.insert(group, host_id);
+        }
+
+        self.group_ids = group_ids;
+        self.guard = Some(net.enter());
+        Ok(())
+    }
+
+    fn ensure_group_host(&mut self, group: &Group) -> io::Result<IpAddr> {
+        let name = group.name();
+        if let Some((_, ip)) = self.hostnames.get(&name).copied() {
+            return Ok(ip);
+        }
+
+        if self.guard.is_some() {
+            return Err(io::Error::other(
+                "adding new groups after turmoil-net initialization is not yet supported",
+            ));
+        }
+
+        let ip = self.ips.allocate();
+        self.hostnames.insert(name, (*group, ip));
+        Ok(ip)
+    }
+}
+
+pub(crate) fn monitor_packet(packet: &turmoil_net::Packet) -> Option<Packet> {
+    let transport = match &packet.payload {
+        turmoil_net::Transport::Udp(udp) => Transport::Udp(transport::Udp {
+            source: udp.src_port,
+            destination: udp.dst_port,
+            payload: Bytes::clone(&udp.payload),
+            checksum: 0,
+        }),
+        turmoil_net::Transport::Tcp(_) => return None,
+    };
+
+    let header = match (packet.src, packet.dst) {
+        (IpAddr::V4(source), IpAddr::V4(destination)) => header::V4 {
+            source,
+            destination,
+            dscp: 0,
+            ecn: 0,
+            df: true,
+            id: 0,
+            ttl: packet.ttl,
+        }
+        .into(),
+        (IpAddr::V6(source), IpAddr::V6(destination)) => header::V6 {
+            source,
+            destination,
+            dscp: 0,
+            ecn: 0,
+            flow_label: 0,
+            hop_limit: packet.ttl,
+        }
+        .into(),
+        (IpAddr::V4(source), IpAddr::V6(destination)) => header::V6 {
+            source: source.to_ipv6_mapped(),
+            destination,
+            dscp: 0,
+            ecn: 0,
+            flow_label: 0,
+            hop_limit: packet.ttl,
+        }
+        .into(),
+        (IpAddr::V6(source), IpAddr::V4(destination)) => header::V6 {
+            source,
+            destination: destination.to_ipv6_mapped(),
+            dscp: 0,
+            ecn: 0,
+            flow_label: 0,
+            hop_limit: packet.ttl,
+        }
+        .into(),
+    };
+
+    let mut packet = Packet { header, transport };
+    packet.update_checksum();
+    Some(packet)
 }
 
-#[derive(Default)]
-struct GroupState {
-    udp: port::Allocator,
+fn bind_udp_socket(local_addr: std::net::SocketAddr) -> io::Result<TurmoilUdpSocket> {
+    let mut future = pin!(TurmoilUdpSocket::bind(local_addr));
+    let waker = crate::task::waker::noop();
+    let mut cx = Context::from_waker(&waker);
+
+    match Future::poll(future.as_mut(), &mut cx) {
+        core::task::Poll::Ready(result) => result,
+        core::task::Poll::Pending => Err(io::Error::other(
+            "turmoil-net UDP bind unexpectedly returned pending",
+        )),
+    }
 }