key-manager/service/whitelist.go

package service

import (
	"bytes"
	"encoding/hex"
	"errors"
	"fmt"
	"github.com/ethereum/go-ethereum/accounts/abi"
	"github.com/ethereum/go-ethereum/common"
	"github.com/ethereum/go-ethereum/core/types"
	tronaddr "github.com/fbsobreira/gotron-sdk/pkg/address"
	troncore "github.com/fbsobreira/gotron-sdk/pkg/proto/core"
	"google.golang.org/protobuf/reflect/protoreflect"
	"math/big"
)

const (
	ethUsdt  = "0xdAC17F958D2ee523a2206206994597C13D831ec7"
	ethUsdc  = "0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48"
	tronUsdt = "TR7NHqjeKQxGTCi8q8ZY4pL8otSzgjLj6t"
	tronUsdc = "TEkxiTehnzSmSe2XqrBj4w32RUN966rdz8"
)

var (
	transferId, _ = hex.DecodeString("a9059cbb")
	approveId, _  = hex.DecodeString("095ea7b3")
)

var (
	addressTy, _ = abi.NewType("address", "", nil)
	uint256Ty, _ = abi.NewType("uint256", "", nil)
	arguments    = abi.Arguments{{Type: addressTy}, {Type: uint256Ty}}
)

func checkErc(tx *types.Transaction, whitelists map[string]struct{}) (string, string, string, error) {
	if tx.To().String() == ethUsdt || tx.To().String() == ethUsdc {
		input := tx.Data()
		var methodId [4]byte
		copy(methodId[:], input[:4])

		if !(bytes.Equal(methodId[:], transferId) || bytes.Equal(methodId[:], approveId)) {
			return "", "", "", errors.New("erc20: only sign transfer and approve")
		}
		unpack, err := arguments.Unpack(input[4:])
		if err != nil {
			return "", "", "", err
		}

		to, ok := unpack[0].(common.Address)
		if !ok {
			return "", "", "", errors.New("erc20: to parameter error")
		}

		amount, ok := unpack[1].(*big.Int)
		if !ok {
			return "", "", "", errors.New("erc20: amount parameter error")
		}

		if _, exist := whitelists[to.String()]; !exist {
			return "", "", "", errors.New("erc20: to is not a whitelist address")
		}

		coin := usdtCoin
		if tx.To().String() == ethUsdc {
			coin = usdcCoin
		}

		return coin, to.String(), amount.String(), nil
	} else {
		if _, exist := whitelists[tx.To().String()]; !exist {
			return "", "", "", errors.New("eth: to is not a whitelist address")
		}
	}

	return ethCoin, tx.To().String(), tx.Value().String(), nil
}

func checkTron(tx *troncore.Transaction, whitelists map[string]struct{}) (string, string, string, error) {
	if len(tx.GetRawData().Contract) != 1 {
		return "", "", "", errors.New("Transaction_Contract length must be 1")
	}
	contract := tx.GetRawData().Contract[0]

	var c interface{}
	switch contract.Type {
	case troncore.Transaction_Contract_TransferContract:
		c = &troncore.TransferContract{}
	case troncore.Transaction_Contract_TriggerSmartContract:
		c = &troncore.TriggerSmartContract{}
	default:
		return "", "", "", errors.New("Transaction_Contract Type must be Transaction_Contract_TransferContract and Transaction_Contract_TriggerSmartContract")
	}

	if err := contract.GetParameter().UnmarshalTo(c.(protoreflect.ProtoMessage)); err != nil {
		return "", "", "", fmt.Errorf("proto unmarshal any: %+w", err)
	}

	if contract.Type == troncore.Transaction_Contract_TransferContract {
		transferContract := c.(*troncore.TransferContract)
		to := tronaddr.Address(transferContract.ToAddress)
		if _, exist := whitelists[to.String()]; !exist {
			return "", "", "", errors.New("trx: to is not a whitelist address")
		}

		return trxCoin, to.String(), big.NewInt(0).SetInt64(transferContract.Amount).String(), nil
	} else {
		triggerSmartContract := c.(*troncore.TriggerSmartContract)
		contractAddr := tronaddr.Address(triggerSmartContract.ContractAddress)
		if contractAddr.String() != tronUsdt && contractAddr.String() != tronUsdc {
			return "", "", "", errors.New("trc20: only sign usdt and usdc")
		}

		input := triggerSmartContract.Data
		var methodId [4]byte
		copy(methodId[:], input[:4])

		if !(bytes.Equal(methodId[:], transferId) || bytes.Equal(methodId[:], approveId)) {
			return "", "", "", errors.New("trc20: only sign transfer and approve")
		}

		unpack, err := arguments.Unpack(input[4:])
		if err != nil {
			return "", "", "", err
		}

		to, ok := unpack[0].(common.Address)
		if !ok {
			return "", "", "", errors.New("trc20: to parameter error")
		}

		amount, ok := unpack[1].(*big.Int)
		if !ok {
			return "", "", "", errors.New("erc20: amount parameter error")
		}

		toTron := make([]byte, 0)
		toTron = append(toTron, tronaddr.TronBytePrefix)
		toTron = append(toTron, to.Bytes()...)
		if _, exist := whitelists[tronaddr.Address(toTron).String()]; !exist {
			return "", "", "", errors.New("trc20: to is not a whitelist address")
		}

		coin := usdtCoin
		if contractAddr.String() == tronUsdc {
			coin = usdcCoin
		}

		return coin, to.String(), amount.String(), nil
	}
}