[Cosmos]: Add all Protobuf messages

* Handle `SignDirect` message
trustwallet · Nov 9, 2023 · baf21cb · baf21cb
1 parent 7430ca2
commit baf21cb
Show file tree

Hide file tree

Showing 8 changed files with 385 additions and 21 deletions.
diff --git a/rust/tw_cosmos_sdk/src/modules/serializer/protobuf_serializer.rs b/rust/tw_cosmos_sdk/src/modules/serializer/protobuf_serializer.rs
@@ -14,6 +14,7 @@ use crate::transaction::{
 };
 use std::marker::PhantomData;
 use tw_coin_entry::error::SigningResult;
+use tw_memory::Data;
 use tw_proto::serialize;
 
 pub fn build_coin(coin: &Coin) -> base_proto::Coin {
@@ -28,6 +29,13 @@ pub struct ProtobufSerializer<Context> {
     _phantom: PhantomData<Context>,
 }
 
+pub struct SignDirectArgs {
+    pub tx_body: Data,
+    pub auth_info: Data,
+    pub chain_id: String,
+    pub account_number: u64,
+}
+
 impl<Context> ProtobufSerializer<Context>
 where
     Context: CosmosContext,
@@ -50,6 +58,14 @@ where
         })
     }
 
+    pub fn build_direct_signed_tx(args: &SignDirectArgs, signature: Data) -> tx_proto::TxRaw {
+        tx_proto::TxRaw {
+            body_bytes: args.tx_body.clone(),
+            auth_info_bytes: args.auth_info.clone(),
+            signatures: vec![signature],
+        }
+    }
+
     /// Serializes an unsigned transaction into the Cosmos [`tx_proto::SignDoc`] message.
     /// [`tx_proto::SignDoc`] is used to generate a transaction prehash and sign it.
     pub fn build_sign_doc(
@@ -70,6 +86,15 @@ where
         })
     }
 
+    pub fn build_direct_sign_doc(args: &SignDirectArgs) -> tx_proto::SignDoc {
+        tx_proto::SignDoc {
+            body_bytes: args.tx_body.clone(),
+            auth_info_bytes: args.auth_info.clone(),
+            chain_id: args.chain_id.clone(),
+            account_number: args.account_number,
+        }
+    }
+
     pub fn build_auth_info(
         signer: &SignerInfo<Context::PublicKey>,
         fee: &Fee<Context::Address>,

diff --git a/rust/tw_cosmos_sdk/src/modules/signer/mod.rs b/rust/tw_cosmos_sdk/src/modules/signer/mod.rs
@@ -47,6 +47,10 @@ where
         coin: &dyn CoinContext,
         mut input: Proto::SigningInput<'_>,
     ) -> SigningResult<Proto::SigningOutput<'static>> {
+        if let Ok(Some(_)) = TxBuilder::<Context>::try_sign_direct_args(&input) {
+            return Self::sign_as_protobuf_direct(coin, input);
+        }
+
         let private_key = Context::PrivateKey::try_from(&input.private_key)?;
         let public_key = Context::PublicKey::from_private_key(coin, private_key.as_ref())?;
         let broadcast_mode = Self::broadcast_mode(input.mode);
@@ -73,6 +77,39 @@ where
         })
     }
 
+    pub fn sign_as_protobuf_direct(
+        coin: &dyn CoinContext,
+        input: Proto::SigningInput<'_>,
+    ) -> SigningResult<Proto::SigningOutput<'static>> {
+        let private_key = Context::PrivateKey::try_from(&input.private_key)?;
+        let public_key = Context::PublicKey::from_private_key(coin, private_key.as_ref())?;
+
+        let sign_direct_args = TxBuilder::<Context>::try_sign_direct_args(&input)?
+            // This function must be called only when there is a `SignDirect` message only.
+            .ok_or(SigningError(SigningErrorType::Error_internal))?;
+
+        let signature_data =
+            ProtobufSigner::<Context>::sign_direct(&private_key, &sign_direct_args)?;
+        let signed_tx_raw = ProtobufSerializer::<Context>::build_direct_signed_tx(
+            &sign_direct_args,
+            signature_data.clone(),
+        );
+
+        let broadcast_mode = Self::broadcast_mode(input.mode);
+        let broadcast_tx = BroadcastMsg::raw(broadcast_mode, &signed_tx_raw).to_json_string();
+
+        let signature_json =
+            JsonSerializer::<Context>::serialize_signature(&public_key, signature_data.clone())
+                .to_json_string();
+
+        Ok(Proto::SigningOutput {
+            signature: Cow::from(signature_data),
+            signature_json: Cow::from(signature_json),
+            serialized: Cow::from(broadcast_tx),
+            ..Proto::SigningOutput::default()
+        })
+    }
+
     pub fn sign_as_json(
         coin: &dyn CoinContext,
         mut input: Proto::SigningInput<'_>,

diff --git a/rust/tw_cosmos_sdk/src/modules/signer/protobuf_signer.rs b/rust/tw_cosmos_sdk/src/modules/signer/protobuf_signer.rs
@@ -6,8 +6,8 @@
 
 use crate::context::CosmosContext;
 use crate::hasher::CosmosHasher;
-use crate::modules::serializer::protobuf_serializer::ProtobufSerializer;
-use crate::private_key::CosmosPrivateKey;
+use crate::modules::serializer::protobuf_serializer::{ProtobufSerializer, SignDirectArgs};
+use crate::private_key::{CosmosPrivateKey, SignatureData};
 use crate::public_key::ProtobufPublicKey;
 use crate::transaction::{SignedTransaction, UnsignedTransaction};
 use std::marker::PhantomData;
@@ -34,4 +34,15 @@ where
 
         Ok(unsigned.into_signed(signature_data))
     }
+
+    pub fn sign_direct(
+        private_key: &Context::PrivateKey,
+        args: &SignDirectArgs,
+    ) -> SigningResult<SignatureData> {
+        let tx_to_sign = ProtobufSerializer::<Context>::build_direct_sign_doc(args);
+        let encoded_tx = serialize(&tx_to_sign)?;
+
+        let hash_to_sign = Context::TxHasher::hash_sign_doc(&encoded_tx);
+        private_key.sign_tx_hash(&hash_to_sign)
+    }
 }