WIP: fix bugs in crosschain transfer and support general assets

5 files changed, 90 insertions, 39 deletions

diff --git a/plugin/evm/export_tx.go b/plugin/evm/export_tx.go
index 0487c44..231afee 100644
--- a/plugin/evm/export_tx.go
+++ b/plugin/evm/export_tx.go
@@ -8,7 +8,6 @@ import (
 	"math/big"
 
 	"github.com/ava-labs/coreth/core/state"
-	"github.com/ethereum/go-ethereum/log"
 
 	"github.com/ava-labs/avalanchego/chains/atomic"
 	"github.com/ava-labs/avalanchego/database"
@@ -18,6 +17,7 @@ import (
 	safemath "github.com/ava-labs/avalanchego/utils/math"
 	"github.com/ava-labs/avalanchego/vms/components/avax"
 	"github.com/ava-labs/avalanchego/vms/secp256k1fx"
+	"github.com/ethereum/go-ethereum/log"
 )
 
 // UnsignedExportTx is an unsigned ExportTx
@@ -162,6 +162,7 @@ func (tx *UnsignedExportTx) Accept(ctx *snow.Context, _ database.Batch) error {
 
 // Create a new transaction
 func (vm *VM) newExportTx(
+	assetID ids.ID, // AssetID of the tokens to export
 	amount uint64, // Amount of tokens to export
 	chainID ids.ID, // Chain to send the UTXOs to
 	to ids.ShortID, // Address of chain recipient
@@ -171,15 +172,32 @@ func (vm *VM) newExportTx(
 		return nil, errWrongChainID
 	}
 
-	toBurn, err := safemath.Add64(amount, vm.txFee)
-	if err != nil {
-		return nil, errOverflowExport
+	var toBurn uint64
+	var err error
+	if assetID == vm.ctx.AVAXAssetID {
+		toBurn, err = safemath.Add64(amount, vm.txFee)
+		if err != nil {
+			return nil, errOverflowExport
+		}
+	} else {
+		toBurn = vm.txFee
 	}
-	ins, signers, err := vm.GetSpendableCanonical(keys, toBurn)
+	// burn AVAX
+	ins, signers, err := vm.GetSpendableCanonical(keys, vm.ctx.AVAXAssetID, toBurn)
 	if err != nil {
 		return nil, fmt.Errorf("couldn't generate tx inputs/outputs: %w", err)
 	}
 
+	// burn non-AVAX
+	if assetID != vm.ctx.AVAXAssetID {
+		ins2, signers2, err := vm.GetSpendableCanonical(keys, assetID, toBurn)
+		if err != nil {
+			return nil, fmt.Errorf("couldn't generate tx inputs/outputs: %w", err)
+		}
+		ins = append(ins, ins2...)
+		signers = append(signers, signers2...)
+	}
+
 	// Create the transaction
 	utx := &UnsignedExportTx{
 		NetworkID:        vm.ctx.NetworkID,
@@ -205,15 +223,19 @@ func (vm *VM) newExportTx(
 	return tx, utx.Verify(vm.ctx.XChainID, vm.ctx, vm.txFee, vm.ctx.AVAXAssetID)
 }
 
-func (tx *UnsignedExportTx) EVMStateTransfer(state *state.StateDB) error {
+func (tx *UnsignedExportTx) EVMStateTransfer(vm *VM, state *state.StateDB) error {
 	for _, from := range tx.Ins {
-		log.Info("consume", "in", from.Address, "amount", from.Amount, "nonce", from.Nonce, "nonce0", state.GetNonce(from.Address))
+		log.Info("crosschain C->X", "addr", from.Address, "amount", from.Amount)
 		amount := new(big.Int).Mul(
 			new(big.Int).SetUint64(from.Amount), x2cRate)
 		if state.GetBalance(from.Address).Cmp(amount) < 0 {
 			return errInsufficientFunds
 		}
-		state.SubBalance(from.Address, amount)
+		if from.AssetID == vm.ctx.AVAXAssetID {
+			state.SubBalance(from.Address, amount)
+		} else {
+			state.SubBalanceMultiCoin(from.Address, from.AssetID.Key(), amount)
+		}
 		if state.GetNonce(from.Address) != from.Nonce {
 			return errInvalidNonce
 		}
diff --git a/plugin/evm/import_tx.go b/plugin/evm/import_tx.go
index 7d17c4e..1736067 100644
--- a/plugin/evm/import_tx.go
+++ b/plugin/evm/import_tx.go
@@ -17,6 +17,7 @@ import (
 	"github.com/ava-labs/avalanchego/vms/components/avax"
 	"github.com/ava-labs/avalanchego/vms/secp256k1fx"
 	"github.com/ethereum/go-ethereum/common"
+	"github.com/ethereum/go-ethereum/log"
 )
 
 // UnsignedImportTx is an unsigned ImportTx
@@ -102,7 +103,7 @@ func (tx *UnsignedImportTx) SemanticVerify(
 	fc.Produce(vm.ctx.AVAXAssetID, vm.txFee)
 
 	for _, out := range tx.Outs {
-		fc.Produce(vm.ctx.AVAXAssetID, out.Amount)
+		fc.Produce(out.AssetID, out.Amount)
 	}
 
 	for _, in := range tx.ImportedInputs {
@@ -195,12 +196,9 @@ func (vm *VM) newImportTx(
 	importedInputs := []*avax.TransferableInput{}
 	signers := [][]*crypto.PrivateKeySECP256K1R{}
 
-	importedAmount := uint64(0)
+	importedAmount := make(map[ids.ID]uint64)
 	now := vm.clock.Unix()
 	for _, utxo := range atomicUTXOs {
-		if !utxo.AssetID().Equals(vm.ctx.AVAXAssetID) {
-			continue
-		}
 		inputIntf, utxoSigners, err := kc.Spend(utxo.Out, now)
 		if err != nil {
 			continue
@@ -209,7 +207,8 @@ func (vm *VM) newImportTx(
 		if !ok {
 			continue
 		}
-		importedAmount, err = math.Add64(importedAmount, input.Amount())
+		aid := utxo.AssetID()
+		importedAmount[aid], err = math.Add64(importedAmount[aid], input.Amount())
 		if err != nil {
 			return nil, err
 		}
@@ -221,25 +220,35 @@ func (vm *VM) newImportTx(
 		signers = append(signers, utxoSigners)
 	}
 	avax.SortTransferableInputsWithSigners(importedInputs, signers)
+	importedAVAXAmount := importedAmount[vm.ctx.AVAXAssetID]
 
-	if importedAmount == 0 {
+	if importedAVAXAmount == 0 {
 		return nil, errNoFunds // No imported UTXOs were spendable
 	}
 
-	nonce, err := vm.GetAcceptedNonce(to)
-	if err != nil {
-		return nil, err
-	}
-
 	outs := []EVMOutput{}
-	if importedAmount < vm.txFee { // imported amount goes toward paying tx fee
+
+	// AVAX output
+	if importedAVAXAmount < vm.txFee { // imported amount goes toward paying tx fee
 		// TODO: spend EVM balance to compensate vm.txFee-importedAmount
 		return nil, errNoFunds
-	} else if importedAmount > vm.txFee {
+	} else if importedAVAXAmount > vm.txFee {
+		outs = append(outs, EVMOutput{
+			Address: to,
+			Amount:  importedAVAXAmount - vm.txFee,
+			AssetID: vm.ctx.AVAXAssetID,
+		})
+	}
+
+	// non-AVAX asset outputs
+	for aid, amount := range importedAmount {
+		if aid.Equals(vm.ctx.AVAXAssetID) || amount == 0 {
+			continue
+		}
 		outs = append(outs, EVMOutput{
 			Address: to,
-			Amount:  importedAmount - vm.txFee,
-			Nonce:   nonce,
+			Amount:  amount,
+			AssetID: aid,
 		})
 	}
 
@@ -258,15 +267,16 @@ func (vm *VM) newImportTx(
 	return tx, utx.Verify(vm.ctx.XChainID, vm.ctx, vm.txFee, vm.ctx.AVAXAssetID)
 }
 
-func (tx *UnsignedImportTx) EVMStateTransfer(state *state.StateDB) error {
+func (tx *UnsignedImportTx) EVMStateTransfer(vm *VM, state *state.StateDB) error {
 	for _, to := range tx.Outs {
-		state.AddBalance(to.Address,
-			new(big.Int).Mul(
-				new(big.Int).SetUint64(to.Amount), x2cRate))
-		if state.GetNonce(to.Address) != to.Nonce {
-			return errInvalidNonce
+		log.Info("crosschain X->C", "addr", to.Address, "amount", to.Amount)
+		amount := new(big.Int).Mul(
+			new(big.Int).SetUint64(to.Amount), x2cRate)
+		if to.AssetID == vm.ctx.AVAXAssetID {
+			state.AddBalance(to.Address, amount)
+		} else {
+			state.AddBalanceMultiCoin(to.Address, to.AssetID.Key(), amount)
 		}
-		state.SetNonce(to.Address, to.Nonce+1)
 	}
 	return nil
 }
diff --git a/plugin/evm/service.go b/plugin/evm/service.go
index a3325b7..9f2e61a 100644
--- a/plugin/evm/service.go
+++ b/plugin/evm/service.go
@@ -15,6 +15,7 @@ import (
 	"github.com/ava-labs/coreth"
 
 	"github.com/ava-labs/avalanchego/api"
+	"github.com/ava-labs/avalanchego/ids"
 	"github.com/ava-labs/avalanchego/utils/constants"
 	"github.com/ava-labs/avalanchego/utils/crypto"
 	"github.com/ava-labs/avalanchego/utils/formatting"
@@ -267,7 +268,9 @@ func (service *AvaxAPI) ImportAVAX(_ *http.Request, args *ImportAVAXArgs, respon
 type ExportAVAXArgs struct {
 	api.UserPass
 
-	// Amount of AVAX to send
+	// AssetID of the tokens
+	AssetID ids.ID `json:"assetID"`
+	// Amount of asset to send
 	Amount json.Uint64 `json:"amount"`
 
 	// ID of the address that will receive the AVAX. This address includes the
@@ -302,8 +305,13 @@ func (service *AvaxAPI) ExportAVAX(_ *http.Request, args *ExportAVAXArgs, respon
 		return fmt.Errorf("couldn't get addresses controlled by the user: %w", err)
 	}
 
+	assetID := service.vm.ctx.AVAXAssetID
+	if !args.AssetID.IsZero() {
+		assetID = args.AssetID
+	}
 	// Create the transaction
 	tx, err := service.vm.newExportTx(
+		assetID,             // AssetID
 		uint64(args.Amount), // Amount
 		chainID,             // ID of the chain to send the funds to
 		to,                  // Address
diff --git a/plugin/evm/tx.go b/plugin/evm/tx.go
index e39a053..9580bc0 100644
--- a/plugin/evm/tx.go
+++ b/plugin/evm/tx.go
@@ -33,6 +33,13 @@ var (
 type EVMOutput struct {
 	Address common.Address `serialize:"true" json:"address"`
 	Amount  uint64         `serialize:"true" json:"amount"`
+	AssetID ids.ID         `serialize:"true" json:"assetID"`
+}
+
+type EVMInput struct {
+	Address common.Address `serialize:"true" json:"address"`
+	Amount  uint64         `serialize:"true" json:"amount"`
+	AssetID ids.ID         `serialize:"true" json:"assetID"`
 	Nonce   uint64         `serialize:"true" json:"nonce"`
 }
 
@@ -40,8 +47,6 @@ func (out *EVMOutput) Verify() error {
 	return nil
 }
 
-type EVMInput EVMOutput
-
 func (in *EVMInput) Verify() error {
 	return nil
 }
@@ -66,7 +71,7 @@ type UnsignedAtomicTx interface {
 	// Accept this transaction with the additionally provided state transitions.
 	Accept(ctx *snow.Context, batch database.Batch) error
 
-	EVMStateTransfer(state *state.StateDB) error
+	EVMStateTransfer(vm *VM, state *state.StateDB) error
 }
 
 // Tx is a signed transaction
diff --git a/plugin/evm/vm.go b/plugin/evm/vm.go
index 200a08d..1f025ee 100644
--- a/plugin/evm/vm.go
+++ b/plugin/evm/vm.go
@@ -281,7 +281,7 @@ func (vm *VM) Initialize(
 	chain.SetOnFinalizeAndAssemble(func(state *state.StateDB, txs []*types.Transaction) ([]byte, error) {
 		select {
 		case atx := <-vm.pendingAtomicTxs:
-			if err := atx.UnsignedTx.(UnsignedAtomicTx).EVMStateTransfer(state); err != nil {
+			if err := atx.UnsignedTx.(UnsignedAtomicTx).EVMStateTransfer(vm, state); err != nil {
 				vm.newBlockChan <- nil
 				return nil, err
 			}
@@ -323,7 +323,7 @@ func (vm *VM) Initialize(
 		if tx == nil {
 			return nil
 		}
-		return tx.UnsignedTx.(UnsignedAtomicTx).EVMStateTransfer(state)
+		return tx.UnsignedTx.(UnsignedAtomicTx).EVMStateTransfer(vm, state)
 	})
 	vm.blockCache = cache.LRU{Size: blockCacheSize}
 	vm.blockStatusCache = cache.LRU{Size: blockCacheSize}
@@ -882,7 +882,7 @@ func PublicKeyToEthAddress(pubKey crypto.PublicKey) common.Address {
 		(*pubKey.(*crypto.PublicKeySECP256K1R).ToECDSA()))
 }
 
-func (vm *VM) GetSpendableCanonical(keys []*crypto.PrivateKeySECP256K1R, amount uint64) ([]EVMInput, [][]*crypto.PrivateKeySECP256K1R, error) {
+func (vm *VM) GetSpendableCanonical(keys []*crypto.PrivateKeySECP256K1R, assetID ids.ID, amount uint64) ([]EVMInput, [][]*crypto.PrivateKeySECP256K1R, error) {
 	// NOTE: should we use HEAD block or lastAccepted?
 	state, err := vm.chain.BlockState(vm.lastAccepted.ethBlock)
 	if err != nil {
@@ -895,7 +895,12 @@ func (vm *VM) GetSpendableCanonical(keys []*crypto.PrivateKeySECP256K1R, amount
 			break
 		}
 		addr := GetEthAddress(key)
-		balance := new(big.Int).Div(state.GetBalance(addr), x2cRate).Uint64()
+		var balance uint64
+		if assetID == vm.ctx.AVAXAssetID {
+			balance = new(big.Int).Div(state.GetBalance(addr), x2cRate).Uint64()
+		} else {
+			balance = new(big.Int).Div(state.GetBalanceMultiCoin(addr, assetID.Key()), x2cRate).Uint64()
+		}
 		if balance == 0 {
 			continue
 		}
@@ -909,6 +914,7 @@ func (vm *VM) GetSpendableCanonical(keys []*crypto.PrivateKeySECP256K1R, amount
 		inputs = append(inputs, EVMInput{
 			Address: addr,
 			Amount:  balance,
+			AssetID: assetID,
 			Nonce:   nonce,
 		})
 		signers = append(signers, []*crypto.PrivateKeySECP256K1R{key})