crypto: use result instead of unwrap when parsing (MystenLabs#618)

mwtian · Jul 28, 2022 · 3d59d57 · 3d59d57
1 parent 2801b40
commit 3d59d57
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 8 deletions.
diff --git a/narwhal/crypto/src/secp256k1.rs b/narwhal/crypto/src/secp256k1.rs
@@ -7,7 +7,12 @@ use crate::{
 };
 use base64ct::{Base64, Encoding};
 use once_cell::sync::OnceCell;
-use rust_secp256k1::{constants, rand::rngs::OsRng, Message, PublicKey, Secp256k1, SecretKey};
+use rust_secp256k1::{
+    constants,
+    ecdsa::{RecoverableSignature, RecoveryId},
+    rand::rngs::OsRng,
+    Message, PublicKey, Secp256k1, SecretKey,
+};
 use serde::{de, Deserialize, Serialize};
 use signature::{Signature, Signer, Verifier};
 use std::{
@@ -231,14 +236,19 @@ impl<'de> Deserialize<'de> for Secp256k1Signature {
 
 impl Signature for Secp256k1Signature {
     fn from_bytes(bytes: &[u8]) -> Result<Self, signature::Error> {
-        let recovery_id = rust_secp256k1::ecdsa::RecoveryId::from_i32(bytes[64] as i32).unwrap();
-        match rust_secp256k1::ecdsa::RecoverableSignature::from_compact(&bytes[..64], recovery_id) {
-            Ok(sig) => Ok(Secp256k1Signature {
-                sig,
-                bytes: OnceCell::new(),
-            }),
-            Err(_) => Err(signature::Error::new()),
+        if bytes.len() != 65 {
+            return Err(signature::Error::new());
         }
+        RecoveryId::from_i32(bytes[64] as i32)
+            .and_then(|rec_id| {
+                RecoverableSignature::from_compact(&bytes[..64], rec_id).map(|sig| {
+                    Secp256k1Signature {
+                        sig,
+                        bytes: OnceCell::new(),
+                    }
+                })
+            })
+            .map_err(|_| signature::Error::new())
     }
 }
 

diff --git a/narwhal/crypto/src/tests/secp256k1_tests.rs b/narwhal/crypto/src/tests/secp256k1_tests.rs
@@ -80,6 +80,12 @@ fn test_public_key_recovery() {
 
 #[test]
 fn test_public_key_recovery_error() {
+    // incorrect length
+    assert!(<Secp256k1Signature as ToFromBytes>::from_bytes(&[0u8; 1]).is_err());
+
+    // invalid recovery id at index 65
+    assert!(<Secp256k1Signature as ToFromBytes>::from_bytes(&[4u8; 65]).is_err());
+
     let signature = <Secp256k1Signature as ToFromBytes>::from_bytes(&[0u8; 65]).unwrap();
     let message: &[u8] = b"Hello, world!";
     assert!(signature