openmina_node_common/service/
snark_worker.rs

1use ledger::{
2    proofs::{
3        provers::{TransactionProver, ZkappProver},
4        zkapp::ZkappParams,
5    },
6    scan_state::scan_state::transaction_snark::SokMessage,
7};
8use mina_p2p_messages::v2;
9use mina_signer::CompressedPubKey;
10use node::{
11    core::channels::mpsc,
12    event_source::ExternalSnarkWorkerEvent,
13    external_snark_worker::{
14        ExternalSnarkWorkerError, ExternalSnarkWorkerWorkError, SnarkWorkResult, SnarkWorkSpec,
15        SnarkWorkSpecError,
16    },
17    snark::TransactionVerifier,
18};
19
20use crate::NodeService;
21
22use super::EventSender;
23
24pub struct SnarkWorker {
25    cmd_sender: mpsc::UnboundedSender<Cmd>,
26}
27
28enum Cmd {
29    Submit(Box<SnarkWorkSpec>),
30    Cancel,
31    Kill,
32}
33
34impl node::service::ExternalSnarkWorkerService for NodeService {
35    fn start(
36        &mut self,
37        pub_key: v2::NonZeroCurvePoint,
38        fee: v2::CurrencyFeeStableV1,
39        work_verifier: TransactionVerifier,
40    ) -> Result<(), ExternalSnarkWorkerError> {
41        if self.replayer.is_some() {
42            return Ok(());
43        }
44        let (cmd_sender, cmd_receiver) = mpsc::unbounded_channel();
45        // TODO(binier): improve pub key conv
46        let sok_message = SokMessage::create(
47            (&fee).into(),
48            CompressedPubKey::from_address(&pub_key.to_string()).unwrap(),
49        );
50        self.snark_worker = Some(SnarkWorker { cmd_sender });
51        let event_sender = self.event_sender().clone();
52
53        node::core::thread::Builder::new()
54            .name("snark_worker".to_owned())
55            .spawn(move || worker_thread(cmd_receiver, event_sender, sok_message, work_verifier))
56            .map(|_| ())
57            .map_err(|err| ExternalSnarkWorkerError::Error(err.to_string()))
58    }
59
60    fn kill(&mut self) -> Result<(), ExternalSnarkWorkerError> {
61        if self.replayer.is_some() {
62            return Ok(());
63        }
64
65        if self
66            .snark_worker
67            .as_ref()
68            .and_then(|s| s.cmd_sender.send(Cmd::Kill).ok())
69            .is_none()
70        {
71            return Err(ExternalSnarkWorkerError::NotRunning);
72        }
73        Ok(())
74    }
75
76    fn submit(&mut self, spec: SnarkWorkSpec) -> Result<(), ExternalSnarkWorkerError> {
77        if self.replayer.is_some() {
78            return Ok(());
79        }
80
81        if self
82            .snark_worker
83            .as_ref()
84            .and_then(|s| s.cmd_sender.send(Cmd::Submit(spec.into())).ok())
85            .is_none()
86        {
87            return Err(ExternalSnarkWorkerError::NotRunning);
88        }
89        Ok(())
90    }
91
92    fn cancel(&mut self) -> Result<(), ExternalSnarkWorkerError> {
93        if self.replayer.is_some() {
94            return Ok(());
95        }
96
97        // TODO(binier): for wasm threads, call terminate:
98        // <https://developer.mozilla.org/en-US/docs/Web/API/Worker/terminate>
99        if self
100            .snark_worker
101            .as_ref()
102            .and_then(|s| s.cmd_sender.send(Cmd::Cancel).ok())
103            .is_none()
104        {
105            return Err(ExternalSnarkWorkerError::NotRunning);
106        }
107        Ok(())
108    }
109}
110
111fn worker_thread(
112    mut cmd_receiver: mpsc::UnboundedReceiver<Cmd>,
113    event_sender: EventSender,
114    sok_message: SokMessage,
115    work_verifier: TransactionVerifier,
116) {
117    let _ = event_sender.send(ExternalSnarkWorkerEvent::Started.into());
118    let tx_prover = TransactionProver::make(Some(work_verifier.clone()));
119    let zkapp_prover = ZkappProver::make(Some(work_verifier));
120    while let Some(cmd) = cmd_receiver.blocking_recv() {
121        match cmd {
122            Cmd::Kill => {
123                let _ = event_sender.send(ExternalSnarkWorkerEvent::Killed.into());
124                return;
125            }
126            Cmd::Cancel => {
127                // can't cancel as it's a blocking thread. Once this
128                // is moved to another process, kill it.
129                let _ = event_sender.send(ExternalSnarkWorkerEvent::WorkCancelled.into());
130            }
131            Cmd::Submit(spec) => {
132                let event = match prove_spec(&tx_prover, &zkapp_prover, *spec, &sok_message) {
133                    Err(err) => ExternalSnarkWorkerEvent::WorkError(err),
134                    Ok(res) => ExternalSnarkWorkerEvent::WorkResult(res),
135                };
136
137                let _ = event_sender.send(event.into());
138            }
139        }
140    }
141}
142
143fn prove_spec(
144    tx_prover: &TransactionProver,
145    zkapp_prover: &ZkappProver,
146    spec: SnarkWorkSpec,
147    sok_message: &SokMessage,
148) -> Result<SnarkWorkResult, ExternalSnarkWorkerWorkError> {
149    match spec {
150        SnarkWorkSpec::One(single) => prove_single(tx_prover, zkapp_prover, single, sok_message)
151            .map(v2::TransactionSnarkWorkTStableV2Proofs::One),
152        SnarkWorkSpec::Two((one, two)) => Ok(v2::TransactionSnarkWorkTStableV2Proofs::Two((
153            prove_single(tx_prover, zkapp_prover, one, sok_message)?,
154            prove_single(tx_prover, zkapp_prover, two, sok_message)?,
155        ))),
156    }
157    .map(Into::into)
158}
159
160fn invalid_bigint_err() -> ExternalSnarkWorkerWorkError {
161    ExternalSnarkWorkerWorkError::WorkSpecError(SnarkWorkSpecError::InvalidBigInt)
162}
163
164fn prove_single(
165    tx_prover: &TransactionProver,
166    zkapp_prover: &ZkappProver,
167    single: v2::SnarkWorkerWorkerRpcsVersionedGetWorkV2TResponseA0Single,
168    sok_message: &SokMessage,
169) -> Result<v2::LedgerProofProdStableV2, ExternalSnarkWorkerWorkError> {
170    use ledger::proofs::{merge::MergeParams, transaction::TransactionParams};
171
172    let (snarked_ledger_state, res) = match single {
173        v2::SnarkWorkerWorkerRpcsVersionedGetWorkV2TResponseA0Single::Transition(
174            snarked_ledger_state,
175            witness,
176        ) => {
177            if let v2::MinaTransactionTransactionStableV2::Command(cmd) = &witness.transaction {
178                if matches!(&**cmd, v2::MinaBaseUserCommandStableV2::ZkappCommand(_)) {
179                    return prove_zkapp(zkapp_prover, snarked_ledger_state, witness, sok_message);
180                }
181            }
182            let res = ledger::proofs::generate_tx_proof(TransactionParams {
183                statement: &snarked_ledger_state.0,
184                tx_witness: &witness,
185                message: sok_message,
186                tx_step_prover: &tx_prover.tx_step_prover,
187                tx_wrap_prover: &tx_prover.tx_wrap_prover,
188                only_verify_constraints: false,
189                expected_step_proof: None,
190                ocaml_wrap_witness: None,
191            });
192            (snarked_ledger_state.0, res)
193        }
194        v2::SnarkWorkerWorkerRpcsVersionedGetWorkV2TResponseA0Single::Merge(data) => {
195            let (snarked_ledger_state, proof_1, proof_2) = *data;
196            let res = ledger::proofs::generate_merge_proof(MergeParams {
197                statement: (&snarked_ledger_state.0)
198                    .try_into()
199                    .map_err(|_| invalid_bigint_err())?,
200                proofs: &[proof_1, proof_2],
201                message: sok_message,
202                step_prover: &tx_prover.merge_step_prover,
203                wrap_prover: &tx_prover.tx_wrap_prover,
204                only_verify_constraints: false,
205                expected_step_proof: None,
206                ocaml_wrap_witness: None,
207            });
208            (snarked_ledger_state.0, res)
209        }
210    };
211    res.map_err(|err| ExternalSnarkWorkerWorkError::Error(err.to_string()))
212        .map(|proof| {
213            v2::LedgerProofProdStableV2(v2::TransactionSnarkStableV2 {
214                statement: v2::MinaStateSnarkedLedgerStateWithSokStableV2 {
215                    source: snarked_ledger_state.source,
216                    target: snarked_ledger_state.target,
217                    connecting_ledger_left: snarked_ledger_state.connecting_ledger_left,
218                    connecting_ledger_right: snarked_ledger_state.connecting_ledger_right,
219                    supply_increase: snarked_ledger_state.supply_increase,
220                    fee_excess: snarked_ledger_state.fee_excess,
221                    sok_digest: (&sok_message.digest()).into(),
222                },
223                proof: v2::TransactionSnarkProofStableV2((&proof).into()),
224            })
225        })
226}
227
228fn prove_zkapp(
229    zkapp_prover: &ZkappProver,
230    snarked_ledger_state: v2::MinaStateSnarkedLedgerStateStableV2,
231    witness: v2::TransactionWitnessStableV2,
232    sok_message: &SokMessage,
233) -> Result<v2::LedgerProofProdStableV2, ExternalSnarkWorkerWorkError> {
234    ledger::proofs::generate_zkapp_proof(ZkappParams {
235        statement: &snarked_ledger_state.0,
236        tx_witness: &witness,
237        message: sok_message,
238        step_opt_signed_opt_signed_prover: &zkapp_prover.step_opt_signed_opt_signed_prover,
239        step_opt_signed_prover: &zkapp_prover.step_opt_signed_prover,
240        step_proof_prover: &zkapp_prover.step_proof_prover,
241        merge_step_prover: &zkapp_prover.merge_step_prover,
242        tx_wrap_prover: &zkapp_prover.tx_wrap_prover,
243        opt_signed_path: None,
244        proved_path: None,
245    })
246    .map(|proof| (&proof).into())
247    .map_err(|err| ExternalSnarkWorkerWorkError::Error(err.to_string()))
248}
openmina_node_common/service/snark_worker.rs

openmina_node_common/service/
snark_worker.rs
