Cleanup Cetus tx_builder.rs and add tests (#1133)

Author: 0xh3rman

crates/gem_encoding/src/protobuf/decode.rs

@@ -86,7 +86,11 @@ fn read_varint(data: &[u8], position: &mut usize) -> MessageResult<u64> {
     while *position < data.len() {
         let byte = data[*position];
         *position += 1;
-        result |= u64::from(byte & 0x7f) << shift;
+        let value = u64::from(byte & 0x7f);
+        if shift == 63 && value > 1 {
+            return Err("protobuf varint overflows u64".into());
+        }
+        result |= value << shift;
         if byte & 0x80 == 0 {
             return Ok(result);
         }
@@ -186,6 +190,13 @@ mod tests {
         assert_eq!(values, vec!["test", "7"]);
     }
 
+    #[test]
+    fn test_visit_fields_rejects_varint_overflow() {
+        let message = [vec![0x08], vec![0xff; 9], vec![0x02]].concat();
+
+        assert_eq!(visit_fields(&message, |_| Ok(())).unwrap_err().to_string(), "protobuf varint overflows u64");
+    }
+
     #[derive(Debug, Default, PartialEq)]
     struct TestMessage {
         name: Option<String>,

crates/gem_encoding/src/protobuf/encode.rs

@@ -35,7 +35,13 @@ pub fn encode_bytes_field(field_number: u32, data: &[u8]) -> Vec<u8> {
         return Vec::new();
     }
 
-    [field_tag(field_number, WIRE_LENGTH_DELIMITED), encode_varint(data.len() as u64), data.to_vec()].concat()
+    let tag = field_tag(field_number, WIRE_LENGTH_DELIMITED);
+    let len = encode_varint(data.len() as u64);
+    let mut buf = Vec::with_capacity(tag.len() + len.len() + data.len());
+    buf.extend_from_slice(&tag);
+    buf.extend_from_slice(&len);
+    buf.extend_from_slice(data);
+    buf
 }
 
 pub fn encode_string_field(field_number: u32, value: &str) -> Vec<u8> {

crates/gem_sui/src/address.rs

@@ -4,6 +4,12 @@ use sui_types::Address;
 
 pub struct SuiAddress(Address);
 
+impl From<SuiAddress> for Address {
+    fn from(value: SuiAddress) -> Self {
+        value.0
+    }
+}
+
 impl AddressTrait for SuiAddress {
     fn try_parse(address: &str) -> Option<Self> {
         Address::from_str(address).ok().map(Self)

crates/gem_sui/src/rpc/proto/checkpoints.rs

@@ -83,3 +83,24 @@ pub struct CheckpointedTransactionInfo {
 proto_decode!(CheckpointedTransactionInfo {
     1 => transaction: optional_string,
 });
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use gem_encoding::protobuf::MessageDecode;
+
+    #[test]
+    fn test_checkpoint_response_wire_bytes_decode() {
+        let response = hex::decode("0a2a082a1211636865636b706f696e742d6469676573741a0a180728633a04707265762a0722050a03747831").unwrap();
+        let checkpoint = GetCheckpointResponse::decode(&response).unwrap().checkpoint.unwrap();
+        let summary = checkpoint.summary.unwrap();
+        let contents = checkpoint.contents.unwrap();
+
+        assert_eq!(checkpoint.sequence_number, Some(42));
+        assert_eq!(checkpoint.digest.as_deref(), Some("checkpoint-digest"));
+        assert_eq!(summary.epoch, Some(7));
+        assert_eq!(summary.total_network_transactions, Some(99));
+        assert_eq!(summary.previous_digest.as_deref(), Some("prev"));
+        assert_eq!(contents.transactions[0].transaction.as_deref(), Some("tx1"));
+    }
+}

crates/gem_sui/src/rpc/proto/mod.rs

@@ -37,4 +37,4 @@ pub(crate) trait WithMut: Sized {
     }
 }
 
-impl<T> WithMut for T {}
+impl<T: gem_encoding::protobuf::MessageEncode> WithMut for T {}

crates/gem_sui/src/rpc/proto/objects.rs

@@ -160,3 +160,24 @@ impl OwnerKind {
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use gem_encoding::protobuf::MessageDecode;
+
+    #[test]
+    fn test_object_response_wire_bytes_decode() {
+        let response = hex::decode("0a2b120530786162631807220a6f626a2d6469676573742a0b0801120730786f776e65723204636f696ea8067b").unwrap();
+        let object = GetObjectResponse::decode(&response).unwrap().object.unwrap();
+        let owner = object.owner.unwrap();
+
+        assert_eq!(object.object_id.as_deref(), Some("0xabc"));
+        assert_eq!(object.version, Some(7));
+        assert_eq!(object.digest.as_deref(), Some("obj-digest"));
+        assert_eq!(owner.kind(), OwnerKind::Address);
+        assert_eq!(owner.address.as_deref(), Some("0xowner"));
+        assert_eq!(object.object_type.as_deref(), Some("coin"));
+        assert_eq!(object.balance, Some(123));
+    }
+}

crates/gem_sui/src/rpc/proto/transactions.rs

@@ -265,7 +265,7 @@ proto_decode!(Event {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use gem_encoding::protobuf::{MessageDecode, encode_message_field, encode_string_field};
+    use gem_encoding::protobuf::{MessageDecode, MessageEncode, encode_message_field, encode_string_field};
 
     #[test]
     fn test_execute_transaction_response_decode() {
@@ -275,4 +275,32 @@ mod tests {
 
         assert_eq!(ExecuteTransactionResponse::decode(&response).unwrap().transaction.unwrap().digest.unwrap(), digest);
     }
+
+    #[test]
+    fn test_request_wire_bytes() {
+        let transaction = Transaction::from_transaction_bcs(vec![1, 2, 3]);
+        let read_mask = FieldMask::from_paths(["digest"]);
+        let simulate = SimulateTransactionRequest {
+            transaction: Some(transaction.clone()),
+            read_mask: Some(read_mask.clone()),
+            checks: Some(TransactionChecks::Disabled),
+        };
+        let execute = ExecuteTransactionRequest {
+            transaction: Some(transaction),
+            signatures: vec![UserSignature {
+                bcs: Some(Bcs::new("UserSignatureBytes", vec![0xaa, 0xbb])),
+                scheme: Some(0),
+            }],
+            read_mask: Some(read_mask),
+        };
+
+        assert_eq!(
+            hex::encode(simulate.encode()),
+            "0a180a160a0f5472616e73616374696f6e44617461120301020312080a066469676573741801"
+        );
+        assert_eq!(
+            hex::encode(execute.encode()),
+            "0a180a160a0f5472616e73616374696f6e446174611203010203121c0a180a12557365725369676e617475726542797465731202aabb10001a080a06646967657374"
+        );
+    }
 }

crates/gem_sui/src/tx_builder/transaction.rs

@@ -14,8 +14,7 @@ use sui_types::{Address, Identifier, TypeTag};
 const MODULE_COIN: &str = "coin";
 const FUNCTION_ZERO: &str = "zero";
 
-pub fn move_call(txb: &mut TransactionBuilder, package: impl ToString, module: &str, function: &str, type_args: &[&str], arguments: Vec<Argument>) -> Result<Argument, SuiError> {
-    let package = package.to_string();
+pub fn move_call(txb: &mut TransactionBuilder, package: Address, module: &str, function: &str, type_args: &[&str], arguments: Vec<Argument>) -> Result<Argument, SuiError> {
     let type_args = type_args
         .iter()
         .map(|value| {
@@ -25,7 +24,7 @@ pub fn move_call(txb: &mut TransactionBuilder, package: impl ToString, module: &
         })
         .collect::<Result<Vec<_>, _>>()?;
     let function = Function::new(
-        Address::from_str(&package).map_err(|err| SuiError::invalid_input(format!("Invalid Sui address {package}: {err}")))?,
+        package,
         Identifier::new(module).map_err(|err| SuiError::invalid_input(err.to_string()))?,
         Identifier::new(function).map_err(|err| SuiError::invalid_input(err.to_string()))?,
     )

crates/swapper/src/cetus_clmm/tx_builder.rs

@@ -9,6 +9,7 @@ use super::{
 use crate::{Quote, SwapperError, SwapperQuoteData, fees::ReferralFee, fees::apply_slippage_in_bp};
 use gem_sui::{
     EMPTY_ADDRESS, ESTIMATION_GAS_BUDGET, SuiClient,
+    address::SuiAddress,
     gas_budget::GAS_BUDGET_MULTIPLIER,
     is_sui_coin,
     models::{CoinAsset, TxOutput},
@@ -18,7 +19,8 @@ use gem_sui::{
         into_balance, move_call,
     },
 };
-use std::{collections::HashMap, fmt::Display, str::FromStr};
+use primitives::Address as AddressTrait;
+use std::{collections::HashMap, fmt::Display};
 use sui_transaction_builder::{Argument, ObjectInput, TransactionBuilder};
 use sui_types::{Address, Digest};
 
@@ -112,7 +114,7 @@ pub(super) fn build_batch_quote_inspect(quotes: &[(&Hop, u64)]) -> Result<Vec<u8
         let amount = txb.pure(amount_in);
         move_call(
             &mut txb,
-            CETUS_CLMM_PUBLISHED_AT,
+            cetus_clmm_publish_at(),
             MODULE_POOL,
             FUNCTION_CALCULATE_SWAP_RESULT,
             &[&hop.coin_a, &hop.coin_b],
@@ -132,20 +134,20 @@ pub(super) fn build_batch_multi_hop_quote_inspect(routes: &[(&Hop, &Hop, u64)])
         let amount = txb.pure(amount_in);
         let csr1 = move_call(
             &mut txb,
-            CETUS_CLMM_PUBLISHED_AT,
+            cetus_clmm_publish_at(),
             MODULE_POOL,
             FUNCTION_CALCULATE_SWAP_RESULT,
             &[&hop1.coin_a, &hop1.coin_b],
             vec![pool1, a2b1, by_amount_in_1, amount],
         )
         .map_err(error)?;
-        let amount2 = move_call(&mut txb, CETUS_CLMM_PUBLISHED_AT, MODULE_POOL, FUNCTION_CALCULATED_SWAP_RESULT_AMOUNT_OUT, &[], vec![csr1]).map_err(error)?;
+        let amount2 = move_call(&mut txb, cetus_clmm_publish_at(), MODULE_POOL, FUNCTION_CALCULATED_SWAP_RESULT_AMOUNT_OUT, &[], vec![csr1]).map_err(error)?;
         let pool2 = txb.object(shared_object_input(&hop2.pool_id, hop2.pool_init_version, false)?);
         let a2b2 = txb.pure(&hop2.a2b);
         let by_amount_in_2 = txb.pure(&true);
         move_call(
             &mut txb,
-            CETUS_CLMM_PUBLISHED_AT,
+            cetus_clmm_publish_at(),
             MODULE_POOL,
             FUNCTION_CALCULATE_SWAP_RESULT,
             &[&hop2.coin_a, &hop2.coin_b],
@@ -159,10 +161,10 @@ pub(super) fn build_batch_multi_hop_quote_inspect(routes: &[(&Hop, &Hop, u64)])
 pub(super) fn build_pool_id_inspect(coin_a: &str, coin_b: &str, tick_spacing: u32) -> Result<Vec<u8>, SwapperError> {
     let mut txb = TransactionBuilder::new();
     let tick = txb.pure(&tick_spacing);
-    let key = move_call(&mut txb, CETUS_CLMM_PUBLISHED_AT, MODULE_FACTORY, FUNCTION_NEW_POOL_KEY, &[coin_a, coin_b], vec![tick]).map_err(error)?;
+    let key = move_call(&mut txb, cetus_clmm_publish_at(), MODULE_FACTORY, FUNCTION_NEW_POOL_KEY, &[coin_a, coin_b], vec![tick]).map_err(error)?;
     let pools = txb.object(shared_object_input(CETUS_POOLS_REGISTRY, CETUS_SHARED_INIT_VERSION, false)?);
-    let info = move_call(&mut txb, CETUS_CLMM_PUBLISHED_AT, MODULE_FACTORY, FUNCTION_POOL_SIMPLE_INFO, &[], vec![pools, key]).map_err(error)?;
-    move_call(&mut txb, CETUS_CLMM_PUBLISHED_AT, MODULE_FACTORY, FUNCTION_POOL_ID, &[], vec![info]).map_err(error)?;
+    let info = move_call(&mut txb, cetus_clmm_publish_at(), MODULE_FACTORY, FUNCTION_POOL_SIMPLE_INFO, &[], vec![pools, key]).map_err(error)?;
+    move_call(&mut txb, cetus_clmm_publish_at(), MODULE_FACTORY, FUNCTION_POOL_ID, &[], vec![info]).map_err(error)?;
     inspect_transaction_kind_bytes(txb)
 }
 
@@ -183,6 +185,7 @@ fn build_transaction(
     input: &BuildInput<'_>,
 ) -> Result<TxOutput, SwapperError> {
     let mut txb = TransactionBuilder::new();
+    let published_at = SuiAddress::from_str(published_at).map(Address::from)?;
     let input_coin = build_input_coin(&mut txb, route.input_coin_type(), input.amount, input.from_coins).map_err(error)?;
     let (swap_coin, swap_amount) = match route.fee_side {
         FeeSide::Input => {
@@ -266,7 +269,7 @@ fn build_transaction(
         .map_err(error)?;
     }
 
-    let sender = recipient(&quote.request.wallet_address)?;
+    let sender = SuiAddress::from_str(&quote.request.wallet_address).map(Address::from)?;
     let output_coin = from_balance(&mut txb, route.output_coin_type(), current_balance).map_err(error)?;
 
     let pay_output_fee = match route.fee_side {
@@ -279,7 +282,8 @@ fn build_transaction(
             .split_coins(output_coin, vec![fee_amount_arg])
             .pop()
             .ok_or_else(|| SwapperError::TransactionError("Cetus CLMM output fee split failed".into()))?;
-        transfer_coin(&mut txb, fee_coin, recipient(&referral_fee.address)?);
+        let recipient = SuiAddress::from_str(&referral_fee.address).map(Address::from)?;
+        transfer_coin(&mut txb, fee_coin, recipient);
     }
 
     let min_out = apply_slippage_in_bp(&route.net_amount_out(), quote.request.options.slippage.bps);
@@ -290,7 +294,12 @@ fn build_transaction(
         .ok_or_else(|| SwapperError::TransactionError("Cetus CLMM min-out split failed".into()))?;
     txb.merge_coins(output_coin, vec![split_off]);
 
-    let dest = swap_recipient(quote)?;
+    let dest = SuiAddress::from_str(if quote.request.destination_address.is_empty() {
+        &quote.request.wallet_address
+    } else {
+        &quote.request.destination_address
+    })
+    .map(Address::from)?;
     if dest == sender && is_sui_coin(route.output_coin_type()) {
         let gas = txb.gas();
         txb.merge_coins(gas, vec![output_coin]);
@@ -325,7 +334,8 @@ fn pay_referral_fee(txb: &mut TransactionBuilder, input_coin: Argument, fee: u64
         .split_coins(input_coin, vec![fee_amount])
         .pop()
         .ok_or_else(|| SwapperError::TransactionError("Sui referral fee split failed".into()))?;
-    transfer_coin(txb, fee_coin, recipient(&referral_fee.address)?);
+    let recipient = SuiAddress::from_str(&referral_fee.address).map(Address::from)?;
+    transfer_coin(txb, fee_coin, recipient);
     Ok(input_coin)
 }
 
@@ -341,25 +351,17 @@ fn transfer_coin(txb: &mut TransactionBuilder, coin: Argument, recipient: Addres
     txb.transfer_objects(vec![coin], recipient);
 }
 
-fn swap_recipient(quote: &Quote) -> Result<Address, SwapperError> {
-    if quote.request.destination_address.is_empty() {
-        recipient(&quote.request.wallet_address)
-    } else {
-        recipient(&quote.request.destination_address)
-    }
-}
-
-fn recipient(value: &str) -> Result<Address, SwapperError> {
-    Address::from_str(value).map_err(|err| SwapperError::TransactionError(format!("Invalid Sui address {value}: {err}")))
+fn cetus_clmm_publish_at() -> Address {
+    SuiAddress::from_str(CETUS_CLMM_PUBLISHED_AT).map(Address::from).unwrap()
 }
 
 fn shared_object_input(object_id: &str, initial_shared_version: u64, mutable: bool) -> Result<ObjectInput, SwapperError> {
-    let address = Address::from_str(object_id).map_err(|err| SwapperError::ComputeQuoteError(format!("Invalid Sui object id {object_id}: {err}")))?;
+    let address = SuiAddress::from_str(object_id).map(Address::from)?;
     Ok(ObjectInput::shared(address, initial_shared_version, mutable))
 }
 
 fn inspect_transaction_kind_bytes(mut txb: TransactionBuilder) -> Result<Vec<u8>, SwapperError> {
-    txb.set_sender(Address::from_str(EMPTY_ADDRESS).map_err(|err| SwapperError::ComputeQuoteError(format!("Invalid inspect sender: {err}")))?);
+    txb.set_sender(SuiAddress::from_str(EMPTY_ADDRESS).map(Address::from)?);
     txb.set_gas_price(0);
     txb.set_gas_budget(0);
     txb.add_gas_objects(vec![ObjectInput::owned(Address::ZERO, 0, Digest::ZERO)]);