#![allow(dead_code)]
use crate::dkg::{
ParticipantsProvider, Registrar, ReshareBlock, SecretStore,
network::{Addresses, Directory as DkgDirectory, Manager as DkgManager},
orchestrator, reshare,
types::{Payload, SchemeInfo},
};
use bytes::{Buf, BufMut};
use commonware_actor::Feedback;
use commonware_codec::{
Codec, Decode, Encode, EncodeSize, Error as CodecError, Read, ReadExt, Write, varint::UInt,
};
use commonware_consensus::{
Automaton, Block, CertifiableAutomaton, Heightable, Relay, Reporter,
marshal::{
self, Start as MarshalStart, Update,
core::{Actor as MarshalActor, Mailbox as MarshalMailbox},
standard::Standard,
},
simplex::{
self, ForwardPolicy, Plan, SkipPolicy, elector::RoundRobin, mocks::scheme, types::Context,
},
types::{Epoch, FixedEpocher, Height, Round, View, ViewDelta},
};
use commonware_cryptography::{
Digest, Digestible, Hasher, PublicKey as CryptoPublicKey, Signer,
bls12381::{
dkg::feldman_desmedt::DealerPrivMsg,
primitives::{
group::Share,
variant::{MinPk, Variant},
},
},
certificate::{ConstantProvider, Verifier as _},
ed25519::{PrivateKey, PublicKey},
sha256::{Digest as Sha256Digest, Sha256},
transcript::Summary,
};
use commonware_p2p::{
Message as P2pMessage, Provider, Receiver, TrackedPeers,
simulated::{Control, Manager as SimManager},
utils::mux,
};
use commonware_parallel::Sequential;
use commonware_runtime::{Supervisor as _, buffer::paged::CacheRef, deterministic};
use commonware_storage::archive::immutable;
use commonware_utils::{
Acknowledgement, NZU16, NZU64, NZUsize,
acknowledgement::Exact,
channel::{fallible::OneshotExt, oneshot},
ordered::Set,
sequence::Unit,
sync::Mutex,
};
use std::{
collections::{BTreeMap, HashSet},
marker::PhantomData,
num::{NonZeroU32, NonZeroU64},
sync::Arc,
time::Duration,
};
pub(crate) type TestDigest = Sha256Digest;
pub(crate) type TestPublicKey = PublicKey;
pub(crate) type TestSigner = PrivateKey;
pub(crate) type TestContext = Context<TestDigest, TestPublicKey>;
pub(crate) type TestBlock = MockBlock<TestDigest, TestContext>;
pub(crate) type TestMarshalVariant = Standard<TestBlock>;
pub(crate) type TestBlsVariant = MinPk;
pub(crate) type TestScheme = scheme::Scheme<TestPublicKey>;
pub(crate) type TestProvider = ConstantProvider<TestScheme, Epoch>;
pub(crate) type TestElector = RoundRobin;
pub(crate) type TestStrategy = Sequential;
pub(crate) type TestBlocker = Control<TestPublicKey, deterministic::Context>;
pub(crate) type TestManager = SimManager<TestPublicKey, deterministic::Context>;
pub(crate) type TestMailbox = orchestrator::Mailbox<TestBlock>;
pub(crate) type TestMarshalMailbox = MarshalMailbox<TestScheme, TestMarshalVariant>;
#[derive(Clone, Copy, Debug, thiserror::Error)]
#[error("peer set unavailable")]
pub(crate) struct TrackFailed;
#[derive(Clone, Debug)]
pub(crate) struct FailingManager<M>(pub(crate) M);
impl<M: Provider> Provider for FailingManager<M> {
type PublicKey = M::PublicKey;
async fn peer_set(&mut self, id: u64) -> Option<TrackedPeers<Self::PublicKey>> {
self.0.peer_set(id).await
}
async fn subscribe(&mut self) -> commonware_p2p::PeerSetSubscription<Self::PublicKey> {
self.0.subscribe().await
}
}
impl<M: Provider> DkgManager for FailingManager<M> {
type Directory = Unit;
type Error = TrackFailed;
fn track(
&mut self,
_epoch: Epoch,
_peers: TrackedPeers<Self::PublicKey>,
_directory: &Self::Directory,
) -> Result<(), Self::Error> {
Err(TrackFailed)
}
}
type DirectoryTracks = Arc<Mutex<Vec<(Epoch, TrackedPeers<PublicKey>, Addresses<PublicKey>)>>>;
#[derive(Clone, Debug)]
pub(crate) struct DirectoryManager<M> {
inner: M,
tracked: DirectoryTracks,
}
impl<M> DirectoryManager<M> {
pub(crate) fn new(inner: M) -> Self {
Self {
inner,
tracked: Arc::default(),
}
}
pub(crate) fn tracked(&self) -> Vec<(Epoch, TrackedPeers<PublicKey>, Addresses<PublicKey>)> {
self.tracked.lock().clone()
}
}
impl<M: Provider<PublicKey = PublicKey>> Provider for DirectoryManager<M> {
type PublicKey = PublicKey;
async fn peer_set(&mut self, id: u64) -> Option<TrackedPeers<Self::PublicKey>> {
self.inner.peer_set(id).await
}
async fn subscribe(&mut self) -> commonware_p2p::PeerSetSubscription<Self::PublicKey> {
self.inner.subscribe().await
}
}
impl<M> DkgManager for DirectoryManager<M>
where
M: DkgManager<PublicKey = PublicKey, Directory = Unit>,
{
type Directory = Addresses<PublicKey>;
type Error = M::Error;
fn track(
&mut self,
epoch: Epoch,
peers: TrackedPeers<Self::PublicKey>,
directory: &Self::Directory,
) -> Result<(), Self::Error> {
self.tracked
.lock()
.push((epoch, peers.clone(), directory.clone()));
self.inner.track(epoch, peers, &Unit)
}
}
pub(crate) type TestActor = orchestrator::Actor<
deterministic::Context,
TestBlocker,
TestManager,
TestProvider,
TestMarshalVariant,
TestBlsVariant,
TestSigner,
MockApplication,
TestElector,
TestStrategy,
>;
pub(crate) type TestReshareActor = reshare::Actor<
deterministic::Context,
TestBlock,
TestBlsVariant,
TestSigner,
TestManager,
TestBlocker,
StaticParticipants,
MemorySecretStore,
Sequential,
commonware_cryptography::ed25519::Batch,
TestScheme,
TestMarshalVariant,
MockConsumer,
>;
#[derive(Clone)]
pub(crate) struct StaticParticipants(pub(crate) Set<TestPublicKey>);
impl ParticipantsProvider for StaticParticipants {
type PublicKey = TestPublicKey;
type Directory = Unit;
async fn participants(&mut self, _epoch: Epoch) -> Set<Self::PublicKey> {
self.0.clone()
}
async fn directory(&mut self, _: Epoch, _: Set<Self::PublicKey>) -> Self::Directory {
Unit
}
}
const NAMESPACE: &[u8] = b"_COMMONWARE_GLUE_DKG_ORCHESTRATOR_TEST";
#[derive(Debug)]
pub(crate) struct FilteredReceiver<R> {
inner: R,
filter: Filter,
}
#[derive(Debug)]
enum Filter {
None,
All,
Epochs(Arc<HashSet<u64>>),
}
impl<R> FilteredReceiver<R> {
pub(crate) const fn pass(inner: R) -> Self {
Self {
inner,
filter: Filter::None,
}
}
pub(crate) const fn drop_all(inner: R) -> Self {
Self {
inner,
filter: Filter::All,
}
}
pub(crate) const fn epochs(inner: R, epochs: Arc<HashSet<u64>>) -> Self {
Self {
inner,
filter: Filter::Epochs(epochs),
}
}
}
impl<R: Receiver> Receiver for FilteredReceiver<R> {
type Error = R::Error;
type PublicKey = R::PublicKey;
async fn recv(&mut self) -> Result<P2pMessage<Self::PublicKey>, Self::Error> {
loop {
let message = self.inner.recv().await?;
match &self.filter {
Filter::None => return Ok(message),
Filter::All => {}
Filter::Epochs(epochs) => {
let (_, bytes) = &message;
let (epoch, _) =
mux::parse(bytes.clone()).expect("failed to parse mux message");
if !epochs.contains(&epoch) {
return Ok(message);
}
}
}
}
}
}
#[derive(Clone, Debug, PartialEq, Eq)]
#[cfg_attr(feature = "arbitrary", derive(arbitrary::Arbitrary))]
pub(crate) struct MockBlock<D: Digest, C, Dir = Unit> {
context: C,
parent: D,
height: Height,
timestamp: u64,
payload: Option<EncodedPayload>,
digest: D,
_directory: PhantomData<Dir>,
}
#[derive(Clone, Debug, PartialEq, Eq)]
#[cfg_attr(feature = "arbitrary", derive(arbitrary::Arbitrary))]
pub(crate) struct EncodedPayload {
max_participants: NonZeroU32,
bytes: Vec<u8>,
}
impl EncodedPayload {
pub(crate) fn new<V, S, Dir>(max_participants: NonZeroU32, payload: Payload<V, S, Dir>) -> Self
where
V: Variant,
S: Signer,
Dir: DkgDirectory<S::PublicKey>,
{
Self {
max_participants,
bytes: payload.encode().to_vec(),
}
}
fn decode<V, S, Dir>(&self) -> Option<Payload<V, S, Dir>>
where
V: Variant,
S: Signer,
Dir: DkgDirectory<S::PublicKey>,
{
Payload::decode_cfg(
self.bytes.as_slice(),
&(
self.max_participants,
crate::dkg::tests::max_supported_mode(),
),
)
.ok()
}
fn write(&self, writer: &mut impl BufMut) {
UInt(self.max_participants.get()).write(writer);
UInt(u32::try_from(self.bytes.len()).expect("payload too large")).write(writer);
writer.put_slice(&self.bytes);
}
fn read(reader: &mut impl Buf) -> Result<Self, CodecError> {
let max_participants = NonZeroU32::new(UInt::<u32>::read(reader)?.into()).ok_or(
CodecError::Invalid("EncodedPayload", "max participants must be non-zero"),
)?;
let len: u32 = UInt::read(reader)?.into();
let len = len as usize;
if reader.remaining() < len {
return Err(CodecError::EndOfBuffer);
}
let bytes = reader.copy_to_bytes(len).to_vec();
Ok(Self {
max_participants,
bytes,
})
}
fn encode_size(&self) -> usize {
UInt(self.max_participants.get()).encode_size()
+ UInt(u32::try_from(self.bytes.len()).expect("payload too large")).encode_size()
+ self.bytes.len()
}
}
impl<D: Digest, C: Codec, Dir> MockBlock<D, C, Dir> {
pub(crate) fn new<H: Hasher<Digest = D>>(
context: C,
parent: D,
height: Height,
timestamp: u64,
) -> Self {
Self::from_parts::<H>(context, parent, height, timestamp, None)
}
pub(crate) fn with_payload<H, V, S>(
self,
max_participants: NonZeroU32,
payload: Payload<V, S, Dir>,
) -> Self
where
H: Hasher<Digest = D>,
V: Variant,
S: Signer,
Dir: DkgDirectory<S::PublicKey>,
{
Self::from_parts::<H>(
self.context,
self.parent,
self.height,
self.timestamp,
Some(EncodedPayload::new(max_participants, payload)),
)
}
pub(crate) const fn context(&self) -> &C {
&self.context
}
fn from_parts<H: Hasher<Digest = D>>(
context: C,
parent: D,
height: Height,
timestamp: u64,
payload: Option<EncodedPayload>,
) -> Self {
let height_be = height.get().to_be_bytes();
let context_enc = context.encode();
let timestamp_be = timestamp.to_be_bytes();
let digest = payload.as_ref().map_or_else(
|| H::hash(&[&parent, &height_be, &context_enc, ×tamp_be, &[0]]),
|payload| {
H::hash(&[
&parent,
&height_be,
&context_enc,
×tamp_be,
&[1],
&payload.max_participants.get().to_be_bytes(),
&u32::try_from(payload.bytes.len())
.expect("payload too large")
.to_be_bytes(),
&payload.bytes,
])
},
);
Self {
context,
parent,
height,
timestamp,
payload,
digest,
_directory: PhantomData,
}
}
}
impl<D: Digest, C: Write, Dir> Write for MockBlock<D, C, Dir> {
fn write(&self, writer: &mut impl BufMut) {
self.context.write(writer);
self.parent.write(writer);
self.height.write(writer);
UInt(self.timestamp).write(writer);
self.payload.is_some().write(writer);
if let Some(log) = &self.payload {
log.write(writer);
}
self.digest.write(writer);
}
}
impl<D: Digest, C: Read<Cfg = ()>, Dir> Read for MockBlock<D, C, Dir> {
type Cfg = ();
fn read_cfg(reader: &mut impl Buf, _: &Self::Cfg) -> Result<Self, CodecError> {
Ok(Self {
context: C::read(reader)?,
parent: D::read(reader)?,
height: Height::read(reader)?,
timestamp: UInt::read(reader)?.into(),
payload: if bool::read(reader)? {
Some(EncodedPayload::read(reader)?)
} else {
None
},
digest: D::read(reader)?,
_directory: PhantomData,
})
}
}
impl<D: Digest, C: EncodeSize, Dir> EncodeSize for MockBlock<D, C, Dir> {
fn encode_size(&self) -> usize {
self.context.encode_size()
+ self.parent.encode_size()
+ self.height.encode_size()
+ UInt(self.timestamp).encode_size()
+ self.payload.is_some().encode_size()
+ self.payload.as_ref().map_or(0, EncodedPayload::encode_size)
+ self.digest.encode_size()
}
}
impl<D: Digest, C: Clone + Send + Sync + 'static, Dir: Clone + Send + Sync + 'static> Digestible
for MockBlock<D, C, Dir>
{
type Digest = D;
fn digest(&self) -> D {
self.digest
}
}
impl<D: Digest, C: Clone + Send + Sync + 'static, Dir: Clone + Send + Sync + 'static> Heightable
for MockBlock<D, C, Dir>
{
fn height(&self) -> Height {
self.height
}
}
impl<D: Digest, C: Codec<Cfg = ()> + Clone + Send + Sync + 'static, Dir> Block
for MockBlock<D, C, Dir>
where
Dir: Clone + Send + Sync + 'static,
{
fn parent(&self) -> Self::Digest {
self.parent
}
}
impl<D, C, Dir> ReshareBlock for MockBlock<D, C, Dir>
where
D: Digest,
C: Codec<Cfg = ()> + Clone + Send + Sync + 'static,
Dir: DkgDirectory<TestPublicKey>,
{
type Variant = TestBlsVariant;
type Signer = TestSigner;
type Directory = Dir;
fn payload(&self) -> Option<Payload<Self::Variant, Self::Signer, Self::Directory>> {
self.payload.as_ref()?.decode()
}
}
#[derive(Clone, Default)]
pub(crate) struct MockApplication {
broadcasts: Arc<Mutex<Vec<TestDigest>>>,
proposals: Arc<Mutex<Vec<TestContext>>>,
}
impl MockApplication {
pub(crate) fn broadcasts(&self) -> Vec<TestDigest> {
self.broadcasts.lock().clone()
}
pub(crate) fn proposals(&self) -> Vec<TestContext> {
self.proposals.lock().clone()
}
}
impl Automaton for MockApplication {
type Context = TestContext;
type Digest = TestDigest;
async fn propose(&mut self, _context: Self::Context) -> oneshot::Receiver<Self::Digest> {
let (sender, receiver) = oneshot::channel();
self.proposals.lock().push(_context);
sender.send_lossy(Sha256::hash(&[b"proposal"]));
receiver
}
async fn verify(
&mut self,
_context: Self::Context,
_payload: Self::Digest,
) -> oneshot::Receiver<bool> {
let (sender, receiver) = oneshot::channel();
sender.send_lossy(true);
receiver
}
}
impl CertifiableAutomaton for MockApplication {}
impl Relay for MockApplication {
type Digest = TestDigest;
type PublicKey = TestPublicKey;
type Plan = Plan<TestPublicKey>;
fn broadcast(&mut self, payload: Self::Digest, _plan: Self::Plan) -> Feedback {
self.broadcasts.lock().push(payload);
Feedback::Ok
}
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub(crate) enum ConsumerEvent {
Enter(Epoch),
Exit(Epoch),
}
#[derive(Clone, Default)]
pub(crate) struct MockConsumer {
events: Arc<Mutex<Vec<ConsumerEvent>>>,
}
impl MockConsumer {
pub(crate) fn events(&self) -> Vec<ConsumerEvent> {
self.events.lock().clone()
}
}
impl Registrar for MockConsumer {
type Variant = TestBlsVariant;
type PublicKey = TestPublicKey;
async fn register(&self, epoch: Epoch, _info: SchemeInfo<Self::Variant, Self::PublicKey>) {
self.events.lock().push(ConsumerEvent::Enter(epoch));
}
}
#[derive(Clone, Default)]
pub(crate) struct MarshalApplication {
blocks: Arc<Mutex<BTreeMap<Height, Arc<TestBlock>>>>,
}
impl MarshalApplication {
pub(crate) fn blocks(&self) -> BTreeMap<Height, Arc<TestBlock>> {
self.blocks.lock().clone()
}
}
impl Reporter for MarshalApplication {
type Activity = Update<TestBlock>;
fn report(&mut self, activity: Self::Activity) -> Feedback {
if let Update::Block(block, ack) = activity {
self.blocks.lock().insert(block.height(), block);
ack.acknowledge();
}
Feedback::Ok
}
}
pub(crate) struct SchemeFixture {
pub(crate) participants: Vec<TestPublicKey>,
pub(crate) schemes: Vec<TestScheme>,
pub(crate) provider: TestProvider,
}
pub(crate) fn scheme_fixture(context: &mut deterministic::Context) -> SchemeFixture {
scheme_fixture_n(context, 1)
}
pub(crate) fn scheme_fixture_n(context: &mut deterministic::Context, n: u32) -> SchemeFixture {
let fixture = scheme::fixture(context, NAMESPACE, n);
let provider = ConstantProvider::new(fixture.schemes[0].clone());
SchemeFixture {
participants: fixture.participants,
schemes: fixture.schemes,
provider,
}
}
pub(crate) fn genesis_block(leader: TestPublicKey) -> TestBlock {
let digest = Sha256::hash(&[b""]);
let context = TestContext {
round: Round::new(Epoch::zero(), View::zero()),
leader,
parent: (View::zero(), digest),
};
TestBlock::new::<Sha256>(context, digest, Height::zero(), 0)
}
pub(crate) async fn closed_marshal_mailbox(
context: deterministic::Context,
signer: &TestSigner,
scheme: TestScheme,
partition_prefix: &str,
blocks_per_epoch: NonZeroU64,
) -> TestMarshalMailbox {
let page_cache = CacheRef::from_pooler(&context, NZU16!(1024), NZUsize!(8));
let finalizations_by_height =
immutable::Archive::init(context.child("finalizations_by_height"), {
let _: () = TestScheme::certificate_codec_config_unbounded();
archive_config(partition_prefix, "finalizations", page_cache.clone(), ())
})
.await
.expect("finalizations archive");
let finalized_blocks = immutable::Archive::init(
context.child("finalized_blocks"),
archive_config(partition_prefix, "blocks", page_cache.clone(), ()),
)
.await
.expect("blocks archive");
let (actor, mailbox, _) = MarshalActor::<_, _, _, _, _, _, _, Exact>::init(
context.child("marshal"),
finalizations_by_height,
finalized_blocks,
marshal::Config {
provider: TestProvider::new(scheme),
epocher: FixedEpocher::new(blocks_per_epoch),
start: MarshalStart::Genesis(genesis_block(signer.public_key())),
partition_prefix: format!("{partition_prefix}-marshal"),
mailbox_size: NZUsize!(16),
view_retention: ViewDelta::new(8),
prunable_items_per_section: NZU64!(10),
page_cache,
replay_buffer: NZUsize!(1024),
key_write_buffer: NZUsize!(1024),
value_write_buffer: NZUsize!(1024),
block_codec_config: (),
max_repair: NZUsize!(4),
max_pending_acks: NZUsize!(4),
strategy: Sequential,
},
)
.await;
drop(actor);
mailbox
}
fn archive_config<C>(
prefix: &str,
name: &str,
page_cache: CacheRef,
codec_config: C,
) -> immutable::Config<C> {
immutable::Config {
metadata_partition: format!("{prefix}-{name}-metadata"),
freezer_table_partition: format!("{prefix}-{name}-freezer-table"),
freezer_table_initial_size: 64,
freezer_table_resize_frequency: 10,
freezer_table_resize_chunk_size: 10,
freezer_key_partition: format!("{prefix}-{name}-freezer-key"),
freezer_key_page_cache: page_cache,
freezer_value_partition: format!("{prefix}-{name}-freezer-value"),
freezer_value_target_size: 1024,
freezer_value_compression: None,
ordinal_partition: format!("{prefix}-{name}-ordinal"),
items_per_section: NZU64!(10),
codec_config,
replay_buffer: NZUsize!(1024),
freezer_key_write_buffer: NZUsize!(1024),
freezer_value_write_buffer: NZUsize!(1024),
ordinal_write_buffer: NZUsize!(1024),
}
}
pub(crate) fn simplex_config() -> orchestrator::SimplexConfig<TestElector> {
orchestrator::SimplexConfig {
elector: TestElector::default(),
mailbox_size: NZUsize!(16),
replay_buffer: NZUsize!(1024),
write_buffer: NZUsize!(1024),
page_cache_page_size: NZU16!(1024),
page_cache_pages: NZUsize!(8),
leader_timeout: Duration::from_millis(100),
certification_timeout: Duration::from_millis(200),
timeout_retry: Duration::from_millis(500),
fetch_timeout: Duration::from_millis(100),
view_retention: ViewDelta::new(8),
skip: SkipPolicy::Enabled {
timeout: Duration::from_secs(1),
budget: simplex::SkipBudget::Participants,
},
forward: ForwardPolicy::Disabled,
track_historical_votes: false,
}
}
#[derive(Clone, Default)]
pub(crate) struct MemorySecretStore {
inner: Arc<Mutex<MemorySecretStoreInner>>,
}
#[derive(Default)]
struct MemorySecretStoreInner {
shares: BTreeMap<Epoch, Share>,
seeds: BTreeMap<Epoch, Summary>,
dealings: BTreeMap<(Epoch, Vec<u8>), DealerPrivMsg>,
prunes: Vec<Epoch>,
}
impl MemorySecretStore {
pub(crate) fn has_share(&self, epoch: Epoch) -> bool {
self.inner.lock().shares.contains_key(&epoch)
}
pub(crate) fn prunes(&self) -> Vec<Epoch> {
self.inner.lock().prunes.clone()
}
pub(crate) fn seed_share(&self, epoch: Epoch, share: Share) {
self.inner.lock().shares.insert(epoch, share);
}
}
impl SecretStore for MemorySecretStore {
async fn put_share(&mut self, epoch: Epoch, share: Share) {
self.inner.lock().shares.insert(epoch, share);
}
async fn get_share(&mut self, epoch: Epoch) -> Option<Share> {
self.inner.lock().shares.get(&epoch).cloned()
}
async fn put_seed(&mut self, epoch: Epoch, seed: Summary) {
self.inner.lock().seeds.insert(epoch, seed);
}
async fn get_seed(&mut self, epoch: Epoch) -> Option<Summary> {
self.inner.lock().seeds.get(&epoch).cloned()
}
async fn put_dealing<P: CryptoPublicKey>(
&mut self,
epoch: Epoch,
dealer: P,
private: DealerPrivMsg,
) {
self.inner
.lock()
.dealings
.insert((epoch, dealer.encode().to_vec()), private);
}
async fn get_dealing<P: CryptoPublicKey>(
&mut self,
epoch: Epoch,
dealer: &P,
) -> Option<DealerPrivMsg> {
self.inner
.lock()
.dealings
.get(&(epoch, dealer.encode().to_vec()))
.cloned()
}
async fn prune(&mut self, min: Epoch) {
let mut inner = self.inner.lock();
inner.prunes.push(min);
inner.shares.retain(|epoch, _| *epoch >= min);
inner.seeds.retain(|epoch, _| *epoch >= min);
inner.dealings.retain(|(epoch, _), _| *epoch >= min);
}
}