feat: Make CurveType accessible within Base58PublicKey

near-sdk/src/json_types/public_key.rs

@@ -1,19 +1,29 @@
-use borsh::{BorshDeserialize, BorshSerialize};
+use borsh::{maybestd::io, BorshDeserialize, BorshSerialize};
 use bs58::decode::Error as B58Error;
-use serde::{Deserialize, Serialize};
 use std::convert::TryFrom;
 
 /// PublicKey curve
-#[derive(
-    Debug, Clone, Copy, Serialize, Deserialize, PartialEq, BorshDeserialize, BorshSerialize,
-)]
+#[derive(Debug, Clone, Copy, PartialOrd, Ord, Eq, PartialEq, BorshDeserialize, BorshSerialize)]
+#[repr(u8)]
 pub enum CurveType {
     ED25519 = 0,
     SECP256K1 = 1,
 }
 
+impl TryFrom<u8> for CurveType {
+    type Error = ParsePublicKeyError;
+
+    fn try_from(value: u8) -> Result<Self, Self::Error> {
+        match value {
+            0 => Ok(Self::ED25519),
+            1 => Ok(Self::SECP256K1),
+            _ => Err(ParsePublicKeyError { kind: ParsePublicKeyErrorKind::UnknownCurve }),
+        }
+    }
+}
+
 impl TryFrom<String> for CurveType {
-    type Error = Box<dyn std::error::Error>;
+    type Error = ParsePublicKeyError;
 
     fn try_from(value: String) -> Result<Self, Self::Error> {
         Ok(value.parse::<Self>()?)
@@ -50,8 +60,11 @@ impl std::str::FromStr for CurveType {
 ///             .parse()
 ///             .unwrap();
 /// ```
-#[derive(Debug, Clone, PartialEq, PartialOrd, Ord, Eq, BorshDeserialize, BorshSerialize)]
-pub struct Base58PublicKey(pub Vec<u8>);
+#[derive(Debug, Clone, PartialEq, PartialOrd, Ord, Eq)]
+pub struct Base58PublicKey {
+    curve: CurveType,
+    data: Vec<u8>,
+}
 
 impl Base58PublicKey {
     fn split_key_type_data(value: &str) -> Result<(CurveType, &str), ParsePublicKeyError> {
@@ -63,23 +76,45 @@ impl Base58PublicKey {
             Ok((CurveType::ED25519, value))
         }
     }
+
+    fn from_parts(curve: CurveType, data: Vec<u8>) -> Result<Self, ParsePublicKeyError> {
+        let expected_length = match curve {
+            CurveType::ED25519 => 32,
+            CurveType::SECP256K1 => 64,
+        };
+        if data.len() != expected_length {
+            return Err(ParsePublicKeyError {
+                kind: ParsePublicKeyErrorKind::InvalidLength(data.len()),
+            });
+        }
+
+        Ok(Self { curve, data })
+    }
+
+    /// Get info about the CurveType for this public key
+    pub fn curve(&self) -> CurveType {
+        self.curve
+    }
 }
 
 impl From<Base58PublicKey> for Vec<u8> {
-    fn from(v: Base58PublicKey) -> Vec<u8> {
-        v.0
+    fn from(mut v: Base58PublicKey) -> Vec<u8> {
+        v.data.insert(0, v.curve as u8);
+        v.data
     }
 }
 
 impl TryFrom<Vec<u8>> for Base58PublicKey {
-    type Error = Box<dyn std::error::Error>;
+    type Error = ParsePublicKeyError;
 
-    fn try_from(v: Vec<u8>) -> Result<Self, Self::Error> {
-        match v.len() {
-            33 if v[0] == 0 => Ok(Self(v)),
-            65 if v[0] == 1 => Ok(Self(v)),
-            _ => Err("Invalid public key".into()),
+    fn try_from(mut data: Vec<u8>) -> Result<Self, Self::Error> {
+        if data.len() == 0 {
+            return Err(ParsePublicKeyError {
+                kind: ParsePublicKeyErrorKind::InvalidLength(data.len()),
+            });
         }
+        let curve = CurveType::try_from(data.remove(0))?;
+        Self::from_parts(curve, data)
     }
 }
 
@@ -97,31 +132,34 @@ impl<'de> serde::Deserialize<'de> for Base58PublicKey {
     where
         D: serde::Deserializer<'de>,
     {
-        let s: String = Deserialize::deserialize(deserializer)?;
+        let s: String = serde::Deserialize::deserialize(deserializer)?;
         s.parse::<Base58PublicKey>().map_err(serde::de::Error::custom)
     }
 }
 
 impl From<&Base58PublicKey> for String {
     fn from(str_public_key: &Base58PublicKey) -> Self {
-        match str_public_key.0[0] {
-            0 => "ed25519:".to_string() + &bs58::encode(&str_public_key.0[1..]).into_string(),
-            1 => "secp256k1:".to_string() + &bs58::encode(&str_public_key.0[1..]).into_string(),
-            _ => panic!("Unexpected curve"),
+        match str_public_key.curve {
+            CurveType::ED25519 => {
+                ["ed25519:", &bs58::encode(&str_public_key.data).into_string()].concat()
+            }
+            CurveType::SECP256K1 => {
+                ["secp256k1:", &bs58::encode(&str_public_key.data).into_string()].concat()
+            }
         }
     }
 }
 
 impl TryFrom<String> for Base58PublicKey {
-    type Error = Box<dyn std::error::Error>;
+    type Error = ParsePublicKeyError;
 
     fn try_from(value: String) -> Result<Self, Self::Error> {
         Self::try_from(value.as_str())
     }
 }
 
 impl TryFrom<&str> for Base58PublicKey {
-    type Error = Box<dyn std::error::Error>;
+    type Error = ParsePublicKeyError;
 
     fn try_from(value: &str) -> Result<Self, Self::Error> {
         Ok(value.parse::<Self>()?)
@@ -132,24 +170,31 @@ impl std::str::FromStr for Base58PublicKey {
     type Err = ParsePublicKeyError;
 
     fn from_str(value: &str) -> Result<Self, Self::Err> {
-        let (key_type, key_data) = Base58PublicKey::split_key_type_data(&value)?;
-        let expected_length = match key_type {
+        let (curve, key_data) = Base58PublicKey::split_key_type_data(&value)?;
+        let data = bs58::decode(key_data).into_vec()?;
+        Self::from_parts(curve, data)
+    }
+}
+
+impl BorshDeserialize for Base58PublicKey {
+    fn deserialize(buf: &mut &[u8]) -> io::Result<Self> {
+        let buf = Vec::<u8>::deserialize(buf)?;
+        Self::try_from(buf).map_err(|err| io::Error::new(io::ErrorKind::Other, err))
+    }
+}
+
+impl BorshSerialize for Base58PublicKey {
+    fn serialize<W: io::Write>(&self, writer: &mut W) -> io::Result<()> {
+        let len = match self.curve {
             CurveType::ED25519 => 32,
             CurveType::SECP256K1 => 64,
         };
-        let data = bs58::decode(key_data).into_vec()?;
-        if data.len() != expected_length {
-            return Err(ParsePublicKeyError {
-                kind: ParsePublicKeyErrorKind::InvalidLength(data.len()),
-            });
-        }
-        let mut res = Vec::with_capacity(1 + expected_length);
-        match key_type {
-            CurveType::ED25519 => res.push(0),
-            CurveType::SECP256K1 => res.push(1),
-        };
-        res.extend(data);
-        Ok(Self(res))
+
+        let mut data = Vec::with_capacity(1 + len);
+        data.push(self.curve as u8);
+        data.extend(&self.data);
+
+        BorshSerialize::serialize(&data, writer)
     }
 }
 
@@ -190,25 +235,25 @@ mod tests {
     use super::*;
     use std::convert::TryInto;
 
-    fn binary_key() -> Vec<u8> {
-        let mut binary_key = vec![0];
-        binary_key.extend(
+    fn expected_key() -> Base58PublicKey {
+        let mut key = vec![CurveType::ED25519 as u8];
+        key.extend(
             bs58::decode("6E8sCci9badyRkXb3JoRpBj5p8C6Tw41ELDZoiihKEtp").into_vec().unwrap(),
         );
-        binary_key
+        key.try_into().unwrap()
     }
 
     #[test]
     fn test_public_key_deser() {
         let key: Base58PublicKey =
             serde_json::from_str("\"ed25519:6E8sCci9badyRkXb3JoRpBj5p8C6Tw41ELDZoiihKEtp\"")
                 .unwrap();
-        assert_eq!(key.0, binary_key());
+        assert_eq!(key, expected_key());
     }
 
     #[test]
     fn test_public_key_ser() {
-        let key: Base58PublicKey = binary_key().try_into().unwrap();
+        let key: Base58PublicKey = expected_key();
         let actual: String = serde_json::to_string(&key).unwrap();
         assert_eq!(actual, "\"ed25519:6E8sCci9badyRkXb3JoRpBj5p8C6Tw41ELDZoiihKEtp\"");
     }
@@ -217,13 +262,32 @@ mod tests {
     fn test_public_key_from_str() {
         let key = Base58PublicKey::try_from("ed25519:6E8sCci9badyRkXb3JoRpBj5p8C6Tw41ELDZoiihKEtp")
             .unwrap();
-        assert_eq!(key.0, binary_key());
+        assert_eq!(key, expected_key());
     }
 
     #[test]
     fn test_public_key_to_string() {
-        let key: Base58PublicKey = binary_key().try_into().unwrap();
+        let key: Base58PublicKey = expected_key();
         let actual: String = String::try_from(&key).unwrap();
         assert_eq!(actual, "ed25519:6E8sCci9badyRkXb3JoRpBj5p8C6Tw41ELDZoiihKEtp");
     }
+
+    #[test]
+    fn test_public_key_borsh_format_change() {
+        // Original struct to reference Borsh serialization from
+        #[derive(BorshSerialize, BorshDeserialize)]
+        struct Base58PublicKeyRef(Vec<u8>);
+
+        let mut data = vec![CurveType::ED25519 as u8];
+        data.extend(
+            bs58::decode("6E8sCci9badyRkXb3JoRpBj5p8C6Tw41ELDZoiihKEtp").into_vec().unwrap(),
+        );
+
+        // Test internal serialization of Vec<u8> is the same:
+        let old_key = Base58PublicKeyRef(data.clone());
+        let old_encoded_key = old_key.try_to_vec().unwrap();
+        let new_key: Base58PublicKey = data.try_into().unwrap();
+        let new_encoded_key = new_key.clone().try_to_vec().unwrap();
+        assert_eq!(old_encoded_key, new_encoded_key);
+    }
 }