1 files changed, 72 insertions, 13 deletions

diff --git a/src/transport/sim.rs b/src/transport/sim.rs
index 500050b..2f167dc 100644
--- a/src/transport/sim.rs
+++ b/src/transport/sim.rs
@@ -1,8 +1,9 @@
 //! A simulated network for tests: many nodes in one process, each with a
 //! [`SimTransport`] on a shared [`SimNet`].
 //!
-//! Messages go through the Mycelium wire encoding (`mycelium::encode_message`), so the
-//! codec is exercised as in production. The net delays, loses and duplicates them per
+//! Messages go through the wire encoding (`WIRE`, as libp2p and Tor send it), so the
+//! codec is exercised as in production; a net can use a transport's envelope instead
+//! (`SimNet::with_codec`; Mycelium's tests run a swarm on its envelope). The net delays, loses and duplicates them per
 //! link, from a seeded generator; nodes can be taken down (everything to and from them
 //! is lost) and brought up again; and a test can inject a message as if a node had sent
 //! it, e.g. one that arrives late.
@@ -24,11 +25,14 @@ use bytes::Bytes;
 use tokio::sync::{mpsc, oneshot};
 use tracing::{debug, warn};
 
-use super::mycelium::{decode_message, encode_message};
 use super::{EncodedMessage, MessageHandler, MessageKind, ParsedMessage, Transport};
 use crate::locks::MutexExt;
 use crate::testrand::Rng;
 
