[examples/reshare] Set is_dkg Correctly after Restart (#3615)

Author: patrick-ogrady

examples/reshare/Cargo.toml

@@ -38,6 +38,7 @@ thiserror.workspace = true
 tracing.workspace = true
 
 [dev-dependencies]
+commonware-p2p = { workspace = true, features = ["mocks"] }
 tracing-subscriber.workspace = true
 
 [[bin]]

examples/reshare/src/dkg/actor.rs

@@ -233,7 +233,6 @@ where
         mut callback: Box<dyn UpdateCallBack<V, C::PublicKey>>,
     ) {
         let max_read_size = NZU32!(self.peer_config.max_participants_per_round());
-        let is_dkg = output.is_none();
         let epocher = FixedEpocher::new(BLOCKS_PER_EPOCH);
 
         // Initialize persistent state
@@ -262,6 +261,7 @@ where
         'actor: loop {
             // Get latest epoch and state
             let (epoch, epoch_state) = storage.epoch().expect("epoch should be initialized");
+            let is_dkg = epoch_state.output.is_none();
 
             // Prune everything older than the previous epoch
             if let Some(prev) = epoch.previous() {
@@ -660,3 +660,163 @@ where
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::{dkg::ContinueOnUpdate, orchestrator::Message, setup::PeerConfig};
+    use commonware_cryptography::{
+        bls12381::{dkg::deal, primitives::variant::MinSig},
+        ed25519::{PrivateKey, PublicKey as Ed25519PublicKey},
+        transcript::Summary,
+        Sha256, Signer,
+    };
+    use commonware_macros::test_traced;
+    use commonware_math::algebra::Random;
+    use commonware_p2p::{utils::mocks::inert_channel, PeerSetSubscription, Provider};
+    use commonware_runtime::{deterministic, Runner};
+    use commonware_utils::{channel::mpsc, N3f1, TryCollect, NZU32};
+    use core::marker::PhantomData;
+    use std::collections::BTreeMap;
+
+    #[derive(Clone, Debug)]
+    struct NoopManager<P: PublicKey>(PhantomData<P>);
+
+    impl<P: PublicKey> Default for NoopManager<P> {
+        fn default() -> Self {
+            Self(PhantomData)
+        }
+    }
+
+    impl<P: PublicKey> Provider for NoopManager<P> {
+        type PublicKey = P;
+
+        async fn peer_set(&mut self, _: u64) -> Option<TrackedPeers<Self::PublicKey>> {
+            None
+        }
+
+        async fn subscribe(&mut self) -> PeerSetSubscription<Self::PublicKey> {
+            let (_, rx) = mpsc::unbounded_channel();
+            rx
+        }
+    }
+
+    impl<P: PublicKey> Manager for NoopManager<P> {
+        async fn track<R>(&mut self, _: u64, _: R)
+        where
+            R: Into<TrackedPeers<Self::PublicKey>> + Send,
+        {
+        }
+    }
+
+    fn peer_config(
+        total: u64,
+        per_round: Vec<u32>,
+    ) -> (
+        PeerConfig<Ed25519PublicKey>,
+        BTreeMap<Ed25519PublicKey, PrivateKey>,
+    ) {
+        let participants = (0..total)
+            .map(|seed| {
+                let signer = PrivateKey::from_seed(seed);
+                (signer.public_key(), signer)
+            })
+            .collect::<BTreeMap<_, _>>();
+        let peer_config = PeerConfig {
+            num_participants_per_round: per_round,
+            participants: participants.keys().cloned().try_collect().unwrap(),
+        };
+        (peer_config, participants)
+    }
+
+    #[test_traced]
+    fn recovered_storage_controls_dkg_mode_on_restart() {
+        let executor = deterministic::Runner::seeded(8);
+        executor.start(|mut context| async move {
+            // Seed a mid-life state well past the bootstrap epoch so the recovered round is
+            // unambiguously not the initial DKG. Per production semantics, the stored output
+            // carries the current round's dealers as its players (produced by the prior
+            // reshare), so deal with `dealers(RECOVERED_ROUND)`.
+            const RECOVERED_EPOCH: u64 = 5;
+            const RECOVERED_ROUND: u64 = 5;
+            let (peer_config, participants) = peer_config(6, vec![4]);
+            let first_player = peer_config
+                .dealers(RECOVERED_ROUND)
+                .iter()
+                .next()
+                .cloned()
+                .expect("recovered dealer exists");
+            let signer = participants
+                .get(&first_player)
+                .cloned()
+                .expect("signer should exist");
+            let (output, shares) = deal::<MinSig, _, N3f1>(
+                &mut context,
+                Default::default(),
+                peer_config.dealers(RECOVERED_ROUND),
+            )
+            .expect("deal should succeed");
+            let share = shares.get_value(&first_player).cloned();
+            let partition_prefix = format!("recovered_restart_{first_player}");
+
+            // Seed durable state that looks like a completed reshare several rounds in, even
+            // though the restarted actor will be given stale bootstrap inputs below.
+            let mut storage = Storage::<_, MinSig, Ed25519PublicKey>::init(
+                context.with_label("seed_storage"),
+                &partition_prefix,
+                NZU32!(peer_config.max_participants_per_round()),
+                crate::dkg::MAX_SUPPORTED_MODE,
+            )
+            .await;
+            storage
+                .set_epoch(
+                    Epoch::new(RECOVERED_EPOCH),
+                    EpochState {
+                        round: RECOVERED_ROUND,
+                        rng_seed: Summary::random(&mut context),
+                        output: Some(output),
+                        share,
+                    },
+                )
+                .await;
+            drop(storage);
+
+            // Restart the actor with stale bootstrap inputs (output=None, share=None). The
+            // recovered epoch must override these.
+            let (actor, _mailbox) = Actor::<_, _, Sha256, _, MinSig>::new(
+                context.with_label("actor"),
+                Config {
+                    manager: NoopManager::<Ed25519PublicKey>::default(),
+                    signer,
+                    mailbox_size: 8,
+                    partition_prefix,
+                    peer_config: peer_config.clone(),
+                    max_supported_mode: crate::dkg::MAX_SUPPORTED_MODE,
+                },
+            );
+            let (sender, receiver) = inert_channel(&peer_config.participants);
+            let (orchestrator_sender, mut orchestrator_receiver) = mpsc::channel(4);
+            actor.start(
+                None,
+                None,
+                orchestrator::Mailbox::new(orchestrator_sender),
+                (sender, receiver),
+                ContinueOnUpdate::boxed(),
+            );
+
+            // The first epoch transition the actor emits should describe the recovered reshare
+            // round. Under the bug, `is_dkg` was computed from the `None` startup output and the
+            // actor re-entered the bootstrap DKG path, producing a transition with all
+            // participants as dealers and an empty poly.
+            let Some(Message::Enter(transition)) = orchestrator_receiver.recv().await else {
+                panic!("actor should emit an epoch transition");
+            };
+            assert_eq!(transition.epoch, Epoch::new(RECOVERED_EPOCH));
+            assert!(
+                transition.poly.is_some(),
+                "transition should carry the recovered public polynomial",
+            );
+            assert_eq!(transition.dealers, peer_config.dealers(RECOVERED_ROUND));
+        });
+    }
+}

examples/reshare/src/orchestrator/mod.rs

@@ -4,4 +4,6 @@ mod actor;
 pub use actor::{Actor, Config};
 
 mod ingress;
+#[cfg(test)]
+pub use ingress::Message;
 pub use ingress::{EpochTransition, Mailbox};