From ec6336fadcfd863a8d7c6e3fbed1147bcfda3751 Mon Sep 17 00:00:00 2001
From: quake wang <quake.wang@gmail.com>
Date: Sun, 14 Apr 2019 12:59:37 +0900
Subject: [PATCH] fix: CuckooEngine verify invalid length proof should not
 panic

---
 pow/src/cuckoo.rs | 25 ++++++++++++++++++++-----
 1 file changed, 20 insertions(+), 5 deletions(-)

diff --git a/pow/src/cuckoo.rs b/pow/src/cuckoo.rs
index 1e642fde92..1d340d3c5b 100644
--- a/pow/src/cuckoo.rs
+++ b/pow/src/cuckoo.rs
@@ -61,6 +61,9 @@ impl PowEngine for CuckooEngine {
 
     #[inline]
     fn verify(&self, _number: BlockNumber, message: &[u8], proof: &[u8]) -> bool {
+        if proof.len() != self.cuckoo.cycle_length << 2 {
+            return false;
+        }
         let mut proof_u32 = vec![0u32; self.cuckoo.cycle_length];
         LittleEndian::read_u32_into(&proof, &mut proof_u32);
         self.cuckoo.verify(message, &proof_u32)
@@ -69,7 +72,7 @@ impl PowEngine for CuckooEngine {
     #[inline]
     fn solve(&self, _number: BlockNumber, message: &[u8]) -> Option<Vec<u8>> {
         self.cuckoo.solve(message).map(|proof| {
-            let mut proof_u8 = vec![0u8; self.cuckoo.cycle_length * 4];
+            let mut proof_u8 = vec![0u8; self.cuckoo.cycle_length << 2];
             LittleEndian::write_u32_into(&proof, &mut proof_u8);
             proof_u8
         })
@@ -299,14 +302,17 @@ impl Cuckoo {
 
 #[cfg(test)]
 mod test {
-    use super::Cuckoo;
+    use super::*;
 
     use proptest::{collection::size_range, prelude::*};
 
     fn _cuckoo_solve(message: &[u8]) -> Result<(), TestCaseError> {
-        let cuckoo = Cuckoo::new(6, 8);
-        if let Some(proof) = cuckoo.solve(message) {
-            prop_assert!(cuckoo.verify(message, &proof));
+        let engine = CuckooEngine::new(CuckooParams {
+            edge_bits: 6,
+            cycle_length: 8,
+        });
+        if let Some(proof) = engine.solve(0, message) {
+            prop_assert!(engine.verify(0, message, &proof));
         }
         Ok(())
     }
@@ -366,4 +372,13 @@ mod test {
             assert!(cuckoo.verify(message, proof));
         }
     }
+
+    #[test]
+    fn verify_invalid_length_should_not_panic() {
+        let engine = CuckooEngine::new(CuckooParams {
+            edge_bits: 6,
+            cycle_length: 8,
+        });
+        assert!(!engine.verify(0, &[0, 1], &[0, 1]));
+    }
 }