encoding treasury addresss

trustwallet · Oct 7, 2024 · e1947ea · e1947ea
1 parent 01293ec
commit e1947ea
Show file tree

Hide file tree

Showing 3 changed files with 82 additions and 7 deletions.
diff --git a/rust/chains/tw_pactus/src/address.rs b/rust/chains/tw_pactus/src/address.rs
@@ -16,10 +16,12 @@ use tw_memory::Data;
 use crate::encode;
 
 const HRP: &str = "pc";
+const TREASURY_ADDRESS_STRING: &str = "000000000000000000000000000000000000000000";
 
 /// Enum for Pactus address types.
 #[derive(Debug, Clone, PartialEq)]
 pub enum AddressType {
+    Treasury = 0,
     Validator = 1,
     BlsAccount = 2,
     Ed25519Account = 3,
@@ -58,6 +60,10 @@ impl Address {
             pub_hash: pub_hash.try_into().map_err(|_| AddressError::Internal)?,
         })
     }
+
+    pub fn is_treasury(&self) -> bool {
+        self.addr_type == AddressType::Treasury && self.pub_hash == [0; 20]
+    }
 }
 
 impl CoinAddress for Address {
@@ -74,6 +80,10 @@ impl CoinAddress for Address {
 // Pactus addresses are encoded into a string format using the Bech32m encoding scheme.
 impl fmt::Display for Address {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        if self.is_treasury() {
+            return f.write_str(TREASURY_ADDRESS_STRING);
+        }
+
         let mut b32 = Vec::with_capacity(33);
 
         b32.push(bech32::u5::try_from_u8(self.addr_type.clone() as u8).map_err(|_| fmt::Error)?);
@@ -84,10 +94,20 @@ impl fmt::Display for Address {
 
 impl encode::Encodable for Address {
     fn encode(&self, stream: &mut encode::stream::Stream) {
+        if self.is_treasury() {
+            stream.append(&0u8);
+
+            return;
+        }
+
         stream.append_raw_slice(&self.data());
     }
 
     fn encoded_size(&self) -> usize {
+        if self.is_treasury() {
+            return 1;
+        }
+
         21
     }
 }
@@ -96,12 +116,23 @@ impl FromStr for Address {
     type Err = AddressError;
 
     fn from_str(s: &str) -> Result<Self, AddressError> {
+        if s == TREASURY_ADDRESS_STRING {
+            return Ok(Address {
+                addr_type: AddressType::Treasury,
+                pub_hash: [0; 20],
+            });
+        }
+
         let (hrp, b32, _variant) = bech32::decode(s).map_err(|_| AddressError::FromBech32Error)?;
 
         if hrp != HRP {
             return Err(AddressError::InvalidHrp);
         }
 
+        if b32.len() != 33 {
+            return Err(AddressError::InvalidInput)
+        }
+
         let addr_type = AddressType::from(b32[0].to_u8());
         if addr_type == AddressType::Invalid {
             return Err(AddressError::InvalidInput);
@@ -119,12 +150,35 @@ impl FromStr for Address {
 
 #[cfg(test)]
 mod test {
-    use encode::{stream::Stream, Encodable};
+    use encode::{stream::{self, Stream}, Encodable};
     use tw_encoding::hex::DecodeHex;
     use tw_keypair::ed25519::sha512::PrivateKey;
 
     use super::*;
 
+    #[test]
+    fn test_treasury_address_encoding() {
+        let addr = Address::from_str(TREASURY_ADDRESS_STRING).unwrap();
+        assert!(addr.is_treasury());
+
+        let mut  stream = Stream::new();
+        addr.encode(&mut stream);
+        assert_eq!( stream.out(), [0x00]);
+        assert_eq!( addr.encoded_size(), 1);
+    }
+
+
+    #[test]
+    fn test_address_encoding() {
+        let addr = Address::from_str("pc1rqqqsyqcyq5rqwzqfpg9scrgwpuqqzqsr36kkra").unwrap();
+        assert!(!addr.is_treasury());
+
+        let mut  stream = Stream::new();
+        addr.encode(&mut stream);
+        assert_eq!( stream.out(), "03000102030405060708090a0b0c0d0e0f00010203".decode_hex().unwrap());
+        assert_eq!( addr.encoded_size(), 21);
+    }
+
     #[test]
     fn test_address_string() {
         struct TestCase<'a> {
@@ -136,6 +190,12 @@ mod test {
 
         // Define a list of test cases for encoding and decoding
         let test_cases = vec![
+            TestCase {
+                name: "Type Treasury (0)",
+                addr_type: AddressType::Treasury,
+                pub_hash: "0000000000000000000000000000000000000000",
+                expected_addr: TREASURY_ADDRESS_STRING,
+            },
             TestCase {
                 name: "Type Validator (1)",
                 addr_type: AddressType::Validator,

diff --git a/rust/tw_any_coin/tests/chains/pactus/pactus_address.rs b/rust/tw_any_coin/tests/chains/pactus/pactus_address.rs
@@ -10,25 +10,40 @@ use tw_coin_registry::coin_type::CoinType;
 
 #[test]
 fn test_pactus_address_derive() {
-    test_address_derive(CoinType::Pactus, "PRIVATE_KEY", "EXPECTED ADDRESS");
+    test_address_derive(CoinType::Pactus, "2134ae97465505dfd5a1fd05a8a0f146209c601eb3f1b0363b4cfe4b47ba1ab4", "pc1r7jkvfnegf0rf5ua05fzu9krjhjxcrrygl3v4nl");
 }
 
 #[test]
 fn test_pactus_address_normalization() {
-    test_address_normalization(CoinType::Pactus, "DENORMALIZED", "EXPECTED");
+    test_address_normalization(CoinType::Pactus, "pc1r7jkvfnegf0rf5ua05fzu9krjhjxcrrygl3v4nl", "pc1r7jkvfnegf0rf5ua05fzu9krjhjxcrrygl3v4nl");
 }
 
+
 #[test]
 fn test_pactus_address_is_valid() {
-    test_address_valid(CoinType::Pactus, "VALID ADDRESS");
+    test_address_valid(CoinType::Pactus, "000000000000000000000000000000000000000000");
+    test_address_valid(CoinType::Pactus, "pc1p0hrct7eflrpw4ccrttxzs4qud2axex4dcdzdfr");
+    test_address_valid(CoinType::Pactus, "pc1zzqkzzu4vyddss052as6c37qrdcfptegquw826x");
+    test_address_valid(CoinType::Pactus, "pc1r7jkvfnegf0rf5ua05fzu9krjhjxcrrygl3v4nl");
 }
 
 #[test]
 fn test_pactus_address_invalid() {
-    test_address_invalid(CoinType::Pactus, "INVALID ADDRESS");
+    test_address_invalid(CoinType::Pactus, "");
+    test_address_invalid(CoinType::Pactus, "00");
+    test_address_invalid(CoinType::Pactus, "not_proper_encoded");
+    test_address_invalid(CoinType::Pactus, "pc1ioiooi");
+    test_address_invalid(CoinType::Pactus, "pc19p72rf");
+    test_address_invalid(CoinType::Pactus, "qc1z0hrct7eflrpw4ccrttxzs4qud2axex4dh8zz75");
+    test_address_invalid(CoinType::Pactus, "pc1p0hrct7eflrpw4ccrttxzs4qud2axex4dg8xaf5"); // invalid checksum
+    test_address_invalid(CoinType::Pactus, "pc1p0hrct7eflrpw4ccrttxzs4qud2axexs2dhdk8"); // invalid length
+    test_address_invalid(CoinType::Pactus, "pc1y0hrct7eflrpw4ccrttxzs4qud2axex4dksmred"); // invalid type
 }
 
 #[test]
 fn test_pactus_address_get_data() {
-    test_address_get_data(CoinType::Pactus, "ADDRESS", "HEX(DATA)");
+    test_address_get_data(CoinType::Pactus, "000000000000000000000000000000000000000000", "000000000000000000000000000000000000000000");
+    test_address_get_data(CoinType::Pactus, "pc1p0hrct7eflrpw4ccrttxzs4qud2axex4dcdzdfr", "017dc785fb29f8c2eae3035acc28541c6aba6c9aad");
+    test_address_get_data(CoinType::Pactus, "pc1zzqkzzu4vyddss052as6c37qrdcfptegquw826x", "02102c2172ac235b083e8aec3588f8036e1215e500");
+    test_address_get_data(CoinType::Pactus, "pc1r7jkvfnegf0rf5ua05fzu9krjhjxcrrygl3v4nl", "03f4acc4cf284bc69a73afa245c2d872bc8d818c88");
 }
diff --git a/rust/tw_tests/tests/chains/aptos/aptos_address.rs b/rust/tw_tests/tests/chains/aptos/aptos_address.rs
@@ -12,7 +12,7 @@ fn test_aptos_address_normalization() {
     test_address_normalization(
         CoinType::Aptos,
         "0xf3d7f364dd7705824a5ebda9c7aab6cb3fc7bb5b58718249f12defec240b36cc",
-        "0xf3d7f364dd7705824a5ebda9c7aab6cb3fc7bb5b58718249f12defec240b36cc",
+          "0xf3d7f364dd7705824a5ebda9c7aab6cb3fc7bb5b58718249f12defec240b36cc",
     );
 }