key-manager/service/signer.go

package service

import (
	"context"
	"crypto/ecdsa"
	"crypto/sha256"
	"encoding/hex"
	"errors"
	"fmt"
	"github.com/ethereum/go-ethereum/common"
	"github.com/ethereum/go-ethereum/core/types"
	"github.com/ethereum/go-ethereum/crypto"
	tronaddr "github.com/fbsobreira/gotron-sdk/pkg/address"
	troncore "github.com/fbsobreira/gotron-sdk/pkg/proto/core"
	"google.golang.org/protobuf/proto"
	"key-manager/dao"
	"math/big"
)

func format(network, address string) (string, error) {
	if network == eth {
		address = common.HexToAddress(address).String()
	} else {

		addr, err := tronaddr.Base58ToAddress(address)
		if err != nil {
			return "", err
		}
		address = addr.String()
	}
	return address, nil
}

func (km *KeyManager) loadSigner(network, signer string) error {
	var err error
	signer, err = format(network, signer)
	if err != nil {
		return err
	}

	if _, ok := km.signers[signer]; !ok {
		var info dao.AccountInfo
		if network == eth {
			info, err = km.dao.GetEthFromAddress(signer)
			if err != nil {
				return err
			}
		} else {
			info, err = km.dao.GetTronFromAddress(signer)
			if err != nil {
				return err
			}
		}

		if _, exist := km.mnemonics[info.GetName()]; !exist {
			return errors.New("user mnemonic lost")
		}

		account, err := createKey(km.mnemonics[info.GetName()], network, info.GetIndex())
		if err != nil {
			return err
		}
		km.signers[signer] = account.GetPriKey()
	}

	return nil
}

func (km *KeyManager) Sign(ctx context.Context, req *SignRequest) (*SignResponse, error) {
	km.lk.Lock()
	defer km.lk.Unlock()

	if err := checkNetwork(req.Network); err != nil {
		return &SignResponse{
			Code:     networkErrCode,
			Msg:      networkErrMsg,
			Network:  req.Network,
			Sender:   req.Sender,
			SignedTx: "",
		}, nil
	}

	if err := km.loadSigner(req.Network, req.Sender); err != nil {
		return &SignResponse{
			Code:     loadSignerErrCode,
			Msg:      err.Error(),
			Network:  req.Network,
			Sender:   req.Sender,
			SignedTx: "",
		}, nil
	}

	var marshalTx string
	var coin string
	var to string
	var amount string

	if req.Network == eth {
		tx, err := unmarshalJEthTx(req.Tx)
		if err != nil {
			return &SignResponse{
				Code:     unmarshalJEthTxErrCode,
				Msg:      err.Error(),
				Network:  req.Network,
				Sender:   req.Sender,
				SignedTx: "",
			}, nil
		}

		coin, to, amount, err = checkErc(tx, km.whitelists)
		if err != nil {
			return &SignResponse{
				Code:     checkErcTxErrCode,
				Msg:      err.Error(),
				Network:  req.Network,
				Sender:   req.Sender,
				SignedTx: "",
			}, nil
		}

		signedTx, err := signEthTransaction(tx, km.signers[req.Sender])
		if err != nil {
			return &SignResponse{
				Code:     signEthTxErrCode,
				Msg:      err.Error(),
				Network:  req.Network,
				Sender:   req.Sender,
				SignedTx: "",
			}, nil
		}

		marshalTx, err = marshalJEthTx(signedTx)
		if err != nil {
			return &SignResponse{
				Code:     marshalJEthTxErrCode,
				Msg:      err.Error(),
				Network:  req.Network,
				Sender:   req.Sender,
				SignedTx: "",
			}, nil
		}
	} else {
		tx, err := unmarshalJTronTx(req.Tx)
		if err != nil {
			return &SignResponse{
				Code:     unmarshalJTronTxErrCode,
				Msg:      err.Error(),
				Network:  req.Network,
				Sender:   req.Sender,
				SignedTx: "",
			}, nil
		}

		coin, to, amount, err = checkTron(tx, km.whitelists)
		if err != nil {
			return &SignResponse{
				Code:     checkTronTxErrCode,
				Msg:      err.Error(),
				Network:  req.Network,
				Sender:   req.Sender,
				SignedTx: "",
			}, nil
		}

		signedTx, err := signTronTransaction(tx, km.signers[req.Sender])
		if err != nil {
			return &SignResponse{
				Code:     signTronTxErrCode,
				Msg:      err.Error(),
				Network:  req.Network,
				Sender:   req.Sender,
				SignedTx: "",
			}, nil
		}

		marshalTx, err = marshalJTronTx(signedTx)
		if err != nil {
			return &SignResponse{
				Code:     marshalJTronTxErrCode,
				Msg:      err.Error(),
				Network:  req.Network,
				Sender:   req.Sender,
				SignedTx: "",
			}, nil
		}
	}

	err := km.dao.CreateSignTx(&dao.SignTx{
		Network: req.Network,
		Coin:    coin,
		From:    req.Sender,
		To:      to,
		Amount:  amount,
	})

	if err != nil {
		log.Warnw("save Sign Tx to mysql failed", "network", req.Network, "coin", coin, "from", req.Sender, "to", to, "amount", amount, "err", err)
	}

	return &SignResponse{
		Code:     okCode,
		Msg:      okMsg,
		Network:  req.Network,
		Sender:   req.Sender,
		SignedTx: marshalTx,
	}, nil
}

func marshalJEthTx(transaction *types.Transaction) (string, error) {
	b, err := transaction.MarshalJSON()
	if err != nil {
		return "", err
	}

	return hex.EncodeToString(b), nil
}

func unmarshalJEthTx(tx string) (*types.Transaction, error) {
	b, err := hex.DecodeString(tx)
	if err != nil {
		return nil, err
	}

	var transaction = &types.Transaction{}
	err = transaction.UnmarshalJSON(b)
	if err != nil {
		return nil, err
	}

	return transaction, err
}

func marshalJTronTx(transaction *troncore.Transaction) (string, error) {
	b, err := proto.Marshal(transaction)
	if err != nil {
		return "", err
	}

	return hex.EncodeToString(b), nil
}

func unmarshalJTronTx(tx string) (*troncore.Transaction, error) {
	b, err := hex.DecodeString(tx)
	if err != nil {
		return nil, err
	}

	var transaction troncore.Transaction
	err = proto.Unmarshal(b, &transaction)
	if err != nil {
		return nil, err
	}

	return &transaction, nil
}

func signEthTransaction(transaction *types.Transaction, priv *ecdsa.PrivateKey) (*types.Transaction, error) {
	signer := types.NewLondonSigner(big.NewInt(1))
	var err error
	transaction, err = types.SignTx(transaction, signer, priv)
	if err != nil {
		return nil, err
	}

	return transaction, nil
}

func signTronTransaction(transaction *troncore.Transaction, priv *ecdsa.PrivateKey) (*troncore.Transaction, error) {
	rawData, err := proto.Marshal(transaction.GetRawData())
	if err != nil {
		return nil, fmt.Errorf("proto marshal tx raw data error: %v", err)
	}
	h256h := sha256.New()
	h256h.Write(rawData)
	hash := h256h.Sum(nil)
	signature, err := crypto.Sign(hash, priv)
	if err != nil {
		return nil, fmt.Errorf("sign error: %v", err)
	}
	transaction.Signature = append(transaction.Signature, signature)
	return transaction, nil
}