From ca97de01cf40b921a550875f53923835ced94368 Mon Sep 17 00:00:00 2001
From: kafeikui <kkffkk3@protonmail.com>
Date: Thu, 23 May 2024 23:24:25 +0800
Subject: [PATCH] improve DKGNodeError

---
 crates/arpa-node/src/algorithm/dkg.rs | 15 +++++++++------
 crates/dkg-core/src/node.rs           | 13 +------------
 2 files changed, 10 insertions(+), 18 deletions(-)

diff --git a/crates/arpa-node/src/algorithm/dkg.rs b/crates/arpa-node/src/algorithm/dkg.rs
index 901c096..1fc6f34 100644
--- a/crates/arpa-node/src/algorithm/dkg.rs
+++ b/crates/arpa-node/src/algorithm/dkg.rs
@@ -182,12 +182,15 @@ where
                 }
             }
             Err(err) => match err {
-                DKGNodeError::NotEnoughValidShares(_, _, disqualified_node_indices) => {
-                    Ok(DKGOutput::<C>::fail(disqualified_node_indices))
-                }
-                DKGNodeError::NotEnoughJustifications(_, _, disqualified_node_indices) => {
-                    Ok(DKGOutput::<C>::fail(disqualified_node_indices))
-                }
+                DKGNodeError::DKGError(e) => match e {
+                    DKGError::NotEnoughValidShares(_, _, disqualified_node_indices) => {
+                        Ok(DKGOutput::<C>::fail(disqualified_node_indices))
+                    }
+                    DKGError::NotEnoughJustifications(_, _, disqualified_node_indices) => {
+                        Ok(DKGOutput::<C>::fail(disqualified_node_indices))
+                    }
+                    _ => Err(e.into()),
+                },
                 _ => Err(err.into()),
             },
         }
diff --git a/crates/dkg-core/src/node.rs b/crates/dkg-core/src/node.rs
index de1a338..c18414b 100644
--- a/crates/dkg-core/src/node.rs
+++ b/crates/dkg-core/src/node.rs
@@ -23,12 +23,6 @@ pub enum DKGNodeError {
     /// There was an internal error in the DKG
     #[error("DKG Error: {0}")]
     DKGError(#[from] DKGError),
-
-    #[error("only has {0}/{1} valid shares, disqualified: {2:?}")]
-    NotEnoughValidShares(usize, usize, Vec<u32>),
-
-    #[error("only has {0}/{1} required justifications, disqualified: {2:?}")]
-    NotEnoughJustifications(usize, usize, Vec<u32>),
 }
 
 /// Phase2 can either be successful or require going to Phase 3.
@@ -115,12 +109,7 @@ where
 
                 Ok(next)
             }
-            Err(e) => match e {
-                DKGError::NotEnoughValidShares(got, required, disqualified_node_indices) => Err(
-                    DKGNodeError::NotEnoughValidShares(got, required, disqualified_node_indices),
-                ),
-                _ => Err(DKGNodeError::DKGError(e)),
-            },
+            Err(e) => Err(DKGNodeError::DKGError(e)),
         }
     }
 }