chore: tron sign message v2 review problem optimization

consenlabs · Oct 17, 2024 · 2141414 · 2141414
1 parent d2685ac
commit 2141414
Show file tree

Hide file tree

Showing 7 changed files with 98 additions and 39 deletions.
diff --git a/imkey-core/ikc-proto/src/tron.proto b/imkey-core/ikc-proto/src/tron.proto
@@ -12,8 +12,8 @@ message TronTxOutput {
 message TronMessageInput {
     string message = 1;
     bool is_hex = 2;
-    bool is_tron_header = 3;
-    string version = 4;
+    string header = 3;
+    uint32 version = 4;
 }
 
 message TronMessageOutput {

diff --git a/imkey-core/ikc-wallet/coin-tron/src/signer.rs b/imkey-core/ikc-wallet/coin-tron/src/signer.rs
@@ -1,6 +1,7 @@
 use crate::address::TronAddress;
 use crate::tronapi::{TronMessageInput, TronMessageOutput, TronTxInput, TronTxOutput};
 use crate::Result;
+use anyhow::anyhow;
 use ikc_common::apdu::{Apdu, ApduCheck, Secp256k1Apdu};
 use ikc_common::constants::TRON_AID;
 use ikc_common::error::CoinError;
@@ -34,12 +35,14 @@ impl TronSigner {
         };
 
         // this code is from tron wallet
-        let header = match input.is_tron_header {
-            true => match input.version.to_uppercase().as_str() {
-                "V2" => "\x19TRON Signed Message:\n".as_bytes(),
+        let header = match input.header.to_uppercase().as_str() {
+            "TRON" => match input.version {
+                2 => "\x19TRON Signed Message:\n".as_bytes(),
                 _ => "\x19TRON Signed Message:\n32".as_bytes(),
             },
-            false => "\x19Ethereum Signed Message:\n32".as_bytes(),
+            "ETH" => "\x19Ethereum Signed Message:\n32".as_bytes(),
+            "NONE" => "\x19Ethereum Signed Message:\n32".as_bytes(),
+            _ => return Err(anyhow!("sign_message_header_type_incorrect")),
         };
         let mut msg_with_header = Vec::new();
         msg_with_header.extend(header);
@@ -203,31 +206,62 @@ mod tests {
             receiver: "".to_string(),
             sender: "TY2uroBeZ5trA9QT96aEWj32XLkAAhQ9R2".to_string(),
             fee: "".to_string(),
+            seg_wit: "".to_string(),
         };
 
         let input = TronMessageInput {
             message: "645c0b7b58158babbfa6c6cd5a48aa7340a8749176b120e8516216787a13dc76".to_string(),
             is_hex: true,
-            is_tron_header: true,
+            header: "TRON".to_string(),
+            version: 1,
         };
         let res = TronSigner::sign_message(input, &sign_param).unwrap();
         assert_eq!("16417c6489da3a88ef980bf0a42551b9e76181d03e7334548ab3cb36e7622a484482722882a29e2fe4587b95c739a68624ebf9ada5f013a9340d883f03fcf9af1b", &res.signature);
 
         let input2 = TronMessageInput {
             message: "645c0b7b58158babbfa6c6cd5a48aa7340a8749176b120e8516216787a13dc76".to_string(),
             is_hex: true,
-            is_tron_header: false,
+            header: "ETH".to_string(),
+            version: 1,
         };
         let res = TronSigner::sign_message(input2, &sign_param).unwrap();
         assert_eq!("06ff3c5f98b8e8e257f47a66ce8e953c7a7d0f96eb6687da6a98b66a36c2a725759cab3df94d014bd17760328adf860649303c68c4fa6644d9f307e2f32cc3311c", &res.signature);
 
         let input3 = TronMessageInput {
             message: "abcdef".to_string(),
             is_hex: false,
-            is_tron_header: true,
+            header: "TRON".to_string(),
+            version: 1,
         };
         let res = TronSigner::sign_message(input3, &sign_param).unwrap();
         assert_eq!("a87eb6ae7e97621b6ba2e2f70db31fe0c744c6adcfdc005044026506b70ac11a33f415f4478b6cf84af32b3b5d70a13a77e53287613449b345bb16fe012c04081b", &res.signature);
+
+        let input4 = TronMessageInput {
+            message: "hello world".to_string(),
+            is_hex: false,
+            header: "ETH".to_string(),
+            version: 1,
+        };
+        let res = TronSigner::sign_message(input4, &sign_param).unwrap();
+        assert_eq!("e14f6aab4b87af398917c8a0fd6d065029df9ecc01afbc4d789eefd6c2de1e243272d630992b470c2bbb7f52024280af9bbd2e62d96ecab333c91f527b059ffe1c", &res.signature);
+
+        let input5 = TronMessageInput {
+            message: "hello world".to_string(),
+            is_hex: false,
+            header: "TRON".to_string(),
+            version: 2,
+        };
+        let res = TronSigner::sign_message(input5, &sign_param).unwrap();
+        assert_eq!("bca12bfcc9f0e23ff1d3567c4ef04ff83ac93346d6b3062d56922cc15b7669436c1aaa6a3f1ec4013545ba7d3bb79ab4b1125159d251a910f92ea198cbc469a21c", &res.signature);
+
+        let input6 = TronMessageInput {
+            message: "hello world".to_string(),
+            is_hex: false,
+            header: "NONE".to_string(),
+            version: 1,
+        };
+        let res = TronSigner::sign_message(input6, &sign_param).unwrap();
+        assert_eq!("e14f6aab4b87af398917c8a0fd6d065029df9ecc01afbc4d789eefd6c2de1e243272d630992b470c2bbb7f52024280af9bbd2e62d96ecab333c91f527b059ffe1c", &res.signature);
     }
 
     #[test]
@@ -242,6 +276,7 @@ mod tests {
             receiver: "TDQqJsFsStSy5fjG52KuiWW7HhJGAKGJLb".to_string(),
             sender: "TY2uroBeZ5trA9QT96aEWj32XLkAAhQ9R2".to_string(),
             fee: "20 dd".to_string(),
+            seg_wit: "".to_string(),
         };
         let input = TronTxInput{
             raw_data: "0a0208312208b02efdc02638b61e40f083c3a7c92d5a65080112610a2d747970652e676f6f676c65617069732e636f6d2f70726f746f636f6c2e5472616e73666572436f6e747261637412300a1541a1e81654258bf14f63feb2e8d1380075d45b0dac1215410b3e84ec677b3e63c99affcadb91a6b4e086798f186470a0bfbfa7c92d".to_string(),

diff --git a/imkey-core/ikc-wallet/coin-tron/src/tronapi.rs b/imkey-core/ikc-wallet/coin-tron/src/tronapi.rs
@@ -17,10 +17,10 @@ pub struct TronMessageInput {
     pub message: ::prost::alloc::string::String,
     #[prost(bool, tag = "2")]
     pub is_hex: bool,
-    #[prost(bool, tag = "3")]
-    pub is_tron_header: bool,
-    #[prost(string, tag = "4")]
-    pub version: ::prost::alloc::string::String,
+    #[prost(string, tag = "3")]
+    pub header: ::prost::alloc::string::String,
+    #[prost(uint32, tag = "4")]
+    pub version: u32,
 }
 #[allow(clippy::derive_partial_eq_without_eq)]
 #[derive(Clone, PartialEq, ::prost::Message)]

diff --git a/token-core/tcx-proto/src/tron.proto b/token-core/tcx-proto/src/tron.proto
@@ -21,8 +21,8 @@ message TronTxOutput {
 // build the `TronMessageInput` and put it in the `input` field
 message TronMessageInput {
     string value = 1;
-    bool isTronHeader = 2;
-    string version = 3;
+    string header = 2;
+    uint32 version = 3;
 }
 
 message TronMessageOutput {

diff --git a/token-core/tcx-tron/src/signer.rs b/token-core/tcx-tron/src/signer.rs
@@ -7,6 +7,7 @@ use tcx_keystore::{
 use bitcoin_hashes::sha256::Hash;
 use bitcoin_hashes::Hash as TraitHash;
 
+use anyhow::anyhow;
 use tcx_common::{keccak256, FromHex, ToHex};
 
 // http://jsoneditoronline.org/index.html?id=2b86a8503ba641bebed73f32b4ac9c42
@@ -65,13 +66,16 @@ impl TraitMessageSigner<TronMessageInput, TronMessageOutput> for Keystore {
     ) -> Result<TronMessageOutput> {
         let data = Vec::from_hex_auto(&message.value)?;
 
-        let header = match message.is_tron_header {
-            true => match message.version.to_uppercase().as_str() {
-                "V2" => "\x19TRON Signed Message:\n".as_bytes(),
+        let header = match message.header.to_uppercase().as_str() {
+            "TRON" => match message.version {
+                2 => "\x19TRON Signed Message:\n".as_bytes(),
                 _ => "\x19TRON Signed Message:\n32".as_bytes(),
             },
-            false => "\x19Ethereum Signed Message:\n32".as_bytes(),
+            "ETH" => "\x19Ethereum Signed Message:\n32".as_bytes(),
+            "NONE" => "\x19Ethereum Signed Message:\n32".as_bytes(),
+            _ => return Err(anyhow!("sign_message_header_type_incorrect")),
         };
+
         let to_hash = [header, &data].concat();
 
         let hash = keccak256(&to_hash);
@@ -132,7 +136,8 @@ mod tests {
 
         let message = TronMessageInput {
             value: "hello world".as_bytes().to_hex(),
-            is_tron_header: true,
+            header: "TRON".to_string(),
+            version: 1,
         };
 
         let signed = keystore.sign_message(&params, &message).unwrap();
@@ -141,12 +146,33 @@ mod tests {
 
         let message = TronMessageInput {
             value: "hello world".as_bytes().to_hex(),
-            is_tron_header: false,
+            header: "NONE".to_string(),
+            version: 1,
+        };
+
+        let signed = keystore.sign_message(&params, &message).unwrap();
+
+        assert_eq!(signed.signature, "0xe14f6aab4b87af398917c8a0fd6d065029df9ecc01afbc4d789eefd6c2de1e243272d630992b470c2bbb7f52024280af9bbd2e62d96ecab333c91f527b059ffe1c");
+
+        let message = TronMessageInput {
+            value: "hello world".as_bytes().to_hex(),
+            header: "ETH".to_string(),
+            version: 1,
         };
 
         let signed = keystore.sign_message(&params, &message).unwrap();
 
         assert_eq!(signed.signature, "0xe14f6aab4b87af398917c8a0fd6d065029df9ecc01afbc4d789eefd6c2de1e243272d630992b470c2bbb7f52024280af9bbd2e62d96ecab333c91f527b059ffe1c");
+
+        let message = TronMessageInput {
+            value: "hello world".as_bytes().to_hex(),
+            header: "TRON".to_string(),
+            version: 2,
+        };
+
+        let signed = keystore.sign_message(&params, &message).unwrap();
+
+        assert_eq!(signed.signature, "0xbca12bfcc9f0e23ff1d3567c4ef04ff83ac93346d6b3062d56922cc15b7669436c1aaa6a3f1ec4013545ba7d3bb79ab4b1125159d251a910f92ea198cbc469a21c");
     }
 
     #[test]

diff --git a/token-core/tcx-tron/src/transaction.rs b/token-core/tcx-tron/src/transaction.rs
@@ -22,10 +22,10 @@ pub struct TronTxOutput {
 pub struct TronMessageInput {
     #[prost(string, tag = "1")]
     pub value: ::prost::alloc::string::String,
-    #[prost(bool, tag = "2")]
-    pub is_tron_header: bool,
-    #[prost(string, tag = "3")]
-    pub version: ::prost::alloc::string::String,
+    #[prost(string, tag = "2")]
+    pub header: ::prost::alloc::string::String,
+    #[prost(uint32, tag = "3")]
+    pub version: u32,
 }
 #[allow(clippy::derive_partial_eq_without_eq)]
 #[derive(Clone, PartialEq, ::prost::Message)]

diff --git a/token-core/tcx/tests/sign_test.rs b/token-core/tcx/tests/sign_test.rs
@@ -506,20 +506,20 @@ fn test_tron_sign_message() {
                 (TronMessageInput {
                     value: "0x645c0b7b58158babbfa6c6cd5a48aa7340a8749176b120e8516216787a13dc76"
                         .to_string(),
-                    is_tron_header: true,
-                    version: "V1".to_string(),
+                    header: "TRON".to_string(),
+                    version: 1,
                 }, "0x16417c6489da3a88ef980bf0a42551b9e76181d03e7334548ab3cb36e7622a484482722882a29e2fe4587b95c739a68624ebf9ada5f013a9340d883f03fcf9af1b"),
                 (TronMessageInput {
                     value: "645c0b7b58158babbfa6c6cd5a48aa7340a8749176b120e8516216787a13dc76"
                         .to_string(),
-                    is_tron_header: true,
-                    version: "V1".to_string(),
+                    header: "TRON".to_string(),
+                    version: 1,
                 }, "0x16417c6489da3a88ef980bf0a42551b9e76181d03e7334548ab3cb36e7622a484482722882a29e2fe4587b95c739a68624ebf9ada5f013a9340d883f03fcf9af1b"),
                 (TronMessageInput {
                     value: "abcdef"
                         .to_string(),
-                    is_tron_header: true,
-                    version: "V1".to_string(),
+                    header: "TRON".to_string(),
+                    version: 1,
                 }, "0x13e407627e584c821ba527d23d64163d458447dfea1c3bfc92be660aa8d093ee5cfa3881870c4c51f157828eb9d4f7fad8112761f3b51cf76c7a4a3f241033d51b"),
             ];
         for (input, expected) in input_expects {
@@ -551,12 +551,11 @@ fn test_tron_sign_message_v2() {
         let wallet = import_default_wallet();
 
         let input = TronMessageInput {
-            value: "0x645c0b7b58158babbfa6c6cd5a48aa7340a8749176b120e8516216787a13dc76"
-                .to_string(),
-            is_tron_header: true,
-            version: "V2".to_string(),
+            value: "0x645c0b7b58158babbfa6c6cd5a48aa7340a8749176b120e8516216787a13dc76".to_string(),
+            header: "TRON".to_string(),
+            version: 2,
         };
-        
+
         let tx = SignParam {
             id: wallet.id.to_string(),
             key: Some(Key::Password(TEST_PASSWORD.to_string())),
@@ -574,7 +573,6 @@ fn test_tron_sign_message_v2() {
         let sign_result = call_api("sign_msg", tx).unwrap();
         let ret: TronMessageOutput = TronMessageOutput::decode(sign_result.as_slice()).unwrap();
         assert_eq!("0x9e7a691647c02fad5fe939a50df0351a58be67b3cdd87619c37f316b913d0be92ecf190f5e0c3640d54d9be731e8ab4bea4894ca9e7267b6c86d852e5c5dd71d1c", ret.signature);
-
     });
 }
 
@@ -625,8 +623,8 @@ fn test_sign_by_dk_hd_store() {
         let wallet = import_default_wallet();
         let input = TronMessageInput {
             value: "0x645c0b7b58158babbfa6c6cd5a48aa7340a8749176b120e8516216787a13dc76".to_string(),
-            is_tron_header: true,
-            version: "V1".to_string()
+            header: "TRON".to_string(),
+            version: 1,
         };
 
         let dk_param = WalletKeyParam {