Rosetta: Handle single GasCoin transfers

crates/sui-rosetta/src/operations.rs

@@ -558,6 +558,83 @@ impl Operations {
         };
         balance_change.chain(gas)
     }
+
+    fn is_single_gascoin_transfer(tx: SuiTransactionBlockKind) -> bool {
+        match tx {
+            SuiTransactionBlockKind::ProgrammableTransaction(pt) => {
+                let SuiProgrammableTransactionBlock {
+                    inputs: _,
+                    commands,
+                } = &pt;
+                return commands
+                    .into_iter()
+                    .find(|command| match command {
+                        SuiCommand::TransferObjects(objs, _) => {
+                            objs.len() > 0 && objs[0] == SuiArgument::GasCoin
+                        }
+                        _ => false,
+                    })
+                    .is_some();
+            }
+            _ => {}
+        }
+        false
+    }
+
+    fn process_single_gascoin_transfer(
+        coin_change_operations: &mut impl Iterator<Item = crate::operations::Operation>,
+        tx: SuiTransactionBlockKind,
+        prev_gas_owner: SuiAddress,
+        new_gas_owner: SuiAddress,
+        gas_used: i128,
+    ) -> Vec<Operation> {
+        let mut operations = vec![];
+        if Self::is_single_gascoin_transfer(tx) {
+            coin_change_operations.into_iter().for_each(|operation| {
+                match operation.type_ {
+                    OperationType::Gas => {
+                        // change gas account back to the previous owner as it is the one
+                        // who paid for the txn (this is the format Rosetta wants to process)
+                        operations.push(Operation::gas(prev_gas_owner, gas_used))
+                    }
+                    OperationType::SuiBalanceChange => {
+                        operation.account.map(|account| {
+                            let mut amount = match operation.amount {
+                                Some(amount) => amount,
+                                None => return,
+                            };
+                            let mut is_convert_to_pay_sui = false;
+                            if account.address == prev_gas_owner {
+                                // previous owner's balance needs to be adjusted for gas
+                                amount.value -= gas_used;
+                                is_convert_to_pay_sui = true;
+                            } else if account.address == new_gas_owner {
+                                // new owner's balance needs to be adjusted for gas
+                                amount.value += gas_used;
+                                is_convert_to_pay_sui = true;
+                            }
+                            if is_convert_to_pay_sui {
+                                operations.push(Operation::pay_sui(
+                                    operation.status,
+                                    account.address,
+                                    amount.value,
+                                ));
+                            } else {
+                                operations.push(Operation::balance_change(
+                                    operation.status,
+                                    account.address,
+                                    amount.value,
+                                    amount.currency,
+                                ));
+                            }
+                        });
+                    }
+                    _ => {}
+                }
+            });
+        }
+        operations
+    }
 }
 
 impl Operations {
@@ -592,7 +669,7 @@ impl Operations {
             - gas_summary.computation_cost as i128;
 
         let status = Some(effect.into_status().into());
-        let ops = Operations::try_from_data(tx.data, status)?;
+        let ops = Operations::try_from_data(tx.data.clone(), status)?;
         let ops = ops.into_iter();
 
         // We will need to subtract the operation amounts from the actual balance
@@ -662,17 +739,26 @@ impl Operations {
         }
 
         // Extract coin change operations from balance changes
-        let coin_change_operations = Self::process_balance_change(
+        let mut coin_change_operations = Self::process_balance_change(
             gas_owner,
             gas_used,
             balance_changes,
             status,
-            accounted_balances,
+            accounted_balances.clone(),
+        );
+
+        let single_gascoin_transfer = Self::process_single_gascoin_transfer(
+            &mut coin_change_operations,
+            tx.data.transaction().clone(),
+            tx.data.gas_data().owner,
+            gas_owner,
+            gas_used,
         );
 
         let ops: Operations = ops
             .into_iter()
             .chain(coin_change_operations)
+            .chain(single_gascoin_transfer)
             .chain(staking_balance)
             .collect();
 

crates/sui-rosetta/tests/end_to_end_tests.rs

@@ -1,23 +1,34 @@
 // Copyright (c) Mysten Labs, Inc.
 // SPDX-License-Identifier: Apache-2.0
 
+use anyhow::anyhow;
+use rand::rngs::OsRng;
+use rand::seq::IteratorRandom;
+use rosetta_client::start_rosetta_test_server;
 use serde_json::json;
+use shared_crypto::intent::Intent;
 use std::num::NonZeroUsize;
 use std::time::Duration;
-
-use rosetta_client::start_rosetta_test_server;
-use sui_json_rpc_types::SuiTransactionBlockResponseOptions;
+use sui_json_rpc_types::{
+    SuiObjectDataOptions, SuiObjectResponseQuery, SuiTransactionBlockResponseOptions,
+};
 use sui_keys::keystore::AccountKeystore;
 use sui_rosetta::operations::Operations;
-use sui_rosetta::types::Currencies;
 use sui_rosetta::types::{
     AccountBalanceRequest, AccountBalanceResponse, AccountIdentifier, Currency, NetworkIdentifier,
     SubAccount, SubAccountType, SuiEnv,
 };
+use sui_rosetta::types::{Currencies, OperationType};
 use sui_rosetta::CoinMetadataCache;
 use sui_sdk::rpc_types::{SuiExecutionStatus, SuiTransactionBlockEffectsAPI};
+use sui_sdk::SuiClient;
 use sui_swarm_config::genesis_config::{DEFAULT_GAS_AMOUNT, DEFAULT_NUMBER_OF_OBJECT_PER_ACCOUNT};
+use sui_types::base_types::{ObjectID, ObjectRef, SuiAddress};
+use sui_types::programmable_transaction_builder::ProgrammableTransactionBuilder;
 use sui_types::quorum_driver_types::ExecuteTransactionRequestType;
+use sui_types::transaction::{
+    Argument, InputObjectKind, Transaction, TransactionData, TEST_ONLY_GAS_UNIT_FOR_TRANSFER,
+};
 use sui_types::utils::to_sender_signed_transaction;
 use test_cluster::TestClusterBuilder;
 
@@ -501,3 +512,117 @@ async fn test_pay_sui_multiple_times() {
         );
     }
 }
+
+async fn get_random_sui(
+    client: &SuiClient,
+    sender: SuiAddress,
+    except: Vec<ObjectID>,
+) -> ObjectRef {
+    let coins = client
+        .read_api()
+        .get_owned_objects(
+            sender,
+            Some(SuiObjectResponseQuery::new_with_options(
+                SuiObjectDataOptions::new()
+                    .with_type()
+                    .with_owner()
+                    .with_previous_transaction(),
+            )),
+            /* cursor */ None,
+            /* limit */ None,
+        )
+        .await
+        .unwrap()
+        .data;
+
+    let coin_resp = coins
+        .iter()
+        .filter(|object| {
+            let obj = object.object().unwrap();
+            obj.is_gas_coin() && !except.contains(&obj.object_id)
+        })
+        .choose(&mut OsRng)
+        .unwrap();
+
+    let coin = coin_resp.object().unwrap();
+    (coin.object_id, coin.version, coin.digest)
+}
+#[tokio::test]
+async fn test_transfer_single_gas_coin() {
+    let test_cluster = TestClusterBuilder::new().build().await;
+    let sender = test_cluster.get_address_0();
+    let recipient = test_cluster.get_address_1();
+    let client = test_cluster.wallet.get_client().await.unwrap();
+    let keystore = &test_cluster.wallet.config.keystore;
+
+    let pt = {
+        let mut builder = ProgrammableTransactionBuilder::new();
+        builder.transfer_arg(recipient, Argument::GasCoin);
+        builder.finish()
+    };
+
+    let input_objects = pt
+        .input_objects()
+        .unwrap_or_default()
+        .iter()
+        .flat_map(|obj| {
+            if let InputObjectKind::ImmOrOwnedMoveObject((id, ..)) = obj {
+                Some(*id)
+            } else {
+                None
+            }
+        })
+        .collect::<Vec<_>>();
+    let gas = vec![get_random_sui(&client, sender, input_objects).await];
+    let gas_price = client
+        .governance_api()
+        .get_reference_gas_price()
+        .await
+        .unwrap();
+
+    let data = TransactionData::new_programmable(
+        sender,
+        gas,
+        pt,
+        TEST_ONLY_GAS_UNIT_FOR_TRANSFER * gas_price,
+        gas_price,
+    );
+
+    let signature = keystore
+        .sign_secure(&sender, &data, Intent::sui_transaction())
+        .unwrap();
+
+    let response = client
+        .quorum_driver_api()
+        .execute_transaction_block(
+            Transaction::from_data(data.clone(), vec![signature]),
+            SuiTransactionBlockResponseOptions::new()
+                .with_effects()
+                .with_object_changes()
+                .with_balance_changes()
+                .with_input(),
+            Some(ExecuteTransactionRequestType::WaitForLocalExecution),
+        )
+        .await
+        .map_err(|e| anyhow!("TX execution failed for {data:#?}, error : {e}"))
+        .unwrap();
+
+    let coin_cache = CoinMetadataCache::new(client, NonZeroUsize::new(2).unwrap());
+    let operations = Operations::try_from_response(response, &coin_cache)
+        .await
+        .unwrap();
+    println!("operations: {operations:#?}");
+
+    let mut balance = 0;
+    operations
+        .into_iter()
+        .for_each(|op| {
+            if op.type_ == OperationType::Gas {
+                assert_eq!(op.account.unwrap().address, sender);
+            }
+            if op.type_ == OperationType::PaySui {
+                balance += op.amount.unwrap().value;
+            }
+        });
+    assert_eq!(balance, 0);
+}