[Sui Rosetta] - Delegation and balance queries (MystenLabs#7282)

MystenLabs#6423
MystenLabs#6422

This PR added 2 main features to Sui-rosetta,
1, Account balance query for `LockedSui`, `DelegatedSui` and
`PendingDelegation` by using rosetta's sub account.
2, added `Delegation` operation which allow user to delegate Sui and
Locked Sui to a validator using rosetta api.

example rosette delegation operation request:
```
curl --location --request POST '0.0.0.0:9003/construction/preprocess' \
--header 'Content-Type: application/json' \
--data-raw '{
    "network_identifier": {
        "blockchain": "sui",
        "network": "localnet"
    },
    "operations": [
        {
            "operation_identifier": {
                "index": 0
            },
            "type": "Delegation",
            "account": {
                "address": "0xb5e92aa6556fa0c27c4f6b5a659768b2d4141680"
            },
            "amount": {
                "value": "-1000000",
                "currency": {
                    "symbol": "SUI",
                    "decimals": 9
                }
            },
            "metadata": {
                "Delegation": {
                    "validator": "0x4a6f0a87a4cd4ba13184a924d8f73e8c081113cf"
                }
            }
        }
    ]
}'
```

example rosetta balance query request:
```
curl --location --request POST '0.0.0.0:9002/account/balance' \
--header 'Content-Type: application/json' \
--data-raw '{
    "network_identifier": {
        "blockchain": "sui",
        "network": "localnet"
    },
    "account_identifier": {
        "address": "0xe357b7fb38d261eff3818337f6ba40c156b203cd",
        "sub_account": {
            "address": "PendingDelegation"
        }
    }
}'
```
example rosetta balance query response:
```
{
    "block_identifier": {
        "index": 37,
        "hash": "7zZHnQuvMDjDeCxNt7Nb77WbWXmzE2uSaxbi5fs1pzzw"
    },
    "balances": [
        {
            "value": "1000000",
            "currency": {
                "symbol": "SUI",
                "decimals": 9
            }
        }
    ]
}
```

crates/sui-rosetta/Cargo.toml

@@ -55,3 +55,4 @@ sui-framework-build = { path = "../sui-framework-build" }
 test-utils = { path = "../test-utils" }
 tempfile = "3.3.0"
 rand = "0.8.5"
+reqwest = "0.11.13"

crates/sui-rosetta/src/account.rs

@@ -6,60 +6,115 @@ use axum::extract::State;
 use axum::{Extension, Json};
 use axum_extra::extract::WithRejection;
 use futures::StreamExt;
+use std::collections::HashMap;
 
-use sui_sdk::SUI_COIN_TYPE;
+use sui_sdk::{SuiClient, SUI_COIN_TYPE};
+use sui_types::base_types::SuiAddress;
+use sui_types::governance::DelegationStatus;
 
 use crate::errors::Error;
 use crate::types::{
     AccountBalanceRequest, AccountBalanceResponse, AccountCoinsRequest, AccountCoinsResponse,
-    Amount, Coin,
+    Amount, Coin, SubAccount, SubAccountType,
 };
 use crate::{OnlineServerContext, SuiEnv};
 
 /// Get an array of all AccountBalances for an AccountIdentifier and the BlockIdentifier
 /// at which the balance lookup was performed.
 /// [Rosetta API Spec](https://www.rosetta-api.org/docs/AccountApi.html#accountbalance)
 pub async fn balance(
-    State(context): State<OnlineServerContext>,
+    State(ctx): State<OnlineServerContext>,
     Extension(env): Extension<SuiEnv>,
     WithRejection(Json(request), _): WithRejection<Json<AccountBalanceRequest>, Error>,
 ) -> Result<AccountBalanceResponse, Error> {
     env.check_network_identifier(&request.network_identifier)?;
-
-    let block_id = if let Some(index) = request.block_identifier.index {
-        context.blocks().get_block_by_index(index).await.ok()
-    } else if let Some(hash) = request.block_identifier.hash {
-        context.blocks().get_block_by_hash(hash).await.ok()
+    let address = request.account_identifier.address;
+    if let Some(SubAccount { account_type }) = request.account_identifier.sub_account {
+        let balances = get_sub_account_balances(account_type, &ctx.client, address).await?;
+        Ok(AccountBalanceResponse {
+            block_identifier: ctx.blocks().current_block_identifier()?,
+            balances,
+        })
     } else {
-        None
-    }
-    .map(|b| b.block.block_identifier);
+        let block_identifier = if let Some(index) = request.block_identifier.index {
+            let response = ctx.blocks().get_block_by_index(index).await?;
+            response.block.block_identifier
+        } else if let Some(hash) = request.block_identifier.hash {
+            let response = ctx.blocks().get_block_by_hash(hash).await?;
+            response.block.block_identifier
+        } else {
+            ctx.blocks().current_block_identifier()?
+        };
 
-    if let Some(block_identifier) = block_id {
-        context
-            .blocks()
-            .get_balance_at_block(request.account_identifier.address, block_identifier.index)
+        ctx.blocks()
+            .get_balance_at_block(address, block_identifier.index)
             .await
             .map(|balance| AccountBalanceResponse {
                 block_identifier,
                 balances: vec![Amount::new(balance as i128)],
             })
-    } else {
-        let amount = context
-            .client
-            .coin_read_api()
-            .get_coins_stream(
-                request.account_identifier.address,
-                Some(SUI_COIN_TYPE.to_string()),
-            )
-            .fold(0u128, |acc, coin| async move { acc + coin.balance as u128 })
-            .await;
+    }
+}
 
-        Ok(AccountBalanceResponse {
-            block_identifier: context.blocks().current_block_identifier().await?,
-            balances: vec![Amount::new(amount as i128)],
+async fn get_sub_account_balances(
+    account_type: SubAccountType,
+    client: &SuiClient,
+    address: SuiAddress,
+) -> Result<Vec<Amount>, Error> {
+    let balances: HashMap<_, u128> = match account_type {
+        SubAccountType::DelegatedSui => {
+            let delegations = client
+                .governance_api()
+                .get_delegated_stakes(address)
+                .await?;
+            delegations
+                .into_iter()
+                .fold(HashMap::new(), |mut balances, delegation| {
+                    if let DelegationStatus::Active(_) = delegation.delegation_status {
+                        *balances
+                            .entry(delegation.staked_sui.sui_token_lock())
+                            .or_default() += delegation.staked_sui.principal() as u128;
+                    }
+                    balances
+                })
+        }
+        SubAccountType::PendingDelegation => {
+            let delegations = client
+                .governance_api()
+                .get_delegated_stakes(address)
+                .await?;
+            delegations
+                .into_iter()
+                .fold(HashMap::new(), |mut balances, delegation| {
+                    if let DelegationStatus::Pending = delegation.delegation_status {
+                        *balances
+                            .entry(delegation.staked_sui.sui_token_lock())
+                            .or_default() += delegation.staked_sui.principal() as u128;
+                    }
+                    balances
+                })
+        }
+        SubAccountType::LockedSui => {
+            let sui_balance = client.coin_read_api().get_balance(address, None).await?;
+            sui_balance
+                .locked_balance
+                .into_iter()
+                .map(|(lock, amount)| (Some(lock), amount))
+                .collect()
+        }
+    };
+
+    Ok(balances
+        .into_iter()
+        .map(|(lock, balance)| {
+            if let Some(lock) = lock {
+                // Safe to cast to i128 as total issued SUI will be less then u64.
+                Amount::new_locked(lock, balance as i128)
+            } else {
+                Amount::new(balance as i128)
+            }
         })
-    }
+        .collect())
 }
 
 /// Get an array of all unspent coins for an AccountIdentifier and the BlockIdentifier at which the lookup was performed. .
@@ -82,7 +137,7 @@ pub async fn coins(
         .await;
 
     Ok(AccountCoinsResponse {
-        block_identifier: context.blocks().current_block_identifier().await?,
+        block_identifier: context.blocks().current_block_identifier()?,
         coins,
     })
 }

crates/sui-rosetta/src/construction.rs

@@ -6,7 +6,10 @@ use fastcrypto::encoding::{Encoding, Hex};
 use sui_types::base_types::SuiAddress;
 use sui_types::crypto;
 use sui_types::crypto::{SignatureScheme, ToFromBytes};
-use sui_types::messages::{ExecuteTransactionRequestType, Transaction, TransactionData};
+use sui_types::messages::{
+    ExecuteTransactionRequestType, SingleTransactionKind, Transaction, TransactionData,
+    TransactionKind,
+};
 
 use crate::errors::Error;
 use crate::types::{
@@ -19,7 +22,6 @@ use crate::types::{
     TransactionIdentifierResponse, TransactionMetadata,
 };
 use crate::{OnlineServerContext, SuiEnv};
-use anyhow::anyhow;
 use axum::extract::State;
 use axum_extra::extract::WithRejection;
 use sui_sdk::rpc_types::SuiExecutionStatus;
@@ -54,7 +56,10 @@ pub async fn payloads(
     let metadata = request.metadata.ok_or(Error::MissingMetadata)?;
     let address = metadata.sender;
 
-    let data = request.operations.into_internal()?.into_data(metadata);
+    let data = request
+        .operations
+        .into_internal(Some(metadata.tx_metadata.clone().into()))?
+        .try_into_data(metadata)?;
     let intent_msg = IntentMessage::new(Intent::default(), data);
     let intent_msg_bytes = bcs::to_bytes(&intent_msg)?;
 
@@ -77,10 +82,7 @@ pub async fn combine(
     WithRejection(Json(request), _): WithRejection<Json<ConstructionCombineRequest>, Error>,
 ) -> Result<ConstructionCombineResponse, Error> {
     env.check_network_identifier(&request.network_identifier)?;
-    let unsigned_tx = request
-        .unsigned_transaction
-        .to_vec()
-        .map_err(|e| anyhow!(e))?;
+    let unsigned_tx = request.unsigned_transaction.to_vec()?;
     let intent_msg: IntentMessage<TransactionData> = bcs::from_bytes(&unsigned_tx)?;
     let sig = request
         .signatures
@@ -152,7 +154,7 @@ pub async fn preprocess(
 ) -> Result<ConstructionPreprocessResponse, Error> {
     env.check_network_identifier(&request.network_identifier)?;
 
-    let internal_operation = request.operations.into_internal()?;
+    let internal_operation = request.operations.into_internal(request.metadata)?;
     let sender = internal_operation.sender();
 
     Ok(ConstructionPreprocessResponse {
@@ -169,10 +171,7 @@ pub async fn hash(
     WithRejection(Json(request), _): WithRejection<Json<ConstructionHashRequest>, Error>,
 ) -> Result<TransactionIdentifierResponse, Error> {
     env.check_network_identifier(&request.network_identifier)?;
-    let tx_bytes = request
-        .signed_transaction
-        .to_vec()
-        .map_err(|e| anyhow::anyhow!(e))?;
+    let tx_bytes = request.signed_transaction.to_vec()?;
     let tx: Transaction = bcs::from_bytes(&tx_bytes)?;
 
     Ok(TransactionIdentifierResponse {
@@ -211,14 +210,59 @@ pub async fn metadata(
             let gas = sender_coins[0];
             (TransactionMetadata::PaySui(sender_coins), gas)
         }
+        InternalOperation::Delegation {
+            sender,
+            validator,
+            amount,
+            locked_until_epoch,
+        } => {
+            let coins = context
+                .client
+                .coin_read_api()
+                .select_coins(*sender, None, *amount as u128, *locked_until_epoch, vec![])
+                .await?
+                .into_iter()
+                .map(|coin| coin.object_ref())
+                .collect::<Vec<_>>();
+
+            let data = context
+                .client
+                .transaction_builder()
+                .request_add_delegation(
+                    *sender,
+                    coins.iter().map(|coin| coin.0).collect(),
+                    Some(*amount as u64),
+                    *validator,
+                    None,
+                    2000,
+                )
+                .await?;
+
+            let gas = data.gas();
+            let TransactionKind::Single(SingleTransactionKind::Call(call)) = data.kind else{
+                // This will not happen because `request_add_delegation` call creates a move call transaction.
+                panic!("Malformed transaction received from TransactionBuilder.")
+            };
+
+            (
+                TransactionMetadata::Delegation {
+                    sui_framework: call.package,
+                    coins,
+                    locked_until_epoch: *locked_until_epoch,
+                },
+                gas,
+            )
+        }
     };
     // get gas estimation from dry-run, this will also return any tx error.
-    let data = option.internal_operation.into_data(ConstructionMetadata {
-        tx_metadata: tx_metadata.clone(),
-        sender,
-        gas,
-        budget: 1000,
-    });
+    let data = option
+        .internal_operation
+        .try_into_data(ConstructionMetadata {
+            tx_metadata: tx_metadata.clone(),
+            sender,
+            gas,
+            budget: 1000,
+        })?;
     let dry_run = context.client.read_api().dry_run_transaction(data).await?;
 
     let budget = dry_run.gas_used.computation_cost + dry_run.gas_used.storage_cost;
@@ -245,16 +289,11 @@ pub async fn parse(
     env.check_network_identifier(&request.network_identifier)?;
 
     let data = if request.signed {
-        let tx: Transaction = bcs::from_bytes(
-            &request
-                .transaction
-                .to_vec()
-                .map_err(|e| anyhow::anyhow!(e))?,
-        )?;
+        let tx: Transaction = bcs::from_bytes(&request.transaction.to_vec()?)?;
         tx.into_data().intent_message.value
     } else {
         let intent: IntentMessage<TransactionData> =
-            bcs::from_bytes(&request.transaction.to_vec().map_err(|e| anyhow!(e))?)?;
+            bcs::from_bytes(&request.transaction.to_vec()?)?;
         intent.value
     };
     let account_identifier_signers = if request.signed {