Optimize genesis and vault transfer

x/btcbridge/abci.go

@@ -9,6 +9,12 @@ import (
 
 // EndBlocker called at every block to handle DKG requests
 func EndBlocker(ctx sdk.Context, k keeper.Keeper) {
+	handleDKGRequests(ctx, k)
+	handleVaultTransfer(ctx, k)
+}
+
+// handleDKGRequests performs the DKG request handling
+func handleDKGRequests(ctx sdk.Context, k keeper.Keeper) {
 	pendingDKGRequests := k.GetPendingDKGRequests(ctx)
 
 	for _, req := range pendingDKGRequests {
@@ -40,34 +46,48 @@ func EndBlocker(ctx sdk.Context, k keeper.Keeper) {
 		// update status
 		req.Status = types.DKGRequestStatus_DKG_REQUEST_STATUS_COMPLETED
 		k.SetDKGRequest(ctx, req)
+	}
+}
 
-		// transfer vaults if the EnableTransfer flag set
+// handleVaultTransfer performs the vault asset transfer
+func handleVaultTransfer(ctx sdk.Context, k keeper.Keeper) {
+	completedDKGRequests := k.GetDKGRequests(ctx, types.DKGRequestStatus_DKG_REQUEST_STATUS_COMPLETED)
+
+	for _, req := range completedDKGRequests {
 		if req.EnableTransfer {
-			err := transferVaults(ctx, k, req.TargetUtxoNum, req.FeeRate)
+			completions := k.GetDKGCompletionRequests(ctx, req.Id)
+			dkgVaultVersion, _ := k.GetVaultVersionByAddress(ctx, completions[0].Vaults[0])
 
-			// reenable bridge when successfully transferred
-			if err == nil && req.DisableBridge {
-				k.EnableBridge(ctx)
-			}
-		}
-	}
-}
+			sourceVersion := dkgVaultVersion - 1
+			destVersion := k.GetLatestVaultVersion(ctx)
 
-// transferVaults performs the vault asset transfer (possibly partially)
-func transferVaults(ctx sdk.Context, k keeper.Keeper, targetUtxoNum uint32, feeRate string) error {
-	latestVaultVersion := k.GetLatestVaultVersion(ctx)
+			if k.VaultsTransferCompleted(ctx, sourceVersion) {
+				return
+			}
 
-	if err := k.TransferVault(ctx, latestVaultVersion-1, latestVaultVersion, types.AssetType_ASSET_TYPE_RUNES, nil, targetUtxoNum, feeRate); err != nil {
-		k.Logger(ctx).Error("transfer vault errored", "source version", latestVaultVersion-1, "destination version", latestVaultVersion, "asset type", types.AssetType_ASSET_TYPE_RUNES, "target utxo num", targetUtxoNum, "fee rate", feeRate, "err", err)
+			sourceBtcVault := k.GetVaultByAssetTypeAndVersion(ctx, types.AssetType_ASSET_TYPE_BTC, sourceVersion).Address
+			sourceRunesVault := k.GetVaultByAssetTypeAndVersion(ctx, types.AssetType_ASSET_TYPE_RUNES, sourceVersion).Address
 
-		return err
-	}
+			// transfer runes
+			if !k.VaultTransferCompleted(ctx, sourceRunesVault) {
+				if err := k.TransferVault(ctx, sourceVersion, destVersion, types.AssetType_ASSET_TYPE_RUNES, nil, req.TargetUtxoNum, req.FeeRate); err != nil {
+					k.Logger(ctx).Error("transfer vault errored", "source version", sourceVersion, "destination version", destVersion, "asset type", types.AssetType_ASSET_TYPE_RUNES, "target utxo num", req.TargetUtxoNum, "fee rate", req.FeeRate, "err", err)
+				}
+			}
 
-	if err := k.TransferVault(ctx, latestVaultVersion-1, latestVaultVersion, types.AssetType_ASSET_TYPE_BTC, nil, targetUtxoNum, feeRate); err != nil {
-		k.Logger(ctx).Error("transfer vault errored", "source version", latestVaultVersion-1, "destination version", latestVaultVersion, "asset type", types.AssetType_ASSET_TYPE_BTC, "target utxo num", targetUtxoNum, "fee rate", feeRate, "err", err)
+			// transfer btc only when runes transfer completed
+			if k.VaultTransferCompleted(ctx, sourceRunesVault) && !k.VaultTransferCompleted(ctx, sourceBtcVault) {
+				if err := k.TransferVault(ctx, sourceVersion, destVersion, types.AssetType_ASSET_TYPE_BTC, nil, req.TargetUtxoNum, req.FeeRate); err != nil {
+					k.Logger(ctx).Error("transfer vault errored", "source version", sourceVersion, "destination version", destVersion, "asset type", types.AssetType_ASSET_TYPE_BTC, "target utxo num", req.TargetUtxoNum, "fee rate", req.FeeRate, "err", err)
+				}
+			}
 
-		return err
+			// reenable bridge functions if disabled
+			if k.VaultsTransferCompleted(ctx, sourceVersion) {
+				if req.DisableBridge {
+					k.EnableBridge(ctx)
+				}
+			}
+		}
 	}
-
-	return nil
 }

x/btcbridge/genesis.go

@@ -1,6 +1,8 @@
 package btcbridge
 
 import (
+	"sort"
+
 	sdk "github.com/cosmos/cosmos-sdk/types"
 
 	"github.com/sideprotocol/side/x/btcbridge/keeper"
@@ -23,6 +25,14 @@ func InitGenesis(ctx sdk.Context, k keeper.Keeper, genState types.GenesisState)
 	for _, utxo := range genState.Utxos {
 		k.SaveUTXO(ctx, utxo)
 	}
+
+	// sort vaults and set the latest vault version
+	if len(genState.Params.Vaults) > 0 {
+		vaults := genState.Params.Vaults
+		sort.Slice(vaults, func(i, j int) bool { return vaults[i].Version < vaults[j].Version })
+
+		k.SetVaultVersion(ctx, vaults[len(vaults)-1].Version)
+	}
 }
 
 // ExportGenesis returns the module's exported genesis

x/btcbridge/keeper/params.go

@@ -48,6 +48,17 @@ func (k Keeper) GetVaultByAssetTypeAndVersion(ctx sdk.Context, assetType types.A
 	return nil
 }
 
+// GetVaultVersionByAddress gets the vault version of the given address
+func (k Keeper) GetVaultVersionByAddress(ctx sdk.Context, address string) (uint64, bool) {
+	for _, v := range k.GetParams(ctx).Vaults {
+		if v.Address == address {
+			return v.Version, true
+		}
+	}
+
+	return 0, false
+}
+
 // EnableBridge enables the bridge deposit and withdrawal
 func (k Keeper) EnableBridge(ctx sdk.Context) {
 	params := k.GetParams(ctx)

x/btcbridge/keeper/tss.go

@@ -593,3 +593,25 @@ func (k Keeper) GetLatestVaultVersion(ctx sdk.Context) uint64 {
 
 	return 0
 }
+
+// SetVaultVersion sets the vault version
+func (k Keeper) SetVaultVersion(ctx sdk.Context, version uint64) {
+	store := ctx.KVStore(k.storeKey)
+
+	store.Set(types.VaultVersionKey, sdk.Uint64ToBigEndian(version))
+}
+
+// VaultTransferCompleted returns true if the asset transfer completed for the given vault, false otherwise
+func (k Keeper) VaultTransferCompleted(ctx sdk.Context, vault string) bool {
+	count, _, _ := k.GetUnlockedUTXOCountAndBalancesByAddr(ctx, vault)
+
+	return count == 0
+}
+
+// VaultsTransferCompleted returns true if all asset transfer completed for the given vault version, false otherwise
+func (k Keeper) VaultsTransferCompleted(ctx sdk.Context, version uint64) bool {
+	btcVault := k.GetVaultByAssetTypeAndVersion(ctx, types.AssetType_ASSET_TYPE_BTC, version).Address
+	runesVault := k.GetVaultByAssetTypeAndVersion(ctx, types.AssetType_ASSET_TYPE_RUNES, version).Address
+
+	return k.VaultTransferCompleted(ctx, btcVault) && k.VaultTransferCompleted(ctx, runesVault)
+}