use crate::dkg::{network::Directory, types::EpochInfo};
use bytes::{Buf, BufMut};
use commonware_codec::{Decode as _, Encode as _, EncodeSize, Error as CodecError, Read, Write};
use commonware_consensus::{
Epochable as _,
marshal::core::{Mailbox as MarshalMailbox, Variant as MarshalVariant},
simplex::{scheme::Scheme, types::Finalization},
types::{Epoch, Epocher, FixedEpocher},
};
#[cfg(feature = "arbitrary")]
use commonware_cryptography::bls12381::dkg::feldman_desmedt::Output;
use commonware_cryptography::{
Digest,
bls12381::primitives::{sharing::ModeVersion, variant::Variant},
};
use commonware_storage::{
Context,
metadata::{self, Metadata},
};
use commonware_utils::{
fixed_bytes,
sequence::{FixedBytes, Unit},
sync::AsyncMutex,
};
use std::{fmt, num::NonZeroU32, sync::Arc};
const STATE_SYNC_KEY: FixedBytes<1> = fixed_bytes!("00");
const STATE_SYNC_SUFFIX: &str = "_dkg_state_sync";
type EpochInfoCodecConfig = (NonZeroU32, ModeVersion);
#[derive(Clone, Debug)]
pub struct Config {
pub partition_prefix: String,
pub max_participants: NonZeroU32,
pub max_supported_mode: ModeVersion,
}
pub struct StateSync<S, D, V, Dir = Unit>
where
S: Scheme<D>,
D: Digest,
V: Variant,
Dir: Directory<S::PublicKey>,
{
pub info: EpochInfo<V, S::PublicKey, Dir>,
pub floor: Finalization<S, D>,
}
impl<S, D, V, Dir> Clone for StateSync<S, D, V, Dir>
where
S: Scheme<D>,
D: Digest,
V: Variant,
Dir: Directory<S::PublicKey>,
{
fn clone(&self) -> Self {
Self {
info: self.info.clone(),
floor: self.floor.clone(),
}
}
}
impl<S, D, V, Dir> fmt::Debug for StateSync<S, D, V, Dir>
where
S: Scheme<D>,
D: Digest,
V: Variant,
Dir: Directory<S::PublicKey>,
{
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter
.debug_struct("StateSync")
.field("info", &self.info)
.field("floor", &self.floor)
.finish()
}
}
impl<S, D, V, Dir> PartialEq for StateSync<S, D, V, Dir>
where
S: Scheme<D>,
D: Digest,
V: Variant,
Dir: Directory<S::PublicKey>,
{
fn eq(&self, other: &Self) -> bool {
self.info == other.info && self.floor == other.floor
}
}
impl<S, D, V, Dir> Eq for StateSync<S, D, V, Dir>
where
S: Scheme<D>,
D: Digest,
V: Variant,
Dir: Directory<S::PublicKey>,
{
}
impl<S, D, V, Dir> Write for StateSync<S, D, V, Dir>
where
S: Scheme<D>,
D: Digest,
V: Variant,
Dir: Directory<S::PublicKey>,
{
fn write(&self, writer: &mut impl BufMut) {
self.info.write(writer);
self.floor.write(writer);
}
}
impl<S, D, V, Dir> EncodeSize for StateSync<S, D, V, Dir>
where
S: Scheme<D>,
D: Digest,
V: Variant,
Dir: Directory<S::PublicKey>,
{
fn encode_size(&self) -> usize {
self.info.encode_size() + self.floor.encode_size()
}
}
impl<S, D, V, Dir> Read for StateSync<S, D, V, Dir>
where
S: Scheme<D>,
D: Digest,
V: Variant,
Dir: Directory<S::PublicKey>,
{
type Cfg = (EpochInfoCodecConfig, <S::Certificate as Read>::Cfg);
fn read_cfg(
reader: &mut impl Buf,
(epoch_info, certificate): &Self::Cfg,
) -> Result<Self, CodecError> {
Ok(Self {
info: EpochInfo::read_cfg(reader, epoch_info)?,
floor: Finalization::read_cfg(reader, certificate)?,
})
}
}
#[cfg(feature = "arbitrary")]
impl<S, D, V, Dir> arbitrary::Arbitrary<'_> for StateSync<S, D, V, Dir>
where
S: Scheme<D>,
D: Digest + for<'a> arbitrary::Arbitrary<'a>,
V: Variant,
Dir: Directory<S::PublicKey> + for<'a> arbitrary::Arbitrary<'a>,
S::PublicKey: for<'a> arbitrary::Arbitrary<'a>,
S::Certificate: for<'a> arbitrary::Arbitrary<'a>,
Output<V, S::PublicKey>: for<'a> arbitrary::Arbitrary<'a>,
{
fn arbitrary(u: &mut arbitrary::Unstructured<'_>) -> arbitrary::Result<Self> {
Ok(Self {
info: u.arbitrary()?,
floor: u.arbitrary()?,
})
}
}
pub(crate) async fn recovered_epoch<S, V>(
marshal: &MarshalMailbox<S, V>,
epocher: &FixedEpocher,
) -> Option<Epoch>
where
S: commonware_cryptography::certificate::Scheme,
V: MarshalVariant,
{
let height = marshal.get_processed_height().await?.next();
Some(
epocher
.containing(height)
.expect("epocher must know recovered height")
.epoch(),
)
}
enum PlanState<S, D, V, Dir>
where
S: Scheme<D>,
D: Digest,
V: Variant,
Dir: Directory<S::PublicKey>,
{
Pending {
candidate: Option<StateSync<S, D, V, Dir>>,
partition: String,
codec_config: EpochInfoCodecConfig,
},
Resolved(Option<StateSync<S, D, V, Dir>>),
}
pub struct Plan<S, D, V, Dir = Unit>
where
S: Scheme<D>,
D: Digest,
V: Variant,
Dir: Directory<S::PublicKey>,
{
state: Arc<AsyncMutex<PlanState<S, D, V, Dir>>>,
}
impl<S, D, V, Dir> Clone for Plan<S, D, V, Dir>
where
S: Scheme<D>,
D: Digest,
V: Variant,
Dir: Directory<S::PublicKey>,
{
fn clone(&self) -> Self {
Self {
state: self.state.clone(),
}
}
}
impl<S, D, V, Dir> fmt::Debug for Plan<S, D, V, Dir>
where
S: Scheme<D>,
D: Digest,
V: Variant,
Dir: Directory<S::PublicKey>,
{
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter.debug_struct("Plan").finish_non_exhaustive()
}
}
impl<S, D, V, Dir> Plan<S, D, V, Dir>
where
S: Scheme<D>,
D: Digest,
V: Variant,
Dir: Directory<S::PublicKey>,
{
pub async fn init<E: Context>(
context: E,
config: Config,
provided: Option<StateSync<S, D, V, Dir>>,
) -> Self {
let codec_config = (config.max_participants, config.max_supported_mode);
if let Some(provided) = &provided {
assert_provided(provided, codec_config);
}
let partition = format!("{}{STATE_SYNC_SUFFIX}", config.partition_prefix);
let mut store =
open_store::<E, S, D, V, Dir>(context, partition.clone(), codec_config).await;
if let Some(provided) = provided {
store.put(STATE_SYNC_KEY, provided);
store = store
.sync()
.await
.expect("failed to persist DKG state sync metadata");
}
let candidate = store.get(&STATE_SYNC_KEY).cloned();
if let Some(candidate) = &candidate {
assert_epoch(candidate);
}
drop(store);
Self {
state: Arc::new(AsyncMutex::new(PlanState::Pending {
candidate,
partition,
codec_config,
})),
}
}
pub fn disabled() -> Self {
Self {
state: Arc::new(AsyncMutex::new(PlanState::Resolved(None))),
}
}
pub(crate) async fn resolve<E: Context>(
&self,
context: E,
recovered_epoch: Option<Epoch>,
) -> Option<StateSync<S, D, V, Dir>> {
let mut state = self.state.lock().await;
let (candidate, partition, codec_config) = match &*state {
PlanState::Resolved(resolved) => return resolved.clone(),
PlanState::Pending {
candidate,
partition,
codec_config,
} => (candidate.clone(), partition.clone(), *codec_config),
};
let Some(candidate) = candidate else {
*state = PlanState::Resolved(None);
return None;
};
if recovered_epoch.is_none_or(|epoch| epoch <= candidate.floor.epoch()) {
*state = PlanState::Resolved(Some(candidate.clone()));
return Some(candidate);
}
let mut store = open_store::<E, S, D, V, Dir>(context, partition, codec_config).await;
store.remove(&STATE_SYNC_KEY);
store
.sync()
.await
.expect("failed to delete stale DKG state sync metadata");
*state = PlanState::Resolved(None);
None
}
}
async fn open_store<E, S, D, V, Dir>(
context: E,
partition: String,
epoch_info_codec_config: EpochInfoCodecConfig,
) -> Metadata<E, FixedBytes<1>, StateSync<S, D, V, Dir>>
where
E: Context,
S: Scheme<D>,
D: Digest,
V: Variant,
Dir: Directory<S::PublicKey>,
{
Metadata::init(
context,
metadata::Config {
partition,
codec_config: (
epoch_info_codec_config,
S::certificate_codec_config_unbounded(),
),
},
)
.await
.expect("failed to load DKG state sync metadata")
}
fn assert_provided<S, D, V, Dir>(
state_sync: &StateSync<S, D, V, Dir>,
codec_config: EpochInfoCodecConfig,
) where
S: Scheme<D>,
D: Digest,
V: Variant,
Dir: Directory<S::PublicKey>,
{
assert_epoch(state_sync);
StateSync::<S, D, V, Dir>::decode_cfg(
state_sync.encode(),
&(codec_config, S::certificate_codec_config_unbounded()),
)
.expect("provided state sync material must satisfy codec config");
}
fn assert_epoch<S, D, V, Dir>(state_sync: &StateSync<S, D, V, Dir>)
where
S: Scheme<D>,
D: Digest,
V: Variant,
Dir: Directory<S::PublicKey>,
{
assert!(
state_sync.info.epoch == state_sync.floor.epoch(),
"state sync artifact and floor must be in the same epoch"
);
}
#[cfg(test)]
mod tests {
use super::*;
use crate::dkg::{network::Addresses, tests::mocks, types::EpochOutcome};
use commonware_consensus::{
simplex::types::{Finalize, Proposal},
types::{Round, View},
};
use commonware_cryptography::{
Hasher as _, Sha256, Signer as _,
bls12381::{dkg::feldman_desmedt::deal, primitives::sharing::Mode},
ed25519,
};
use commonware_p2p::Address;
use commonware_parallel::Sequential;
use commonware_runtime::{Runner as _, Supervisor as _, deterministic};
use commonware_utils::{N3f1, NZU32, TestRng, non_empty, ordered::Set};
use std::net::SocketAddr;
type TestStateSync = StateSync<mocks::TestScheme, mocks::TestDigest, mocks::TestBlsVariant>;
fn state_sync(context: &mut deterministic::Context, epoch: Epoch) -> TestStateSync {
let fixture = mocks::scheme_fixture_n(context, 4);
let participants = Set::from_iter_dedup(fixture.participants);
let (output, _) = deal::<mocks::TestBlsVariant, _, N3f1>(
TestRng::new(epoch.get() + 1),
Mode::NonZeroCounter,
participants.clone(),
)
.expect("test DKG output");
let proposal = Proposal::new(
Round::new(epoch, View::new(1)),
View::zero(),
Sha256::hash(&[b"state sync floor"]),
);
let finalizes = fixture
.schemes
.iter()
.map(|scheme| Finalize::sign(scheme, proposal.clone()).expect("test finalize"))
.collect::<Vec<_>>();
let floor = Finalization::from_finalizes(
&fixture.schemes[0],
non_empty![@finalizes.iter()],
&Sequential,
)
.expect("test finalization quorum");
TestStateSync {
info: EpochInfo {
outcome: EpochOutcome::Success,
epoch,
output,
players: participants.clone(),
next_players: participants,
directory: Unit,
},
floor,
}
}
type TestPlan = Plan<mocks::TestScheme, mocks::TestDigest, mocks::TestBlsVariant>;
fn config(partition: &str) -> Config {
Config {
partition_prefix: partition.into(),
max_participants: NZU32!(16),
max_supported_mode: crate::dkg::tests::max_supported_mode(),
}
}
async fn plan(
context: deterministic::Context,
partition: &str,
provided: Option<TestStateSync>,
) -> TestPlan {
Plan::init(context, config(partition), provided).await
}
fn assert_state_sync(actual: Option<TestStateSync>, expected: &TestStateSync) {
assert_eq!(actual.as_ref(), Some(expected));
}
#[test]
fn init_persists_before_resolve() {
deterministic::Runner::default().start(|mut context| async move {
let expected = state_sync(&mut context, Epoch::new(2));
let initialized = plan(
context.child("initialized"),
"persist-before-resolve",
Some(expected.clone()),
)
.await;
drop(initialized);
let reopened = plan(context.child("reopened"), "persist-before-resolve", None).await;
assert_state_sync(
reopened.resolve(context.child("resolve"), None).await,
&expected,
);
});
}
#[test]
fn two_clones_resolve_identically() {
deterministic::Runner::default().start(|mut context| async move {
let expected = state_sync(&mut context, Epoch::new(2));
let first = plan(context.child("init"), "clones", Some(expected.clone())).await;
let second = first.clone();
assert_state_sync(first.resolve(context.child("first"), None).await, &expected);
assert_state_sync(
second.resolve(context.child("second"), None).await,
&expected,
);
});
}
#[test]
fn stale_resolution_deletes_once_and_reopen_sees_none() {
deterministic::Runner::default().start(|mut context| async move {
let expected = state_sync(&mut context, Epoch::new(2));
let first = plan(context.child("init"), "stale", Some(expected)).await;
let second = first.clone();
assert!(
first
.resolve(context.child("first"), Some(Epoch::new(3)))
.await
.is_none()
);
assert!(
second
.resolve(context.child("second"), Some(Epoch::new(3)))
.await
.is_none()
);
drop(first);
drop(second);
let reopened = plan(context.child("reopened"), "stale", None).await;
assert!(
reopened
.resolve(context.child("reopened_resolve"), None)
.await
.is_none()
);
});
}
#[test]
fn provided_material_overwrites_persisted_record() {
deterministic::Runner::default().start(|mut context| async move {
let first_value = state_sync(&mut context, Epoch::new(1));
let replacement = state_sync(&mut context, Epoch::new(2));
drop(plan(context.child("first"), "overwrite", Some(first_value)).await);
drop(
plan(
context.child("replacement"),
"overwrite",
Some(replacement.clone()),
)
.await,
);
let reopened = plan(context.child("reopened"), "overwrite", None).await;
assert_state_sync(
reopened.resolve(context.child("resolve"), None).await,
&replacement,
);
});
}
#[test]
#[should_panic(expected = "state sync artifact and floor must be in the same epoch")]
fn artifact_beyond_floor_panics_at_init() {
deterministic::Runner::default().start(|mut context| async move {
let mut mismatched = state_sync(&mut context, Epoch::new(2));
mismatched.info.epoch = Epoch::new(3);
let _ = plan(context.child("init"), "mismatch", Some(mismatched)).await;
});
}
#[test]
#[should_panic(expected = "state sync artifact and floor must be in the same epoch")]
fn artifact_below_floor_panics_at_init() {
deterministic::Runner::default().start(|mut context| async move {
let mut mismatched = state_sync(&mut context, Epoch::new(3));
mismatched.info.epoch = Epoch::new(2);
let _ = plan(context.child("init"), "mismatch-below", Some(mismatched)).await;
});
}
#[test]
#[should_panic(expected = "provided state sync material must satisfy codec config")]
fn oversized_provided_material_panics_at_init() {
deterministic::Runner::default().start(|mut context| async move {
let provided = state_sync(&mut context, Epoch::new(2));
let mut config = config("oversized-provided");
config.max_participants = NZU32!(3);
let _ = TestPlan::init(context.child("init"), config, Some(provided)).await;
});
}
#[test]
#[should_panic(expected = "provided state sync material must satisfy codec config")]
fn invalid_addressed_directory_panics_at_init() {
deterministic::Runner::default().start(|mut context| async move {
let StateSync { info, floor } = state_sync(&mut context, Epoch::new(2));
let peers = info
.participants()
.tracked_peers()
.union()
.into_iter()
.chain([ed25519::PrivateKey::from_seed(10_000).public_key()]);
let directory = peers
.enumerate()
.map(|(index, peer)| {
let socket = SocketAddr::from(([127, 0, 0, 1], index as u16 + 1));
(peer, Address::Symmetric(socket))
})
.collect::<Addresses<_>>();
let EpochInfo {
outcome,
epoch,
output,
players,
next_players,
..
} = info;
let invalid = StateSync {
info: EpochInfo {
outcome,
epoch,
output,
players,
next_players,
directory,
},
floor,
};
let _ = Plan::<
mocks::TestScheme,
mocks::TestDigest,
mocks::TestBlsVariant,
Addresses<mocks::TestPublicKey>,
>::init(
context.child("init"),
config("invalid-directory"),
Some(invalid),
)
.await;
});
}
#[test]
fn disabled_resolves_none_without_storage() {
deterministic::Runner::default().start(|context| async move {
let plan = TestPlan::disabled();
assert!(plan.resolve(context, None).await.is_none());
});
}
}
#[cfg(all(test, feature = "arbitrary"))]
mod conformance {
use super::*;
use crate::dkg::{network::Addresses, tests::mocks};
use commonware_codec::conformance::CodecConformance;
commonware_conformance::conformance_tests! {
CodecConformance<StateSync<mocks::TestScheme, mocks::TestDigest, mocks::TestBlsVariant>> => 8192,
CodecConformance<StateSync<mocks::TestScheme, mocks::TestDigest, mocks::TestBlsVariant, Addresses<mocks::TestPublicKey>>> => 8192,
}
}