Skip to main content

commonware_glue/dkg/
types.rs

1//! Shared types for the DKG module.
2
3use crate::dkg::network::Directory;
4use bytes::{Buf, BufMut};
5use commonware_codec::{EncodeSize, Error as CodecError, RangeCfg, Read, ReadExt, Write};
6use commonware_consensus::types::Epoch;
7use commonware_cryptography::{
8    PublicKey, Signer,
9    bls12381::{
10        dkg::feldman_desmedt::{DealerPrivMsg, DealerPubMsg, Output, PlayerAck, SignedDealerLog},
11        primitives::{
12            group::Share,
13            sharing::{ModeVersion, Sharing},
14            variant::Variant,
15        },
16    },
17};
18use commonware_p2p::TrackedPeers;
19use commonware_utils::{Faults as _, N3f1, ordered::Set, sequence::Unit};
20use std::num::{NonZeroU32, NonZeroU64};
21use thiserror::Error;
22
23/// Information required to construct an epoch-scoped threshold scheme that may
24/// or may not be capable of signing messages.
25#[derive(Clone, Debug, PartialEq, Eq)]
26pub enum SchemeInfo<V: Variant, P: PublicKey> {
27    /// Information required for constructing a verifier scheme.
28    Verifier {
29        /// The participants.
30        participants: Set<P>,
31        /// The public group polynomial.
32        sharing: Sharing<V>,
33    },
34    /// Information required for constructing a signer scheme.
35    Signer {
36        /// The participants.
37        participants: Set<P>,
38        /// The public group polynomial.
39        sharing: Sharing<V>,
40        /// A BLS [`Share`].
41        share: Share,
42    },
43}
44
45/// Result of a completed DKG/reshare epoch.
46#[derive(Clone, Copy, Debug, PartialEq, Eq)]
47#[cfg_attr(feature = "arbitrary", derive(arbitrary::Arbitrary))]
48pub enum EpochOutcome {
49    /// The epoch produced a new public output.
50    Success,
51    /// The epoch failed and carried the previous public state forward.
52    Failure,
53}
54
55impl Write for EpochOutcome {
56    fn write(&self, writer: &mut impl BufMut) {
57        let tag = match self {
58            Self::Success => 0u8,
59            Self::Failure => 1u8,
60        };
61        tag.write(writer);
62    }
63}
64
65impl EncodeSize for EpochOutcome {
66    fn encode_size(&self) -> usize {
67        1
68    }
69}
70
71impl Read for EpochOutcome {
72    type Cfg = ();
73
74    fn read_cfg(reader: &mut impl Buf, _: &Self::Cfg) -> Result<Self, CodecError> {
75        match u8::read(reader)? {
76            0 => Ok(Self::Success),
77            1 => Ok(Self::Failure),
78            n => Err(CodecError::InvalidEnum(n)),
79        }
80    }
81}
82
83/// Participants for a DKG/reshare epoch.
84#[derive(Clone, Debug, PartialEq, Eq)]
85pub struct Participants<P: PublicKey> {
86    /// Peers that distribute dealings in this epoch.
87    pub dealers: Set<P>,
88    /// Peers that receive shares in this epoch.
89    pub players: Set<P>,
90    /// Players of the next epoch, tracked early for connectivity.
91    pub next_players: Set<P>,
92}
93
94/// Errors produced while validating DKG/reshare participants.
95#[derive(Debug, Error, PartialEq, Eq)]
96pub enum ParticipantsError {
97    /// No dealers were provided.
98    #[error("dealers must not be empty")]
99    EmptyDealers,
100    /// No players were provided.
101    #[error("players must not be empty")]
102    EmptyPlayers,
103    /// A participant set exceeds the configured maximum.
104    #[error("too many participants: {actual} > {max}")]
105    TooManyParticipants { actual: usize, max: usize },
106    /// Round-zero reshare dealers differ from the previous output players.
107    #[error("round-zero reshare dealers must equal previous output players")]
108    InitialReshareDealers,
109    /// A later reshare dealer does not own a previous share.
110    #[error("reshare dealer is not a previous player")]
111    UnknownReshareDealer,
112}
113
114#[derive(Debug, PartialEq, Eq)]
115pub(crate) struct EpochCapacityError {
116    available: u64,
117    required: u64,
118}
119
120impl<P: PublicKey> Participants<P> {
121    /// Builds the peer set used by the DKG channel.
122    ///
123    /// Dealers are the primary tracked peers because they send protocol data in the
124    /// current round. Current and next players are tracked as secondary peers so the
125    /// actor keeps enough connectivity to receive its own messages and prepare the
126    /// next epoch without allowing next players to act as dealers.
127    pub fn tracked_peers(&self) -> TrackedPeers<P> {
128        TrackedPeers::new(
129            self.dealers.clone(),
130            Set::from_iter_dedup(self.players.iter().chain(self.next_players.iter()).cloned()),
131        )
132    }
133
134    /// Checks that a participant snapshot is usable for the requested reshare round.
135    ///
136    /// Reshare requires non-empty dealer and player sets, caps every participant set
137    /// at `max_participants`, and verifies that reshare dealers are authorized by the
138    /// previous epoch output. Round zero must start from exactly the previous player
139    /// set. Later rounds may use any subset of previous players as dealers.
140    pub fn validate<V: Variant>(
141        &self,
142        max_participants: NonZeroU32,
143        previous: Option<&Output<V, P>>,
144        round: u64,
145    ) -> Result<(), ParticipantsError> {
146        if self.dealers.is_empty() {
147            return Err(ParticipantsError::EmptyDealers);
148        }
149        if self.players.is_empty() {
150            return Err(ParticipantsError::EmptyPlayers);
151        }
152
153        let max = max_participants.get() as usize;
154        for actual in [
155            self.dealers.len(),
156            self.players.len(),
157            self.next_players.len(),
158        ] {
159            if actual > max {
160                return Err(ParticipantsError::TooManyParticipants { actual, max });
161            }
162        }
163
164        let Some(previous) = previous else {
165            return Ok(());
166        };
167
168        if round == 0 {
169            if &self.dealers != previous.players() {
170                return Err(ParticipantsError::InitialReshareDealers);
171            }
172            return Ok(());
173        }
174
175        if self
176            .dealers
177            .iter()
178            .any(|dealer| previous.players().position(dealer).is_none())
179        {
180            return Err(ParticipantsError::UnknownReshareDealer);
181        }
182
183        Ok(())
184    }
185
186    pub(crate) fn validate_epoch_capacity<V: Variant>(
187        &self,
188        blocks_per_epoch: NonZeroU64,
189        previous: Option<&Output<V, P>>,
190    ) -> Result<(), EpochCapacityError> {
191        let available = dealer_log_slots(blocks_per_epoch);
192        let required = self.required_dealer_logs(previous);
193        if available < required {
194            return Err(EpochCapacityError {
195                available,
196                required,
197            });
198        }
199        Ok(())
200    }
201
202    fn required_dealer_logs<V: Variant>(&self, previous: Option<&Output<V, P>>) -> u64 {
203        let dealer_quorum = u64::from(N3f1::quorum(self.dealers.len()));
204        let previous_quorum = previous
205            .map(|previous| u64::from(previous.quorum::<N3f1>()))
206            .unwrap_or_default();
207        dealer_quorum.max(previous_quorum)
208    }
209}
210
211const fn dealer_log_slots(blocks_per_epoch: NonZeroU64) -> u64 {
212    let blocks = blocks_per_epoch.get();
213    // Shorter epochs do not leave a usable dealing-to-inclusion window.
214    if blocks < 4 {
215        return 0;
216    }
217    blocks.saturating_sub(blocks / 2 + 1)
218}
219
220impl<P: PublicKey> Write for Participants<P> {
221    fn write(&self, writer: &mut impl BufMut) {
222        self.dealers.write(writer);
223        self.players.write(writer);
224        self.next_players.write(writer);
225    }
226}
227
228impl<P: PublicKey> EncodeSize for Participants<P> {
229    fn encode_size(&self) -> usize {
230        self.dealers.encode_size() + self.players.encode_size() + self.next_players.encode_size()
231    }
232}
233
234impl<P: PublicKey> Read for Participants<P> {
235    /// Maximum number of participants accepted in any single set.
236    type Cfg = NonZeroU32;
237
238    fn read_cfg(reader: &mut impl Buf, max: &Self::Cfg) -> Result<Self, CodecError> {
239        let cfg = (RangeCfg::new(0..=max.get() as usize), ());
240        Ok(Self {
241            dealers: Set::read_cfg(reader, &cfg)?,
242            players: Set::read_cfg(reader, &cfg)?,
243            next_players: Set::read_cfg(reader, &cfg)?,
244        })
245    }
246}
247
248#[cfg(test)]
249mod tests {
250    use super::*;
251    use crate::dkg::network::Addresses;
252    use commonware_codec::{Decode as _, Encode as _};
253    use commonware_cryptography::{
254        bls12381::{
255            dkg::feldman_desmedt::deal,
256            primitives::{sharing::Mode, variant::MinPk},
257        },
258        ed25519,
259    };
260    use commonware_p2p::Address;
261    use commonware_utils::{NZU32, NZU64, TestRng};
262    use std::net::{IpAddr, Ipv4Addr, SocketAddr};
263
264    fn keys(count: u64) -> Set<ed25519::PublicKey> {
265        Set::from_iter_dedup(
266            (0..count).map(|seed| ed25519::PrivateKey::from_seed(seed).public_key()),
267        )
268    }
269
270    fn addresses(keys: &Set<ed25519::PublicKey>) -> Addresses<ed25519::PublicKey> {
271        keys.iter()
272            .enumerate()
273            .map(|(index, key)| {
274                let socket = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), index as u16 + 1);
275                (key.clone(), Address::Symmetric(socket))
276            })
277            .collect()
278    }
279
280    fn addressed_info(
281        participants: u64,
282    ) -> EpochInfo<MinPk, ed25519::PublicKey, Addresses<ed25519::PublicKey>> {
283        let keys = keys(participants);
284        let (output, _) =
285            deal::<MinPk, _, N3f1>(TestRng::new(1), Mode::NonZeroCounter, keys.clone())
286                .expect("trusted deal");
287        EpochInfo {
288            outcome: EpochOutcome::Success,
289            epoch: Epoch::new(2),
290            output,
291            players: keys.clone(),
292            next_players: keys.clone(),
293            directory: addresses(&keys),
294        }
295    }
296
297    #[test]
298    fn addressed_epoch_info_roundtrips() {
299        let info = addressed_info(4);
300        let decoded =
301            EpochInfo::<MinPk, ed25519::PublicKey, Addresses<ed25519::PublicKey>>::decode_cfg(
302                info.encode(),
303                &(NZU32!(4), crate::dkg::tests::max_supported_mode()),
304            )
305            .expect("decode addressed epoch info");
306        assert_eq!(decoded, info);
307    }
308
309    #[test]
310    fn addressed_epoch_info_roundtrips_disjoint_participant_sets() {
311        let dealers = keys(4);
312        let players = Set::from_iter_dedup(
313            (4..8).map(|seed| ed25519::PrivateKey::from_seed(seed).public_key()),
314        );
315        let next_players = Set::from_iter_dedup(
316            (8..12).map(|seed| ed25519::PrivateKey::from_seed(seed).public_key()),
317        );
318        let peers = Set::from_iter_dedup(
319            dealers
320                .iter()
321                .chain(players.iter())
322                .chain(next_players.iter())
323                .cloned(),
324        );
325        let (output, _) = deal::<MinPk, _, N3f1>(TestRng::new(1), Mode::NonZeroCounter, dealers)
326            .expect("trusted deal");
327        let info = EpochInfo {
328            outcome: EpochOutcome::Success,
329            epoch: Epoch::new(2),
330            output,
331            players,
332            next_players,
333            directory: addresses(&peers),
334        };
335
336        let decoded =
337            EpochInfo::<MinPk, ed25519::PublicKey, Addresses<ed25519::PublicKey>>::decode_cfg(
338                info.encode(),
339                &(NZU32!(4), crate::dkg::tests::max_supported_mode()),
340            )
341            .expect("decode addressed epoch info");
342        assert_eq!(decoded, info);
343    }
344
345    #[test]
346    fn addressed_epoch_info_rejects_extra_directory_peer() {
347        let mut info = addressed_info(1);
348        info.directory = addresses(&keys(2));
349        assert!(
350            EpochInfo::<MinPk, ed25519::PublicKey, Addresses<ed25519::PublicKey>>::decode_cfg(
351                info.encode(),
352                &(NZU32!(1), crate::dkg::tests::max_supported_mode()),
353            )
354            .is_err()
355        );
356    }
357
358    #[test]
359    fn addressed_epoch_info_rejects_missing_directory_peer() {
360        let mut info = addressed_info(1);
361        info.directory = addresses(&keys(0));
362        assert!(
363            EpochInfo::<MinPk, ed25519::PublicKey, Addresses<ed25519::PublicKey>>::decode_cfg(
364                info.encode(),
365                &(NZU32!(1), crate::dkg::tests::max_supported_mode()),
366            )
367            .is_err()
368        );
369    }
370
371    #[test]
372    fn addressed_epoch_info_rejects_wrong_directory_peer() {
373        let mut info = addressed_info(1);
374        let wrong = Set::from_iter_dedup([ed25519::PrivateKey::from_seed(1).public_key()]);
375        info.directory = addresses(&wrong);
376        assert!(
377            EpochInfo::<MinPk, ed25519::PublicKey, Addresses<ed25519::PublicKey>>::decode_cfg(
378                info.encode(),
379                &(NZU32!(1), crate::dkg::tests::max_supported_mode()),
380            )
381            .is_err()
382        );
383    }
384
385    fn participants(count: u64) -> Participants<ed25519::PublicKey> {
386        let keys = keys(count);
387        Participants {
388            dealers: keys.clone(),
389            players: keys,
390            next_players: Set::default(),
391        }
392    }
393
394    #[test]
395    fn epoch_capacity_rejects_insufficient_bootstrap_slots() {
396        assert_eq!(
397            participants(2).validate_epoch_capacity::<MinPk>(NZU64!(4), None),
398            Err(EpochCapacityError {
399                available: 1,
400                required: 2,
401            })
402        );
403    }
404
405    #[test]
406    fn epoch_capacity_accepts_exact_bootstrap_slots() {
407        assert!(
408            participants(2)
409                .validate_epoch_capacity::<MinPk>(NZU64!(5), None)
410                .is_ok()
411        );
412    }
413
414    #[test]
415    fn epoch_capacity_rejects_short_epoch_with_single_dealer() {
416        assert_eq!(
417            participants(1).validate_epoch_capacity::<MinPk>(NZU64!(3), None),
418            Err(EpochCapacityError {
419                available: 0,
420                required: 1,
421            })
422        );
423    }
424
425    #[test]
426    fn epoch_capacity_rejects_insufficient_reshare_slots() {
427        let players = keys(4);
428        let (previous, _) =
429            deal::<MinPk, _, N3f1>(TestRng::new(0), Mode::NonZeroCounter, players.clone())
430                .expect("trusted deal");
431
432        assert_eq!(
433            Participants {
434                dealers: players.clone(),
435                players,
436                next_players: Set::default(),
437            }
438            .validate_epoch_capacity(NZU64!(6), Some(&previous)),
439            Err(EpochCapacityError {
440                available: 2,
441                required: 3,
442            })
443        );
444    }
445
446    #[test]
447    fn epoch_capacity_accepts_exact_reshare_slots() {
448        let players = keys(4);
449        let (previous, _) =
450            deal::<MinPk, _, N3f1>(TestRng::new(1), Mode::NonZeroCounter, players.clone())
451                .expect("trusted deal");
452
453        assert!(
454            Participants {
455                dealers: players.clone(),
456                players,
457                next_players: Set::default(),
458            }
459            .validate_epoch_capacity(NZU64!(7), Some(&previous))
460            .is_ok()
461        );
462    }
463}
464
465#[cfg(feature = "arbitrary")]
466impl<P: PublicKey> arbitrary::Arbitrary<'_> for Participants<P>
467where
468    P: for<'a> arbitrary::Arbitrary<'a>,
469{
470    fn arbitrary(u: &mut arbitrary::Unstructured<'_>) -> arbitrary::Result<Self> {
471        Ok(Self {
472            dealers: u.arbitrary()?,
473            players: u.arbitrary()?,
474            next_players: u.arbitrary()?,
475        })
476    }
477}
478
479/// Canonical public epoch artifact.
480///
481/// This is the public truth needed to start an epoch: the latest public output,
482/// the participant sets not already carried by that output, and the transport
483/// [`Directory`] for every peer of the epoch. The genesis block carries the
484/// [`EpochInfo`] for epoch 0; the final block of each epoch carries the
485/// [`EpochInfo`] for the following epoch. The reshare actor never invents this;
486/// it reads it back from finalized block ancestry.
487///
488/// Because the directory rides in the artifact, a node recovering through a
489/// certificate-backed route (a finalized boundary block, a
490/// [`probe`](crate::dkg::probe) artifact, or persisted
491/// [`state_sync`](crate::dkg::state_sync) material) can activate the epoch's
492/// peers without access to application state.
493#[derive(Clone, Debug, PartialEq, Eq)]
494pub struct EpochInfo<V: Variant, P: PublicKey, D: Directory<P> = Unit> {
495    /// Whether or not the reshare ceremony in this epoch was successful.
496    pub outcome: EpochOutcome,
497    /// Epoch this artifact describes.
498    pub epoch: Epoch,
499    /// Latest public DKG output.
500    pub output: Output<V, P>,
501    /// Peers that receive shares in this epoch.
502    pub players: Set<P>,
503    /// Players of the next epoch, tracked early for connectivity.
504    pub next_players: Set<P>,
505    /// Transport directory containing exactly this epoch's dealers, players,
506    /// and next players.
507    pub directory: D,
508}
509
510impl<V: Variant, P: PublicKey, D: Directory<P>> EpochInfo<V, P, D> {
511    /// Reconstructs the complete participant snapshot for this epoch.
512    pub fn participants(&self) -> Participants<P> {
513        Participants {
514            dealers: self.output.players().clone(),
515            players: self.players.clone(),
516            next_players: self.next_players.clone(),
517        }
518    }
519}
520
521impl<V: Variant, P: PublicKey, D: Directory<P>> Write for EpochInfo<V, P, D> {
522    fn write(&self, buf: &mut impl BufMut) {
523        self.outcome.write(buf);
524        self.epoch.write(buf);
525        self.output.write(buf);
526        self.players.write(buf);
527        self.next_players.write(buf);
528        self.directory.write(buf);
529    }
530}
531
532impl<V: Variant, P: PublicKey, D: Directory<P>> EncodeSize for EpochInfo<V, P, D> {
533    fn encode_size(&self) -> usize {
534        self.outcome.encode_size()
535            + self.epoch.encode_size()
536            + self.output.encode_size()
537            + self.players.encode_size()
538            + self.next_players.encode_size()
539            + self.directory.encode_size()
540    }
541}
542
543impl<V: Variant, P: PublicKey, D: Directory<P>> Read for EpochInfo<V, P, D> {
544    /// Maximum entries accepted in each participant set and maximum supported
545    /// sharing mode version.
546    type Cfg = (NonZeroU32, ModeVersion);
547
548    fn read_cfg(
549        buf: &mut impl Buf,
550        (max_participants, max_supported_mode): &Self::Cfg,
551    ) -> Result<Self, CodecError> {
552        let outcome = EpochOutcome::read(buf)?;
553        let epoch = Epoch::read(buf)?;
554        let output = Output::<V, P>::read_cfg(buf, &(*max_participants, *max_supported_mode))?;
555        let players = Set::read_cfg(
556            buf,
557            &(RangeCfg::new(0..=max_participants.get() as usize), ()),
558        )?;
559        let next_players = Set::read_cfg(
560            buf,
561            &(RangeCfg::new(0..=max_participants.get() as usize), ()),
562        )?;
563        let peers = Set::from_iter_dedup(
564            output
565                .players()
566                .iter()
567                .chain(players.iter())
568                .chain(next_players.iter())
569                .cloned(),
570        );
571        let directory = D::read_cfg(buf, &D::codec_config(&peers))?;
572        if !directory.matches(&peers) {
573            return Err(CodecError::Invalid(
574                "EpochInfo",
575                "directory does not match participants",
576            ));
577        }
578        Ok(Self {
579            outcome,
580            epoch,
581            output,
582            players,
583            next_players,
584            directory,
585        })
586    }
587}
588
589#[cfg(feature = "arbitrary")]
590impl<V: Variant, P: PublicKey, D: Directory<P>> arbitrary::Arbitrary<'_> for EpochInfo<V, P, D>
591where
592    P: for<'a> arbitrary::Arbitrary<'a>,
593    D: for<'a> arbitrary::Arbitrary<'a>,
594    Output<V, P>: for<'a> arbitrary::Arbitrary<'a>,
595{
596    fn arbitrary(u: &mut arbitrary::Unstructured<'_>) -> arbitrary::Result<Self> {
597        Ok(Self {
598            outcome: u.arbitrary()?,
599            epoch: u.arbitrary()?,
600            output: u.arbitrary()?,
601            players: u.arbitrary()?,
602            next_players: u.arbitrary()?,
603            directory: u.arbitrary()?,
604        })
605    }
606}
607
608/// A public artifact published by the reshare actor into a block.
609///
610/// During the dealing and inclusion window of an epoch the actor publishes
611/// finalized dealer logs. The final block of an epoch instead carries the
612/// canonical [`EpochInfo`] for the following epoch.
613#[allow(clippy::large_enum_variant)]
614pub enum Payload<V: Variant, C: Signer, D: Directory<C::PublicKey> = Unit> {
615    /// A finalized signed dealer log for inclusion mid-epoch.
616    DealerLog(SignedDealerLog<V, C>),
617    /// The canonical public epoch artifact for the next epoch, carried by the
618    /// final block of the current epoch.
619    EpochInfo(EpochInfo<V, C::PublicKey, D>),
620}
621
622impl<V: Variant, C: Signer, D: Directory<C::PublicKey>> Clone for Payload<V, C, D> {
623    fn clone(&self) -> Self {
624        match self {
625            Self::DealerLog(log) => Self::DealerLog(log.clone()),
626            Self::EpochInfo(info) => Self::EpochInfo(info.clone()),
627        }
628    }
629}
630
631impl<V: Variant, C: Signer, D: Directory<C::PublicKey>> PartialEq for Payload<V, C, D> {
632    fn eq(&self, other: &Self) -> bool {
633        match (self, other) {
634            (Self::DealerLog(a), Self::DealerLog(b)) => a == b,
635            (Self::EpochInfo(a), Self::EpochInfo(b)) => a == b,
636            _ => false,
637        }
638    }
639}
640
641impl<V: Variant, C: Signer, D: Directory<C::PublicKey>> Eq for Payload<V, C, D> {}
642
643impl<V: Variant, C: Signer, D: Directory<C::PublicKey>> Write for Payload<V, C, D> {
644    fn write(&self, writer: &mut impl BufMut) {
645        match self {
646            Self::DealerLog(log) => {
647                0u8.write(writer);
648                log.write(writer);
649            }
650            Self::EpochInfo(info) => {
651                1u8.write(writer);
652                info.write(writer);
653            }
654        }
655    }
656}
657
658impl<V: Variant, C: Signer, D: Directory<C::PublicKey>> EncodeSize for Payload<V, C, D> {
659    fn encode_size(&self) -> usize {
660        1 + match self {
661            Self::DealerLog(log) => log.encode_size(),
662            Self::EpochInfo(info) => info.encode_size(),
663        }
664    }
665}
666
667impl<V: Variant, C: Signer, D: Directory<C::PublicKey>> Read for Payload<V, C, D> {
668    /// Maximum entries accepted in each participant set and maximum supported
669    /// sharing mode version.
670    type Cfg = (NonZeroU32, ModeVersion);
671
672    fn read_cfg(reader: &mut impl Buf, cfg: &Self::Cfg) -> Result<Self, CodecError> {
673        match u8::read(reader)? {
674            0 => Ok(Self::DealerLog(SignedDealerLog::read_cfg(reader, &cfg.0)?)),
675            1 => Ok(Self::EpochInfo(EpochInfo::read_cfg(reader, cfg)?)),
676            n => Err(CodecError::InvalidEnum(n)),
677        }
678    }
679}
680
681#[cfg(feature = "arbitrary")]
682impl<V: Variant, C: Signer, D: Directory<C::PublicKey>> arbitrary::Arbitrary<'_>
683    for Payload<V, C, D>
684where
685    SignedDealerLog<V, C>: for<'a> arbitrary::Arbitrary<'a>,
686    EpochInfo<V, C::PublicKey, D>: for<'a> arbitrary::Arbitrary<'a>,
687{
688    fn arbitrary(u: &mut arbitrary::Unstructured<'_>) -> arbitrary::Result<Self> {
689        Ok(if u.arbitrary::<bool>()? {
690            Self::DealerLog(u.arbitrary()?)
691        } else {
692            Self::EpochInfo(u.arbitrary()?)
693        })
694    }
695}
696
697/// Wire message type for DKG protocol communication.
698pub enum Message<V: Variant, P: PublicKey> {
699    /// A dealer message containing public and private components for a player.
700    Dealer(DealerPubMsg<V>, DealerPrivMsg),
701    /// A player acknowledgment sent back to a dealer.
702    Ack(PlayerAck<P>),
703}
704
705impl<V: Variant, P: PublicKey> Write for Message<V, P> {
706    fn write(&self, writer: &mut impl BufMut) {
707        match self {
708            Self::Dealer(pub_msg, priv_msg) => {
709                0u8.write(writer);
710                pub_msg.write(writer);
711                priv_msg.write(writer);
712            }
713            Self::Ack(ack) => {
714                1u8.write(writer);
715                ack.write(writer);
716            }
717        }
718    }
719}
720
721impl<V: Variant, P: PublicKey> EncodeSize for Message<V, P> {
722    fn encode_size(&self) -> usize {
723        1 + match self {
724            Self::Dealer(pub_msg, priv_msg) => pub_msg.encode_size() + priv_msg.encode_size(),
725            Self::Ack(ack) => ack.encode_size(),
726        }
727    }
728}
729
730impl<V: Variant, P: PublicKey> Read for Message<V, P> {
731    type Cfg = NonZeroU32;
732
733    fn read_cfg(reader: &mut impl Buf, cfg: &Self::Cfg) -> Result<Self, CodecError> {
734        let tag = u8::read(reader)?;
735        match tag {
736            0 => {
737                let pub_msg = DealerPubMsg::read_cfg(reader, cfg)?;
738                let priv_msg = DealerPrivMsg::read(reader)?;
739                Ok(Self::Dealer(pub_msg, priv_msg))
740            }
741            1 => {
742                let ack = PlayerAck::read(reader)?;
743                Ok(Self::Ack(ack))
744            }
745            n => Err(CodecError::InvalidEnum(n)),
746        }
747    }
748}
749
750#[cfg(feature = "arbitrary")]
751impl<V: Variant, P: PublicKey> arbitrary::Arbitrary<'_> for Message<V, P>
752where
753    DealerPubMsg<V>: for<'a> arbitrary::Arbitrary<'a>,
754    PlayerAck<P>: for<'a> arbitrary::Arbitrary<'a>,
755{
756    fn arbitrary(u: &mut arbitrary::Unstructured<'_>) -> arbitrary::Result<Self> {
757        Ok(if u.arbitrary::<bool>()? {
758            Self::Dealer(u.arbitrary()?, u.arbitrary()?)
759        } else {
760            Self::Ack(u.arbitrary()?)
761        })
762    }
763}
764
765#[cfg(all(test, feature = "arbitrary"))]
766mod conformance {
767    use super::*;
768    use crate::dkg::network::Addresses;
769    use commonware_codec::conformance::CodecConformance;
770    use commonware_cryptography::{bls12381::primitives::variant::MinSig, ed25519};
771
772    commonware_conformance::conformance_tests! {
773        CodecConformance<EpochOutcome>,
774        CodecConformance<Participants<ed25519::PublicKey>>,
775        CodecConformance<EpochInfo<MinSig, ed25519::PublicKey>> => 8192,
776        CodecConformance<EpochInfo<MinSig, ed25519::PublicKey, Addresses<ed25519::PublicKey>>> => 8192,
777        CodecConformance<Payload<MinSig, ed25519::PrivateKey>> => 8192,
778        CodecConformance<Payload<MinSig, ed25519::PrivateKey, Addresses<ed25519::PublicKey>>> => 8192,
779        CodecConformance<Message<MinSig, ed25519::PublicKey>>,
780    }
781}