+/// How long the simulated net may take to deliver a message: well above its delays.
+/// A catch-up source that stops sending sync messages is given up after this.
+pub(crate) const DELIVERY_LIMIT: Duration = Duration::from_secs(3);
+
 /// How one direction of a link behaves.
 #[derive(Clone, Debug)]
 pub(crate) struct LinkParams {
@@ -79,6 +83,9 @@ struct NetInner {
     links: HashMap<(String, String), LinkParams>,
     down: HashSet<String>,
     inboxes: HashMap<String, mpsc::UnboundedSender<Delivery>>,
+    /// Per node, its current run (`restart`): what a transport of an earlier run sends
+    /// is lost, as from a process that is gone.
+    runs: HashMap<String, u64>,
     counts: NetCounts,
 }
 
@@ -86,17 +93,26 @@ struct NetInner {
 #[derive(Clone)]
 pub(crate) struct SimNet {
     inner: Arc<Mutex<NetInner>>,
+    codec: Codec,
 }
 
 impl SimNet {
+    /// Messages in the wire encoding alone (`WIRE`).
     pub(crate) fn new(seed: u64) -> Self {
+        Self::with_codec(seed, WIRE)
+    }
+
+    /// Messages put into bytes by `codec`.
+    pub(crate) fn with_codec(seed: u64, codec: Codec) -> Self {
         Self {
+            codec,
             inner: Arc::new(Mutex::new(NetInner {
                 rng: Rng::new(seed),
                 default_link: LinkParams::default(),
                 links: HashMap::new(),
                 down: HashSet::new(),
                 inboxes: HashMap::new(),
+                runs: HashMap::new(),
                 counts: NetCounts::default(),
             })),
         }
@@ -109,15 +125,31 @@ impl SimNet {
     /// A transport for a new node `id` on this net.
     pub(crate) fn transport(&self, id: &str) -> Arc<SimTransport> {
         let (tx, rx) = mpsc::unbounded_channel();
-        let previous = self.lock().inboxes.insert(id.to_owned(), tx);
-        assert!(previous.is_none(), "node {id} is already on the net");
+        let run = {
+            let mut inner = self.lock();
+            let previous = inner.inboxes.insert(id.to_owned(), tx);
+            assert!(previous.is_none(), "node {id} is already on the net");
+            *inner.runs.entry(id.to_owned()).or_default()
+        };
         Arc::new(SimTransport {
             id: id.to_owned(),
+            run,
             net: self.clone(),
             inbox: tokio::sync::Mutex::new(Some(rx)),
         })
     }
 
+    /// Node `id`'s process is gone, as if it crashed: its transport neither sends nor
+    /// receives any more (messages on their way to it are lost), and [`transport`]
+    /// may give `id` a new one, for its next run.
+    ///
+    /// [`transport`]: Self::transport
+    pub(crate) fn restart(&self, id: &str) {
+        let mut inner = self.lock();
+        inner.inboxes.remove(id);
+        *inner.runs.entry(id.to_owned()).or_default() += 1;
+    }
+
     /// How links without their own parameters behave.
     pub(crate) fn set_default_link(&self, params: LinkParams) {
         self.lock().default_link = params;
@@ -145,12 +177,13 @@ impl SimNet {
     }
 
     /// Hands a message to the net: decides now whether and when it arrives.
-    fn send(&self, from: &str, message: &EncodedMessage) {
+    fn send(&self, from: &str, run: u64, message: &EncodedMessage) {
         let delays = {
             let mut guard = self.lock();
             let inner = &mut *guard;
             inner.counts.sent += 1;
-            if inner.down.contains(from) || inner.down.contains(&message.peer) {
+            let gone = inner.runs.get(from).is_some_and(|current| *current != run);
+            if gone || inner.down.contains(from) || inner.down.contains(&message.peer) {
                 inner.counts.lost += 1;
                 return;
             }
@@ -215,7 +248,7 @@ impl SimNet {
     /// net. For messages a real node wouldn't send at that moment, e.g. one sent long
     /// ago and delivered late.
     pub(crate) async fn inject(&self, from: &str, to: &str, message: &ParsedMessage) -> Result<()> {
-        let encoded = encode_message(to, message)?;
+        let encoded = (self.codec.encode)(to, message)?;
         let (handled_tx, handled_rx) = oneshot::channel();
         let delivery = Delivery {
             from: from.to_owned(),
@@ -245,10 +278,35 @@ impl SimNet {
 /// are node IDs.
 pub(crate) struct SimTransport {
     id: String,
+    /// The run of its node it belongs to (`SimNet::restart`).
+    run: u64,
     net: SimNet,
     inbox: tokio::sync::Mutex<Option<mpsc::UnboundedReceiver<Delivery>>>,
 }
 
+/// How a simulated net puts messages into bytes and back (`SimNet::with_codec`): the
+/// wire encoding alone (`WIRE`, as libp2p and Tor send it), or a transport's envelope
+/// around it (Mycelium's: `mycelium::CODEC`), so a whole swarm runs on that envelope.
+#[derive(Clone, Copy)]
+pub(crate) struct Codec {
+    pub(crate) encode: fn(&str, &ParsedMessage) -> Result<EncodedMessage>,
+    pub(crate) decode: fn(&[u8]) -> Result<ParsedMessage>,
+}
+
+/// The wire encoding alone (`transport::wire`).
+pub(crate) const WIRE: Codec = Codec {
+    encode: wire_encode,
+    decode: wire_decode,
+};
+
+fn wire_encode(peer: &str, message: &ParsedMessage) -> Result<EncodedMessage> {
+    super::wire::encoded(peer, message)
+}
+
+fn wire_decode(body: &[u8]) -> Result<ParsedMessage> {
+    Ok(super::wire::decode(body)?)
+}
+
 #[async_trait]
 impl Transport for SimTransport {
     fn local_id(&self) -> &str {
@@ -260,8 +318,8 @@ impl Transport for SimTransport {
         !id.is_empty()
     }
 
-    fn delivery_limit(&self) -> std::time::Duration {
-        crate::simtest::DELIVERY_LIMIT
+    fn delivery_limit(&self) -> Duration {
+        DELIVERY_LIMIT
     }
 
     /// A seed is a node ID; one no node has doesn't resolve (yet).
@@ -274,13 +332,13 @@ impl Transport for SimTransport {
     }
 
     fn encode(&self, peer: &str, message: &ParsedMessage) -> Result<EncodedMessage> {
-        encode_message(peer, message)
+        (self.net.codec.encode)(peer, message)
     }
 
     /// Always accepted: the net loses messages silently, like Mycelium once it has
     /// taken them.
     async fn send_encoded(&self, message: &EncodedMessage) -> Result<()> {
-        self.net.send(&self.id, message);
+        self.net.send(&self.id, self.run, message);
         Ok(())
     }
 
@@ -290,7 +348,7 @@ impl Transport for SimTransport {
             bail!("consumer of {} already started", self.id);
         };
         while let Some(delivery) = inbox.recv().await {
-            match decode_message(&delivery.body) {
+            match (self.net.codec.decode)(&delivery.body) {
                 Ok(message) => {
                     debug!(node = %self.id, from = %delivery.from, message = ?message, "Simulated delivery");
                     if let Err(error) = handler(delivery.from, message, delivery.payload_len).await
@@ -328,6 +386,7 @@ mod tests {
             epoch_seq: 0,
             ready: true,
             syncs_with_you: false,
+            leaving: false,
         })
     }