From 8b9e5c830ff6e4d847d877ffb3092a1f8e50207d Mon Sep 17 00:00:00 2001
From: ii-cruz <ii.pintocruz@gmail.com>
Date: Wed, 27 Nov 2024 17:25:31 -0600
Subject: [PATCH] Make errors more informative

---
 zkp-component/src/proof_utils.rs  |  4 ++--
 zkp-component/src/types.rs        | 10 ++++++++--
 zkp-component/src/zkp_requests.rs | 18 +++++++++---------
 3 files changed, 19 insertions(+), 13 deletions(-)

diff --git a/zkp-component/src/proof_utils.rs b/zkp-component/src/proof_utils.rs
index 0de3ea2c13..f86522326c 100644
--- a/zkp-component/src/proof_utils.rs
+++ b/zkp-component/src/proof_utils.rs
@@ -48,10 +48,10 @@ pub(crate) fn get_proof_macro_blocks(
         let new_block = blockchain
             .read()
             .get_block_at(block_number, true)
-            .map_err(|_| Error::InvalidBlock)?;
+            .map_err(|_| Error::UnknownBlock(block_number))?;
 
         if !new_block.is_election() {
-            return Err(Error::InvalidBlock);
+            return Err(Error::NotElectionBlock);
         }
         new_block.unwrap_macro()
     };
diff --git a/zkp-component/src/types.rs b/zkp-component/src/types.rs
index a62c3f5d3f..b9e09fadf1 100644
--- a/zkp-component/src/types.rs
+++ b/zkp-component/src/types.rs
@@ -152,8 +152,14 @@ pub enum Error {
     #[error("Nano Zkp Error: {0}")]
     NanoZKP(#[from] NanoZKPError),
 
-    #[error("Proof's blocks are not valid")]
-    InvalidBlock,
+    #[error("Proof's blocks is not an election block")]
+    NotElectionBlock,
+
+    #[error("Proof's block is not known: {0}")]
+    UnknownBlock(u32),
+
+    #[error("Response was invalid: requested_block: {0}")]
+    InvalidResponse(bool),
 
     #[error("Outdated proof")]
     OutdatedProof,
diff --git a/zkp-component/src/zkp_requests.rs b/zkp-component/src/zkp_requests.rs
index 4d55b6edb5..c5cab890fc 100644
--- a/zkp-component/src/zkp_requests.rs
+++ b/zkp-component/src/zkp_requests.rs
@@ -119,18 +119,18 @@ impl<N: Network + 'static> Stream for ZKPRequests<N> {
         while let Poll::Ready(result) = self.zkp_request_results.poll_next_unpin(cx) {
             match result {
                 Some((peer_id, request_election_block, response_channel, result)) => match result {
-                    Ok(RequestZKPResponse::Proof(proof, mut election_block)) => {
+                    Ok(RequestZKPResponse::Proof(proof, election_block)) => {
                         // Check that the response is in-line with whether we asked for the election block or not.
-                        if request_election_block {
-                            if election_block.is_none() {
-                                if let Some(tx) = response_channel {
-                                    let _ = tx.send(Err(Error::InvalidBlock));
-                                }
-                                continue;
+                        if (request_election_block && election_block.is_none())
+                            || (!request_election_block && !election_block.is_none())
+                        {
+                            if let Some(tx) = response_channel {
+                                let _ =
+                                    tx.send(Err(Error::InvalidResponse(request_election_block)));
                             }
-                        } else {
-                            election_block = None;
+                            continue;
                         }
+
                         return Poll::Ready(Some(ZKPRequestsItem {
                             peer_id,
                             proof,