From b1daefca967d409699fecff2ca90ea50d27cdc94 Mon Sep 17 00:00:00 2001
From: oskarth <ot@oskarthoren.com>
Date: Thu, 18 Nov 2021 18:03:37 +0800
Subject: [PATCH] Early exit in error callback from Wasm (#9)

* Early exit in error callback from Wasm

This avoids Wasm execution hanging due to problems such as wrong public
input.

Mimics circom_runtime behaviour with less detailed debug information.

See https://github.com/iden3/circom_runtime/blob/master/js/witness_calculator.js#L52-L64

Adds test for wrong public input. Without early exit, the test stalls.
With it, the Circom build step fails as expected.

* chore: clean up error handling

* ci: add caching

Co-authored-by: Georgios Konstantopoulos <me@gakonst.com>
---
 .github/workflows/ci.yml          |  6 +++++-
 src/witness/witness_calculator.rs | 18 ++++++++++++++++--
 tests/groth16.rs                  | 21 +++++++++++++++++++++
 3 files changed, 42 insertions(+), 3 deletions(-)

diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index 0f38b7c..43d2848 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -1,6 +1,6 @@
 on:
   push:
-    branches: 
+    branches:
       - master
   pull_request:
 
@@ -35,6 +35,10 @@ jobs:
             export PATH=$HOME/bin:$PATH
             solc --version
 
+      - uses: Swatinem/rust-cache@v1
+        with:
+          cache-on-failure: true
+
       - name: cargo test
         run: |
             export PATH=$HOME/bin:$PATH
diff --git a/src/witness/witness_calculator.rs b/src/witness/witness_calculator.rs
index 005cbfd..46fb6f2 100644
--- a/src/witness/witness_calculator.rs
+++ b/src/witness/witness_calculator.rs
@@ -2,7 +2,7 @@ use color_eyre::Result;
 use num_bigint::BigInt;
 use num_traits::Zero;
 use std::cell::Cell;
-use wasmer::{imports, Function, Instance, Memory, MemoryType, Module, Store};
+use wasmer::{imports, Function, Instance, Memory, MemoryType, Module, RuntimeError, Store};
 
 use super::{fnv, SafeMemory, Wasm};
 
@@ -13,6 +13,12 @@ pub struct WitnessCalculator {
     pub n64: i32,
 }
 
+// Error type to signal end of execution.
+// From https://docs.wasmer.io/integrations/examples/exit-early
+#[derive(thiserror::Error, Debug, Clone, Copy)]
+#[error("{0}")]
+struct ExitCode(u32);
+
 impl WitnessCalculator {
     pub fn new(path: impl AsRef<std::path::Path>) -> Result<Self> {
         let store = Store::default();
@@ -144,7 +150,15 @@ mod runtime {
     pub fn error(store: &Store) -> Function {
         #[allow(unused)]
         #[allow(clippy::many_single_char_names)]
-        fn func(a: i32, b: i32, c: i32, d: i32, e: i32, f: i32) {}
+        fn func(a: i32, b: i32, c: i32, d: i32, e: i32, f: i32) {
+            // NOTE: We can also get more information why it is failing, see p2str etc here:
+            // https://github.com/iden3/circom_runtime/blob/master/js/witness_calculator.js#L52-L64
+            println!(
+                "runtime error, exiting early: {0} {1} {2} {3} {4} {5}",
+                a, b, c, d, e, f
+            );
+            RuntimeError::raise(Box::new(ExitCode(1)));
+        }
         Function::new_native(store, func)
     }
 
diff --git a/tests/groth16.rs b/tests/groth16.rs
index 754bb95..c7eadd7 100644
--- a/tests/groth16.rs
+++ b/tests/groth16.rs
@@ -37,3 +37,24 @@ fn groth16_proof() -> Result<()> {
 
     Ok(())
 }
+
+#[test]
+fn groth16_proof_wrong_input() {
+    let cfg = CircomConfig::<Bn254>::new(
+        "./test-vectors/mycircuit.wasm",
+        "./test-vectors/mycircuit.r1cs",
+    )
+    .unwrap();
+    let mut builder = CircomBuilder::new(cfg);
+    builder.push_input("a", 3);
+    // This isn't a public input to the circuit, should faild
+    builder.push_input("foo", 11);
+
+    // create an empty instance for setting it up
+    let circom = builder.setup();
+
+    let mut rng = thread_rng();
+    let _params = generate_random_parameters::<Bn254, _, _>(circom, &mut rng).unwrap();
+
+    builder.build().unwrap_err();
+}