From 7808984fd1994dab681269781355ca8d9c8f51e3 Mon Sep 17 00:00:00 2001
From: xiaohuo <xiaohuo200@gmail.com>
Date: Thu, 1 Apr 2021 02:07:09 +0800
Subject: [PATCH] feat(rosetta): enhance deleagation/undelegation and balance
 api

---
 rosetta/services/account.go      | 45 +++++++++++++++--
 rosetta/services/tx_operation.go | 83 +++++++++++++++++++++++++++++++-
 2 files changed, 121 insertions(+), 7 deletions(-)

diff --git a/rosetta/services/account.go b/rosetta/services/account.go
index 231e5f2c6..bc3510394 100644
--- a/rosetta/services/account.go
+++ b/rosetta/services/account.go
@@ -3,6 +3,7 @@ package services
 import (
 	"context"
 	"fmt"
+	"math/big"
 
 	"github.com/coinbase/rosetta-sdk-go/server"
 	"github.com/coinbase/rosetta-sdk-go/types"
@@ -53,11 +54,45 @@ func (s *AccountAPI) AccountBalance(
 		})
 	}
 	blockNum := rpc.BlockNumber(block.Header().Header.Number().Int64())
-	balance, err := s.hmy.GetBalance(ctx, addr, blockNum)
-	if err != nil {
-		return nil, common.NewError(common.SanityCheckError, map[string]interface{}{
-			"message": "invalid address",
-		})
+	balance := new(big.Int)
+
+	// delegated balance
+	if request.AccountIdentifier.SubAccount != nil {
+		subAccount := request.AccountIdentifier.SubAccount
+		ty, exist := subAccount.Metadata["type"]
+
+		if exist && ty.(string) == "delegation" {
+			validatorAddr := subAccount.Address
+			validators, delegations := s.hmy.GetDelegationsByDelegatorByBlock(addr, block)
+			for index, validator := range validators {
+				if validatorAddr == internalCommon.MustAddressToBech32(validator) {
+					balance = new(big.Int).Add(balance, delegations[index].Amount)
+				}
+			}
+			// pending undelegated balance
+		} else if exist && ty.(string) == "undelegation" {
+			validatorAddr := subAccount.Address
+			validators, delegations := s.hmy.GetDelegationsByDelegatorByBlock(addr, block)
+			for index, validator := range validators {
+				if validatorAddr == internalCommon.MustAddressToBech32(validator) {
+					undelegations := delegations[index].Undelegations
+					for _, undelegate := range undelegations {
+						balance = new(big.Int).Add(balance, undelegate.Amount)
+					}
+				}
+			}
+		} else {
+			return nil, common.NewError(common.SanityCheckError, map[string]interface{}{
+				"message": "invalid sub account or type",
+			})
+		}
+	} else {
+		balance, err = s.hmy.GetBalance(ctx, addr, blockNum)
+		if err != nil {
+			return nil, common.NewError(common.SanityCheckError, map[string]interface{}{
+				"message": "invalid address",
+			})
+		}
 	}
 
 	amount := types.Amount{
diff --git a/rosetta/services/tx_operation.go b/rosetta/services/tx_operation.go
index d43e41efb..a33c69c81 100644
--- a/rosetta/services/tx_operation.go
+++ b/rosetta/services/tx_operation.go
@@ -118,7 +118,7 @@ func GetNativeOperationsFromStakingTransaction(
 		}
 	}
 
-	return append(gasOperations, &types.Operation{
+	operations := append(gasOperations, &types.Operation{
 		OperationIdentifier: &types.OperationIdentifier{
 			Index: gasOperations[0].OperationIdentifier.Index + 1,
 		},
@@ -127,7 +127,86 @@ func GetNativeOperationsFromStakingTransaction(
 		Account:  accountID,
 		Amount:   amount,
 		Metadata: metadata,
-	}), nil
+	})
+
+	// expose delegated balance
+	if tx.StakingType() == stakingTypes.DirectiveDelegate {
+		op2 := GetDelegateOperationForSubAccount(tx, operations[1])
+		return append(operations, op2), nil
+	}
+
+	if tx.StakingType() == stakingTypes.DirectiveUndelegate {
+		// set sub account
+		validatorAddress := operations[1].Metadata["validatorAddress"]
+		operations[1].Account.SubAccount = &types.SubAccountIdentifier{
+			Address: validatorAddress.(string),
+			Metadata: map[string]interface{}{
+				"type": "delegation",
+			},
+		}
+
+		op2 := &types.Operation{
+			OperationIdentifier: &types.OperationIdentifier{
+				Index: operations[1].OperationIdentifier.Index + 1,
+			},
+			Type:   tx.StakingType().String(),
+			Status: GetTransactionStatus(tx, receipt),
+			Account: &types.AccountIdentifier{
+				Address: operations[1].Account.Address,
+				SubAccount: &types.SubAccountIdentifier{
+					Address: validatorAddress.(string),
+					Metadata: map[string]interface{}{
+						"type": "undelegation",
+					},
+				},
+				Metadata: operations[1].Account.Metadata,
+			},
+			Amount:   operations[1].Amount,
+			Metadata: operations[1].Metadata,
+		}
+
+		return append(operations, op2), nil
+	}
+
+	return operations, nil
+
+}
+
+func GetDelegateOperationForSubAccount(tx *stakingTypes.StakingTransaction, delegateOperation *types.Operation) *types.Operation {
+	amt, _ := new(big.Int).SetString(delegateOperation.Amount.Value, 10)
+	delegateAmt := new(big.Int).Sub(new(big.Int).SetUint64(0), amt)
+	validatorAddress := delegateOperation.Metadata["validatorAddress"]
+	delegation := &types.Operation{
+		OperationIdentifier: &types.OperationIdentifier{
+			Index: delegateOperation.OperationIdentifier.Index + 1,
+		},
+		RelatedOperations: []*types.OperationIdentifier{
+			{
+				Index: delegateOperation.OperationIdentifier.Index,
+			},
+		},
+		Type:   tx.StakingType().String(),
+		Status: delegateOperation.Status,
+		Account: &types.AccountIdentifier{
+			Address: delegateOperation.Account.Address,
+			SubAccount: &types.SubAccountIdentifier{
+				Address: validatorAddress.(string),
+				Metadata: map[string]interface{}{
+					"type": "delegation",
+				},
+			},
+			Metadata: delegateOperation.Account.Metadata,
+		},
+		Amount: &types.Amount{
+			Value:    delegateAmt.String(),
+			Currency: delegateOperation.Amount.Currency,
+			Metadata: delegateOperation.Amount.Metadata,
+		},
+		Metadata: delegateOperation.Metadata,
+	}
+
+	return delegation
+
 }
 
 // GetSideEffectOperationsFromUndelegationPayouts from the given payouts.