use crate::{CertifiableBlock, Epochable, Handoff, marshal::ancestry::Ancestry};
use commonware_cryptography::Digested;
use commonware_parallel::Sequential;
use commonware_runtime::deterministic;
use commonware_utils::{
channel::{fallible::OneshotExt, oneshot},
sync::Mutex,
};
use std::{future::pending, marker::PhantomData, sync::Arc};
#[derive(Clone)]
pub struct MockVerifyingApp<B, S> {
pub propose_result: Option<B>,
pub verify_result: bool,
handoff: Handoff<()>,
proposal_gate: Arc<Mutex<Option<ProposalGate>>>,
pub reject: Option<fn(&B) -> bool>,
_phantom: PhantomData<S>,
}
impl<B, S> MockVerifyingApp<B, S> {
pub fn new() -> Self {
Self::default()
}
pub fn with_verify_result(verify_result: bool) -> Self {
Self {
verify_result,
..Self::default()
}
}
pub fn with_propose_result(mut self, block: B) -> Self {
self.propose_result = Some(block);
self
}
pub fn with_reject(mut self, reject: fn(&B) -> bool) -> Self {
self.reject = Some(reject);
self
}
pub const fn with_handoff(mut self, handoff: Handoff<()>) -> Self {
self.handoff = handoff;
self
}
pub fn with_proposal_gate(mut self) -> (Self, oneshot::Receiver<()>, oneshot::Receiver<()>) {
let (started, started_rx) = oneshot::channel();
let (dropped, dropped_rx) = oneshot::channel();
self.proposal_gate = Arc::new(Mutex::new(Some(ProposalGate { started, dropped })));
(self, started_rx, dropped_rx)
}
}
struct ProposalGate {
started: oneshot::Sender<()>,
dropped: oneshot::Sender<()>,
}
impl<B, S> Default for MockVerifyingApp<B, S> {
fn default() -> Self {
Self {
propose_result: None,
verify_result: true,
handoff: Handoff::Wait,
proposal_gate: Arc::default(),
reject: None,
_phantom: PhantomData,
}
}
}
impl<B, S> crate::Application<deterministic::Context> for MockVerifyingApp<B, S>
where
B: CertifiableBlock,
B::Context: Epochable + Send + Sync + 'static,
S: commonware_cryptography::certificate::Scheme,
{
type Block = B;
type Context = B::Context;
type SigningScheme = S;
type Input = ();
async fn prepare(
&mut self,
context: (deterministic::Context, Self::Context),
ancestry: impl Ancestry<Self::Block>,
input: Self::Input,
) -> Handoff<Digested<Self::Block>> {
if self.handoff.is_wait() {
return Handoff::Wait;
}
let decision = self.handoff;
self.propose(context, ancestry, input)
.await
.map_or(Handoff::Wait, |block| decision.map(|()| block))
}
async fn propose(
&mut self,
_context: (deterministic::Context, Self::Context),
_ancestry: impl Ancestry<Self::Block>,
_input: Self::Input,
) -> Option<Digested<Self::Block>> {
let gate = self.proposal_gate.lock().take();
if let Some(gate) = gate {
let _dropped = gate.dropped;
gate.started.send_lossy(());
pending::<()>().await;
}
match self.propose_result.clone() {
Some(block) => Some(Digested::compute(block, &Sequential).await),
None => None,
}
}
async fn verify(
&mut self,
_context: (deterministic::Context, Self::Context),
ancestry: impl Ancestry<Self::Block>,
) -> bool {
if let (Some(reject), Some(block)) = (self.reject, ancestry.peek())
&& reject(block.inner())
{
return false;
}
self.verify_result
}
}
#[derive(Clone)]
pub struct GatedVerifyingApp<B, S> {
started: Arc<Mutex<Option<oneshot::Sender<()>>>>,
release: Arc<Mutex<Option<oneshot::Receiver<()>>>>,
_phantom: PhantomData<(B, S)>,
}
impl<B, S> GatedVerifyingApp<B, S> {
pub fn new() -> (Self, oneshot::Receiver<()>, oneshot::Sender<()>) {
let (started_tx, started_rx) = oneshot::channel();
let (release_tx, release_rx) = oneshot::channel();
(
Self {
started: Arc::new(Mutex::new(Some(started_tx))),
release: Arc::new(Mutex::new(Some(release_rx))),
_phantom: PhantomData,
},
started_rx,
release_tx,
)
}
}
impl<B, S> crate::Application<deterministic::Context> for GatedVerifyingApp<B, S>
where
B: CertifiableBlock,
B::Context: Epochable + Send + Sync + 'static,
S: commonware_cryptography::certificate::Scheme,
{
type Block = B;
type Context = B::Context;
type SigningScheme = S;
type Input = ();
async fn propose(
&mut self,
_context: (deterministic::Context, Self::Context),
_ancestry: impl Ancestry<Self::Block>,
_input: Self::Input,
) -> Option<Digested<Self::Block>> {
None
}
async fn verify(
&mut self,
_context: (deterministic::Context, Self::Context),
_ancestry: impl Ancestry<Self::Block>,
) -> bool {
if let Some(started) = self.started.lock().take() {
started.send_lossy(());
}
let release = self
.release
.lock()
.take()
.expect("release receiver missing");
let _ = release.await;
true
}
}