From f89a4e399fedeb4cb179800c1cb473386b502cda Mon Sep 17 00:00:00 2001
From: Remco Bloemen <remco@wicked.ventures>
Date: Fri, 18 Mar 2022 15:28:30 -0700
Subject: [PATCH] Mutex guard the witness calculator

---
 src/circuit.rs  | 5 +++--
 src/protocol.rs | 6 ++++--
 2 files changed, 7 insertions(+), 4 deletions(-)

diff --git a/src/circuit.rs b/src/circuit.rs
index 175c1bf..274b9ba 100644
--- a/src/circuit.rs
+++ b/src/circuit.rs
@@ -6,6 +6,7 @@ use core::include_bytes;
 use once_cell::sync::Lazy;
 use std::io::{Cursor, Write};
 use tempfile::NamedTempFile;
+use std::sync::Mutex;
 
 const ZKEY_BYTES: &[u8] = include_bytes!("../semaphore/build/snark/semaphore_final.zkey");
 const WASM: &[u8] = include_bytes!("../semaphore/build/snark/semaphore.wasm");
@@ -15,7 +16,7 @@ pub static ZKEY: Lazy<(ProvingKey<Bn254>, ConstraintMatrices<Fr>)> = Lazy::new(|
     read_zkey(&mut reader).expect("zkey should be valid")
 });
 
-pub static WITNESS_CALCULATOR: Lazy<WitnessCalculator> = Lazy::new(|| {
+pub static WITNESS_CALCULATOR: Lazy<Mutex<WitnessCalculator>> = Lazy::new(|| {
     // HACK: ark-circom requires a file, so we make one!
     let mut tmpfile = NamedTempFile::new().expect("Failed to create temp file");
     let written = tmpfile.write(WASM).expect("Failed to write to temp file");
@@ -23,5 +24,5 @@ pub static WITNESS_CALCULATOR: Lazy<WitnessCalculator> = Lazy::new(|| {
     let path = tmpfile.into_temp_path();
     let result = WitnessCalculator::new(&path).expect("Failed to create witness calculator");
     path.close().expect("Could not remove tempfile");
-    result
+    Mutex::new(result)
 });
diff --git a/src/protocol.rs b/src/protocol.rs
index 1e65361..a75eb67 100644
--- a/src/protocol.rs
+++ b/src/protocol.rs
@@ -116,8 +116,10 @@ pub fn generate_proof(
 
     let now = Instant::now();
 
-    let full_assignment = WITNESS_CALCULATOR
-        .clone()
+    let full_assignment =
+        WITNESS_CALCULATOR
+        .lock()
+        .expect("witness_calculator mutex should not get poisoned")
         .calculate_witness_element::<Bn254, _>(inputs, false)
         .map_err(ProofError::WitnessError)?;