From 7db1420194b9fac30806f73704c15e1e86777787 Mon Sep 17 00:00:00 2001
From: David Rusu <davidrusu.me@gmail.com>
Date: Thu, 13 Jun 2024 13:43:09 -0400
Subject: [PATCH] cl: output proof tests

---
 cl/src/input.rs  |  5 ++-
 cl/src/note.rs   |  2 +-
 cl/src/output.rs | 81 ++++++++++++++++++++++++++++++++++++++++++++++--
 3 files changed, 81 insertions(+), 7 deletions(-)

diff --git a/cl/src/input.rs b/cl/src/input.rs
index 958688f..1055f3e 100644
--- a/cl/src/input.rs
+++ b/cl/src/input.rs
@@ -114,9 +114,8 @@ mod test {
         for wrong_witness in wrong_witnesses {
             assert!(input.prove(&wrong_witness).is_err());
 
-            let wrong_note = Input::from_witness(wrong_witness.clone());
-            let wrong_proof = wrong_note.prove(&wrong_witness).unwrap();
-
+            let wrong_input = Input::from_witness(wrong_witness.clone());
+            let wrong_proof = wrong_input.prove(&wrong_witness).unwrap();
             assert!(!input.verify(&wrong_proof));
         }
     }
diff --git a/cl/src/note.rs b/cl/src/note.rs
index d6f4866..5d0446e 100644
--- a/cl/src/note.rs
+++ b/cl/src/note.rs
@@ -13,7 +13,7 @@ lazy_static! {
         crypto::hash_to_curve(b"NOMOS_CL_PEDERSON_COMMITMENT_BLINDING");
 }
 
-#[derive(Debug, PartialEq, Eq)]
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
 pub struct NoteCommitment([u8; 32]);
 
 #[derive(Debug, Clone, PartialEq, Eq)]
diff --git a/cl/src/output.rs b/cl/src/output.rs
index 111627f..bd5826b 100644
--- a/cl/src/output.rs
+++ b/cl/src/output.rs
@@ -1,15 +1,18 @@
 use jubjub::{ExtendedPoint, Scalar};
 
 use crate::{
+    error::Error,
     note::{Note, NoteCommitment},
     nullifier::{NullifierCommitment, NullifierNonce},
 };
 
+#[derive(Debug, Clone, PartialEq, Eq)]
 pub struct Output {
     pub note_comm: NoteCommitment,
     pub balance: ExtendedPoint,
 }
 
+#[derive(Debug, Clone, PartialEq, Eq)]
 pub struct OutputWitness {
     note: Note,
     nf_pk: NullifierCommitment,
@@ -18,21 +21,93 @@ pub struct OutputWitness {
 }
 
 // as we don't have SNARKS hooked up yet, the witness will be our proof
+#[derive(Debug, Clone, PartialEq, Eq)]
 pub struct OutputProof(OutputWitness);
 
 impl Output {
-    pub fn prove(&self, w: OutputWitness) -> OutputProof {
-        OutputProof(w)
+    pub fn from_witness(w: OutputWitness) -> Self {
+        Self {
+            note_comm: w.note.commit(w.nf_pk, w.nonce),
+            balance: w.note.balance(w.balance_blinding),
+        }
+    }
+
+    pub fn prove(&self, w: &OutputWitness) -> Result<OutputProof, Error> {
+        if &Self::from_witness(w.clone()) == self {
+            Ok(OutputProof(w.clone()))
+        } else {
+            Err(Error::ProofFailed)
+        }
     }
 
     pub fn verify(&self, proof: &OutputProof) -> bool {
         // verification checks the relation
         // - note_comm == commit(note || nf_pk)
         // - balance == v * hash_to_curve(Unit) + blinding * H
-
         let witness = &proof.0;
 
         self.note_comm == witness.note.commit(witness.nf_pk, witness.nonce)
             && self.balance == witness.note.balance(witness.balance_blinding)
     }
 }
+
+#[cfg(test)]
+mod test {
+    use group::ff::Field;
+
+    use super::*;
+    use crate::{nullifier::NullifierSecret, test_util::seed_rng};
+
+    #[test]
+    fn test_output_proof() {
+        let mut rng = seed_rng(0);
+
+        let note = Note::new(10, "NMO");
+        let nf_pk = NullifierSecret::random(&mut rng).commit();
+        let nonce = NullifierNonce::random(&mut rng);
+        let balance_blinding = Scalar::random(&mut rng);
+
+        let witness = OutputWitness {
+            note,
+            nf_pk,
+            nonce,
+            balance_blinding,
+        };
+
+        let output = Output::from_witness(witness.clone());
+        let proof = output.prove(&witness).unwrap();
+
+        assert!(output.verify(&proof));
+
+        let wrong_witnesses = [
+            OutputWitness {
+                note: Note::new(11, "NMO"),
+                ..witness.clone()
+            },
+            OutputWitness {
+                note: Note::new(10, "ETH"),
+                ..witness.clone()
+            },
+            OutputWitness {
+                nf_pk: NullifierSecret::random(&mut rng).commit(),
+                ..witness.clone()
+            },
+            OutputWitness {
+                nonce: NullifierNonce::random(&mut rng),
+                ..witness.clone()
+            },
+            OutputWitness {
+                balance_blinding: Scalar::random(&mut rng),
+                ..witness.clone()
+            },
+        ];
+
+        for wrong_witness in wrong_witnesses {
+            assert!(output.prove(&wrong_witness).is_err());
+
+            let wrong_output = Output::from_witness(wrong_witness.clone());
+            let wrong_proof = wrong_output.prove(&wrong_witness).unwrap();
+            assert!(!output.verify(&wrong_proof));
+        }
+    }
+}