From b57e324236015c8d1a0893dc5b053725c4cbd54a Mon Sep 17 00:00:00 2001 From: Himitsuko Date: Thu, 22 Dec 2022 11:03:28 +0700 Subject: [PATCH] add: Support batch call by multicall only and mixing with native coin --- core/balance_scanner.go | 138 +++++++++++++++++++++++++++++++++-- core/balance_scanner_test.go | 82 +++++++++++++++++++++ 2 files changed, 214 insertions(+), 6 deletions(-) diff --git a/core/balance_scanner.go b/core/balance_scanner.go index a10feac..af09f6d 100644 --- a/core/balance_scanner.go +++ b/core/balance_scanner.go @@ -19,7 +19,12 @@ import ( ) var BalanceScannerMetaData = &bind.MetaData{ - ABI: "[{\"inputs\":[{\"internalType\":\"address\",\"name\":\"account\",\"type\":\"address\"},{\"internalType\":\"contractIBEP20[]\",\"name\":\"tokens\",\"type\":\"address[]\"}],\"name\":\"getBalances\",\"outputs\":[{\"internalType\":\"uint256[]\",\"name\":\"balances\",\"type\":\"uint256[]\"}],\"stateMutability\":\"view\",\"type\":\"function\"}]", + ABI: "[{\"inputs\":[{\"internalType\":\"address\",\"name\":\"account\",\"type\":\"address\"},{\"internalType\":\"contractIBEP20[]\",\"name\":\"tokens\",\"type\":\"address[]\"}],\"name\":\"getBalances\",\"outputs\":[{\"internalType\":\"uint256[]\",\"name\":\"balances\",\"type\":\"uint256[]\"}],\"stateMutability\":\"view\",\"type\":\"function\"},{\"inputs\":[{\"components\":[{\"internalType\":\"address\",\"name\":\"target\",\"type\":\"address\"},{\"internalType\":\"bytes\",\"name\":\"callData\",\"type\":\"bytes\"}],\"internalType\":\"structDePocketCore.Call[]\",\"name\":\"calls\",\"type\":\"tuple[]\"}],\"name\":\"aggregate\",\"outputs\":[{\"internalType\":\"uint256\",\"name\":\"blockNumber\",\"type\":\"uint256\"},{\"internalType\":\"bytes[]\",\"name\":\"returnData\",\"type\":\"bytes[]\"}],\"stateMutability\":\"view\",\"type\":\"function\"},{\"inputs\":[{\"components\":[{\"internalType\":\"address\",\"name\":\"target\",\"type\":\"address\"},{\"internalType\":\"bytes\",\"name\":\"callData\",\"type\":\"bytes\"}],\"internalType\":\"structDePocketCore.Call[]\",\"name\":\"calls\",\"type\":\"tuple[]\"}],\"name\":\"blockAndAggregate\",\"outputs\":[{\"internalType\":\"uint256\",\"name\":\"blockNumber\",\"type\":\"uint256\"},{\"internalType\":\"bytes32\",\"name\":\"blockHash\",\"type\":\"bytes32\"},{\"components\":[{\"internalType\":\"bool\",\"name\":\"success\",\"type\":\"bool\"},{\"internalType\":\"bytes\",\"name\":\"returnData\",\"type\":\"bytes\"}],\"internalType\":\"structDePocketCore.Result[]\",\"name\":\"returnData\",\"type\":\"tuple[]\"}],\"stateMutability\":\"view\",\"type\":\"function\"},{\"inputs\":[{\"internalType\":\"bool\",\"name\":\"requireSuccess\",\"type\":\"bool\"},{\"components\":[{\"internalType\":\"address\",\"name\":\"target\",\"type\":\"address\"},{\"internalType\":\"bytes\",\"name\":\"callData\",\"type\":\"bytes\"}],\"internalType\":\"structDePocketCore.Call[]\",\"name\":\"calls\",\"type\":\"tuple[]\"}],\"name\":\"tryAggregate\",\"outputs\":[{\"components\":[{\"internalType\":\"bool\",\"name\":\"success\",\"type\":\"bool\"},{\"internalType\":\"bytes\",\"name\":\"returnData\",\"type\":\"bytes\"}],\"internalType\":\"structDePocketCore.Result[]\",\"name\":\"returnData\",\"type\":\"tuple[]\"}],\"stateMutability\":\"view\",\"type\":\"function\"},{\"inputs\":[{\"internalType\":\"bool\",\"name\":\"requireSuccess\",\"type\":\"bool\"},{\"components\":[{\"internalType\":\"address\",\"name\":\"target\",\"type\":\"address\"},{\"internalType\":\"bytes\",\"name\":\"callData\",\"type\":\"bytes\"}],\"internalType\":\"structDePocketCore.Call[]\",\"name\":\"calls\",\"type\":\"tuple[]\"}],\"name\":\"tryBlockAndAggregate\",\"outputs\":[{\"internalType\":\"uint256\",\"name\":\"blockNumber\",\"type\":\"uint256\"},{\"internalType\":\"bytes32\",\"name\":\"blockHash\",\"type\":\"bytes32\"},{\"components\":[{\"internalType\":\"bool\",\"name\":\"success\",\"type\":\"bool\"},{\"internalType\":\"bytes\",\"name\":\"returnData\",\"type\":\"bytes\"}],\"internalType\":\"structDePocketCore.Result[]\",\"name\":\"returnData\",\"type\":\"tuple[]\"}],\"stateMutability\":\"view\",\"type\":\"function\"},{\"constant\":true,\"inputs\":[{\"internalType\":\"address\",\"name\":\"account\",\"type\":\"address\"}],\"name\":\"balanceOf\",\"outputs\":[{\"internalType\":\"uint256\",\"name\":\"\",\"type\":\"uint256\"}],\"payable\":false,\"stateMutability\":\"view\",\"type\":\"function\"}]", +} + +type MultiCall struct { + Target common.Address + CallData []byte } type CallArg struct { @@ -51,6 +56,10 @@ func NewBalanceScanner(address string, client *ethclient.Client, rpcClient *rpc. }, nil } + +/* + GetBalances query the ERC20 token balances in batch + */ func (_BalanceScanner *BalanceScanner) GetBalances(account common.Address, tokens []common.Address) ([]*big.Int, error) { var out []interface{} callData, err := _BalanceScanner._abi.Pack("getBalances", account, tokens) @@ -69,6 +78,9 @@ func (_BalanceScanner *BalanceScanner) GetBalances(account common.Address, token return out0, err } +/* + TokenMapBalances query the ERC20 token balances in batch and return the result in mapping by token_address:balance +*/ func (_BalanceScanner *BalanceScanner) TokenMapBalances(address common.Address, tokens []common.Address) (map[string]*big.Int, error) { balances := make(map[string]*big.Int) res, err := _BalanceScanner.GetBalances(address, tokens) @@ -81,12 +93,102 @@ func (_BalanceScanner *BalanceScanner) TokenMapBalances(address common.Address, return balances, nil } -func (_BalanceScanner *BalanceScanner) GetBalancesCallData(account common.Address, tokens []common.Address) ([]byte, error) { - return _BalanceScanner._abi.Pack("getBalances", account, tokens) -} +/* + BatchCallBalancesWithNative query the ERC20 token balances in batch of RPC and multi-call and return the result in mapping by token_address:balance +*/ +func (_BalanceScanner *BalanceScanner) BatchCallBalancesWithNative(account common.Address, tokenChunks [][]common.Address, nativeAddress string) ([]*big.Int, error) { + var calls []rpc.BatchElem + var callResults = make([]hexutil.Bytes, len(tokenChunks)) + var nativeResult = new(hexutil.Big) + var nativeIndex = -1 // Determinate the index of native token + var skipResult = -1 // Determinate the balanceOf result that need to be skipped when decode the result + var count = 0 -func (_BalanceScanner *BalanceScanner) DecodeCallResponse(data hexutil.Bytes) ([]interface{}, error) { - return _BalanceScanner._abi.Unpack("getBalances", data) + for index, chunk := range tokenChunks { + var multiBalanceCalls []MultiCall + if len(chunk) == 1 && strings.ToLower(chunk[0].String()) == nativeAddress { + nativeIndex = count + skipResult = index + count++ + continue + } + for _, tk := range chunk { + data, err := _BalanceScanner.BalanceOfCallData(account) + if err != nil { + return nil, err + } + if strings.ToLower(tk.String()) != nativeAddress { + multiBalanceCalls = append(multiBalanceCalls, MultiCall{ + Target: tk, + CallData: data, + }) + } else { + nativeIndex = count + } + count++ + } + if len(multiBalanceCalls) > 0 { + callData, err := _BalanceScanner.GetMulticallBalanceCallData(multiBalanceCalls) + if err != nil { + return []*big.Int{}, err + } + elem := rpc.BatchElem { + Method: "eth_call", + Args: toCallArgs(ethereum.CallMsg{To: &_BalanceScanner.ContractAddress, Data: callData}), + Result: &callResults[index], + } + calls = append(calls, elem) + } + } + + if nativeIndex != -1 { + nativeElem := rpc.BatchElem { + Method: "eth_getBalance", + Args: []interface{}{account, toBlockNumArg(nil)}, + Result: &nativeResult, + } + calls = append(calls, nativeElem) + } + + err := _BalanceScanner.RpcClient.BatchCall(calls) + + if err != nil { + return []*big.Int{}, err + } + + var res = make([]*big.Int, count) + resCount := 0 + for idx, result := range callResults { + if idx == skipResult { + continue + } + if resCount == nativeIndex { + resCount ++ + } + r, err := _BalanceScanner.DecodeMulticallBalanceResponse(result) + if err != nil { + return []*big.Int{}, err + } + for _, response := range r[0].([]struct { + Success bool `json:"success"` + ReturnData []byte `json:"returnData"` + }) { + if response.Success { + balance, err := _BalanceScanner.DeCodeBalanceOf(response.ReturnData) + if err != nil { + return []*big.Int{}, err + } + res[resCount] = balance[0].(*big.Int) + } else { + res[resCount] = new(big.Int) + } + resCount++ + } + } + if nativeIndex != -1 { + res[nativeIndex] = nativeResult.ToInt() + } + return res, nil } func (_BalanceScanner *BalanceScanner) BatchCallBalances(account common.Address, tokenChunks [][]common.Address) ([]*big.Int, error) { @@ -172,4 +274,28 @@ func toBlockNumArg(number *big.Int) string { return "pending" } return hexutil.EncodeBig(number) +} + +func (_BalanceScanner *BalanceScanner) GetBalancesCallData(account common.Address, tokens []common.Address) ([]byte, error) { + return _BalanceScanner._abi.Pack("getBalances", account, tokens) +} + +func (_BalanceScanner *BalanceScanner) BalanceOfCallData(account common.Address) ([]byte, error) { + return _BalanceScanner._abi.Pack("balanceOf", account) +} + +func (_BalanceScanner *BalanceScanner) DeCodeBalanceOf(data hexutil.Bytes) ([]interface{}, error) { + return _BalanceScanner._abi.Unpack("balanceOf", data) +} + +func (_BalanceScanner *BalanceScanner) DecodeCallResponse(data hexutil.Bytes) ([]interface{}, error) { + return _BalanceScanner._abi.Unpack("getBalances", data) +} + +func (_BalanceScanner *BalanceScanner) DecodeMulticallBalanceResponse(data hexutil.Bytes) ([]interface{}, error) { + return _BalanceScanner._abi.Unpack("tryAggregate", data) +} + +func (_BalanceScanner *BalanceScanner) GetMulticallBalanceCallData(calls []MultiCall) ([]byte, error) { + return _BalanceScanner._abi.Pack("tryAggregate", false, calls) } \ No newline at end of file diff --git a/core/balance_scanner_test.go b/core/balance_scanner_test.go index 0ba8895..65519bb 100644 --- a/core/balance_scanner_test.go +++ b/core/balance_scanner_test.go @@ -59,3 +59,85 @@ func TestBalanceScanner_BatchCall(t *testing.T) { } fmt.Println("Total result", resultCount) } + +func TestBalanceScanner_MultiBatchCall(t *testing.T) { + client, _ := ethclient.Dial("https://bsc-dataseed.binance.org/") + r, _ := rpc.Dial("https://bsc-dataseed.binance.org/") + bep20, _ := NewBalanceScanner("0xcA11bde05977b3631167028862bE2a173976CA11", client, r) + resultCount := 0 + ret, err := bep20.BatchCallBalancesWithNative(common.HexToAddress("0xca0C80122afA57c38BcAa14fC77E056b94288469"), [][]common.Address{ + { + common.HexToAddress("0x7d99eda556388Ad7743A1B658b9C4FC67D7A9d74"), + common.HexToAddress("0x7d99eda556388Ad7743A1B658b9C4FC67D7A9d74"), + }, + { + common.HexToAddress("0x7d99eda556388Ad7743A1B658b9C4FC67D7A9d74"), + common.HexToAddress("0x7d99eda556388Ad7743A1B658b9C4FC67D7A9d74"), + common.HexToAddress("0x7d99eda556388Ad7743A1B658b9C4FC67D7A9d74"), + }, + { + common.HexToAddress("0x7d99eda556388Ad7743A1B658b9C4FC67D7A9d74"), + common.HexToAddress("0x7d99eda556388Ad7743A1B658b9C4FC67D7A9d74"), + common.HexToAddress("0x7d99eda556388Ad7743A1B658b9C4FC67D7A9d74"), + common.HexToAddress("0x7d99eda556388Ad7743A1B658b9C4FC67D7A9d74"), + common.HexToAddress("0x7d99eda556388Ad7743A1B658b9C4FC67D7A9d74"), + common.HexToAddress("0x0000000000000000000000000000000000000000"), + }, + }, "0x0000000000000000000000000000000000000000") + if err != nil { + fmt.Println(err) + } + for j := 0; j < len(ret); j++ { + if j == len(ret)-1 { + assert.Equal(t, ret[j].String(), "207924141787718960") + } else { + assert.Equal(t, ret[j].String(), "1081707507519149458") + } + resultCount++ + } + fmt.Println("Total result", resultCount) +} + +func TestBalanceScanner_MultiBatchCallERC20Only(t *testing.T) { + client, _ := ethclient.Dial("https://bsc-dataseed.binance.org/") + r, _ := rpc.Dial("https://bsc-dataseed.binance.org/") + bep20, _ := NewBalanceScanner("0xcA11bde05977b3631167028862bE2a173976CA11", client, r) + resultCount := 0 + ret, err := bep20.BatchCallBalancesWithNative(common.HexToAddress("0xca0C80122afA57c38BcAa14fC77E056b94288469"), [][]common.Address{ + { + common.HexToAddress("0x7d99eda556388Ad7743A1B658b9C4FC67D7A9d74"), + common.HexToAddress("0x7d99eda556388Ad7743A1B658b9C4FC67D7A9d74"), + common.HexToAddress("0x7d99eda556388Ad7743A1B658b9C4FC67D7A9d74"), + common.HexToAddress("0x7d99eda556388Ad7743A1B658b9C4FC67D7A9d74"), + common.HexToAddress("0x7d99eda556388Ad7743A1B658b9C4FC67D7A9d74"), + }, + }, "") + if err != nil { + fmt.Println(err) + } + for j := 0; j < len(ret); j++ { + assert.Equal(t, ret[j].String(), "1081707507519149458") + resultCount++ + } + fmt.Println("Total result", resultCount) +} + +func TestBalanceScanner_MultiBatchCallNativeOnly(t *testing.T) { + client, _ := ethclient.Dial("https://bsc-dataseed.binance.org/") + r, _ := rpc.Dial("https://bsc-dataseed.binance.org/") + bep20, _ := NewBalanceScanner("0xcA11bde05977b3631167028862bE2a173976CA11", client, r) + resultCount := 0 + ret, err := bep20.BatchCallBalancesWithNative(common.HexToAddress("0xca0C80122afA57c38BcAa14fC77E056b94288469"), [][]common.Address{ + { + common.HexToAddress("0x0000000000000000000000000000000000000000"), + }, + }, "0x0000000000000000000000000000000000000000") + if err != nil { + fmt.Println(err) + } + for j := 0; j < len(ret); j++ { + assert.Equal(t, ret[j].String(), "207924141787718960") + resultCount++ + } + fmt.Println("Total result", resultCount) +}