Improve vault transfer

proto/side/btcbridge/btcbridge.proto

@@ -115,10 +115,18 @@ message DKGRequest {
   uint32 threshold = 3;
   // asset types of vaults to be generated
   repeated AssetType vault_types = 4;
+   // indicates if disabling bridge deposit and withdrawal
+  bool disable_bridge = 5;
+  // indicates if transferring assets to the newly generated vaults when the DKG request is completed
+  bool enable_transfer = 6;
+  // target number of the UTXOs to be transferred each time
+  uint32 target_utxo_num = 7;
+  // fee rate for vault transfer
+  string fee_rate = 8;
   // expiration time
-  google.protobuf.Timestamp expiration = 5 [(gogoproto.stdtime) = true, (gogoproto.nullable) = true];
+  google.protobuf.Timestamp expiration = 9 [(gogoproto.stdtime) = true, (gogoproto.nullable) = true];
   // status
-  DKGRequestStatus status = 6;
+  DKGRequestStatus status = 10;
 }
 
 // DKG Completion Request

proto/side/btcbridge/query.proto

@@ -51,6 +51,10 @@ service Query {
   rpc QueryUTXOsByAddress(QueryUTXOsByAddressRequest) returns (QueryUTXOsByAddressResponse) {
     option (google.api.http).get = "/side/btcbridge/utxos/{address}";
   }
+  // QueryUTXOCountAndBalancesByAddress queries the total count and balances of the unlocked utxos by the given address.
+  rpc QueryUTXOCountAndBalancesByAddress(QueryUTXOCountAndBalancesByAddressRequest) returns (QueryUTXOCountAndBalancesByAddressResponse) {
+    option (google.api.http).get = "/side/btcbridge/utxos/{address}/stats";
+  }
   // QueryDKGRequest queries the DKG request by the given id.
   rpc QueryDKGRequest(QueryDKGRequestRequest) returns (QueryDKGRequestResponse) {
     option (google.api.http).get = "/side/btcbridge/dkg/request/{id}";
@@ -171,6 +175,18 @@ message QueryUTXOsByAddressResponse {
   repeated UTXO utxos = 1;
 }
 
+// QueryUTXOCountAndBalancesByAddressRequest is the request type for the Query/UTXOCountAndBalancesByAddress RPC method.
+message QueryUTXOCountAndBalancesByAddressRequest {
+  string address = 1;
+}
+
+// QueryUTXOCountAndBalancesByAddressResponse is the response type for the Query/UTXOCountAndBalancesByAddress RPC method.
+message QueryUTXOCountAndBalancesByAddressResponse {
+  uint32 count = 1;
+  int64 value = 2;
+  repeated RuneBalance runeBalances = 3;
+}
+
 // QueryDKGRequestRequest is the request type for the Query/DKGRequest RPC method.
 message QueryDKGRequestRequest {
   uint64 id = 1;

proto/side/btcbridge/tx.proto

@@ -120,6 +120,14 @@ message MsgInitiateDKG {
   uint32 threshold = 3;
   // asset types of vaults to be generated
   repeated AssetType vault_types = 4;
+  // indicates if disabling bridge functionalities including deposit and withdrawal
+  bool disable_bridge = 5;
+  // indicates if transferring the current vaults to the newly generated vaults when the DKG request is completed
+  bool enable_transfer = 6;
+  // target number of the UTXOs to be transferred each time
+  uint32 target_utxo_num = 7;
+  // fee rate for vault transfer
+  string fee_rate = 8;
 }
 
 // MsgInitiateDKGResponse defines the Msg/InitiateDKG response type.
@@ -152,8 +160,12 @@ message MsgTransferVault {
   uint64 dest_version = 3;
   // asset type
   AssetType asset_type = 4;
-  // a set of pre-built PSBTs to perform the asset transfer
+  // a set of optional pre-built PSBTs to perform the asset transfer
   repeated string psbts = 5;
+  // target number of the UTXOs to be transferred; only take effect when psbt not provided
+  uint32 target_utxo_num = 6;
+  // fee rate; only take effect when psbt not provided
+  string fee_rate = 7;
 }
 
 // MsgTransferVaultResponse defines the Msg/TransferVault response type.

x/btcbridge/abci.go

@@ -40,5 +40,34 @@ 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
+		if req.EnableTransfer {
+			err := transferVaults(ctx, k, req.TargetUtxoNum, req.FeeRate)
+
+			// reenable bridge when successfully transferred
+			if err == nil && req.DisableBridge {
+				k.EnableBridge(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 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)
+
+		return 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)
+
+		return err
 	}
+
+	return nil
 }

x/btcbridge/client/cli/query.go

@@ -31,8 +31,9 @@ func GetQueryCmd(_ string) *cobra.Command {
 	cmd.AddCommand(CmdQueryParams())
 	cmd.AddCommand(CmdBestBlock())
 	cmd.AddCommand(CmdQueryBlock())
-	cmd.AddCommand(CmdQueryWithdrawRequest())
+	cmd.AddCommand(CmdQueryWithdrawRequests())
 	cmd.AddCommand(CmdQueryUTXOs())
+	cmd.AddCommand(CmdQueryUTXOStats())
 	cmd.AddCommand(CmdQueryDKGRequests())
 	cmd.AddCommand(CmdQueryDKGCompletionRequests())
 	// this line is used by starport scaffolding # 1
@@ -94,7 +95,7 @@ func CmdBestBlock() *cobra.Command {
 	return cmd
 }
 
-// CmdQueryBlock returns the command to query the heights of the light client
+// CmdQueryBlock returns the command to query the block by hash or height
 func CmdQueryBlock() *cobra.Command {
 	cmd := &cobra.Command{
 		Use:   "block [hash or height]",
@@ -132,10 +133,10 @@ func CmdQueryBlock() *cobra.Command {
 	return cmd
 }
 
-// CmdQueryWithdrawRequest returns the command to query withdrawal request
-func CmdQueryWithdrawRequest() *cobra.Command {
+// CmdQueryWithdrawRequests returns the command to query withdrawal requests
+func CmdQueryWithdrawRequests() *cobra.Command {
 	cmd := &cobra.Command{
-		Use:   "withdraw-request [status | address | tx hash]",
+		Use:   "withdraw-requests [status | address | tx hash]",
 		Short: "Query withdrawal requests by status, address or tx hash",
 		Args:  cobra.ExactArgs(1),
 		RunE: func(cmd *cobra.Command, args []string) error {
@@ -185,6 +186,7 @@ func CmdQueryWithdrawRequest() *cobra.Command {
 	return cmd
 }
 
+// CmdQueryUTXOs returns the command to query utxos by the optional address
 func CmdQueryUTXOs() *cobra.Command {
 	cmd := &cobra.Command{
 		Use:   "utxos [address]",
@@ -228,6 +230,41 @@ func CmdQueryUTXOs() *cobra.Command {
 	return cmd
 }
 
+// CmdQueryUTXOStats returns the command to query the utxo statistics by address
+func CmdQueryUTXOStats() *cobra.Command {
+	cmd := &cobra.Command{
+		Use:   "utxo-stats [address]",
+		Short: "Query the utxo statistics by the given address",
+		Args:  cobra.ExactArgs(1),
+		RunE: func(cmd *cobra.Command, args []string) error {
+			clientCtx, err := client.GetClientQueryContext(cmd)
+			if err != nil {
+				return err
+			}
+
+			queryClient := types.NewQueryClient(clientCtx)
+
+			_, err = sdk.AccAddressFromBech32(args[0])
+			if err != nil {
+				return err
+			}
+
+			res, err := queryClient.QueryUTXOCountAndBalancesByAddress(cmd.Context(), &types.QueryUTXOCountAndBalancesByAddressRequest{
+				Address: args[0],
+			})
+			if err != nil {
+				return err
+			}
+
+			return clientCtx.PrintProto(res)
+		},
+	}
+
+	flags.AddQueryFlagsToCmd(cmd)
+
+	return cmd
+}
+
 // CmdQueryDKGRequests returns the command to query DKG requests
 func CmdQueryDKGRequests() *cobra.Command {
 	cmd := &cobra.Command{

x/btcbridge/keeper/msg_server.go

@@ -222,7 +222,7 @@ func (m msgServer) InitiateDKG(goCtx context.Context, msg *types.MsgInitiateDKG)
 
 	ctx := sdk.UnwrapSDKContext(goCtx)
 
-	req, err := m.Keeper.InitiateDKG(ctx, msg.Participants, msg.Threshold, msg.VaultTypes)
+	req, err := m.Keeper.InitiateDKG(ctx, msg.Participants, msg.Threshold, msg.VaultTypes, msg.DisableBridge, msg.EnableTransfer, msg.TargetUtxoNum, msg.FeeRate)
 	if err != nil {
 		return nil, err
 	}
@@ -276,7 +276,7 @@ func (m msgServer) TransferVault(goCtx context.Context, msg *types.MsgTransferVa
 
 	ctx := sdk.UnwrapSDKContext(goCtx)
 
-	if err := m.Keeper.TransferVault(ctx, msg.SourceVersion, msg.DestVersion, msg.AssetType, msg.Psbts); err != nil {
+	if err := m.Keeper.TransferVault(ctx, msg.SourceVersion, msg.DestVersion, msg.AssetType, msg.Psbts, msg.TargetUtxoNum, msg.FeeRate); err != nil {
 		return nil, err
 	}
 

x/btcbridge/keeper/params.go

@@ -47,3 +47,23 @@ func (k Keeper) GetVaultByAssetTypeAndVersion(ctx sdk.Context, assetType types.A
 
 	return nil
 }
+
+// EnableBridge enables the bridge deposit and withdrawal
+func (k Keeper) EnableBridge(ctx sdk.Context) {
+	params := k.GetParams(ctx)
+
+	params.DepositEnabled = true
+	params.WithdrawEnabled = true
+
+	k.SetParams(ctx, params)
+}
+
+// DisableBridge disables the bridge deposit and withdrawal
+func (k Keeper) DisableBridge(ctx sdk.Context) {
+	params := k.GetParams(ctx)
+
+	params.DepositEnabled = false
+	params.WithdrawEnabled = false
+
+	k.SetParams(ctx, params)
+}

x/btcbridge/keeper/queries.go

@@ -165,6 +165,26 @@ func (k Keeper) QueryUTXOsByAddress(goCtx context.Context, req *types.QueryUTXOs
 	return &types.QueryUTXOsByAddressResponse{Utxos: utxos}, nil
 }
 
+func (k Keeper) QueryUTXOCountAndBalancesByAddress(goCtx context.Context, req *types.QueryUTXOCountAndBalancesByAddressRequest) (*types.QueryUTXOCountAndBalancesByAddressResponse, error) {
+	if req == nil {
+		return nil, status.Error(codes.InvalidArgument, "invalid request")
+	}
+	ctx := sdk.UnwrapSDKContext(goCtx)
+
+	_, err := sdk.AccAddressFromBech32(req.Address)
+	if err != nil {
+		return nil, err
+	}
+
+	count, value, runeBalances := k.GetUnlockedUTXOCountAndBalancesByAddr(ctx, req.Address)
+
+	return &types.QueryUTXOCountAndBalancesByAddressResponse{
+		Count:        count,
+		Value:        value,
+		RuneBalances: runeBalances,
+	}, nil
+}
+
 func (k Keeper) QueryDKGRequest(goCtx context.Context, req *types.QueryDKGRequestRequest) (*types.QueryDKGRequestResponse, error) {
 	if req == nil {
 		return nil, status.Error(codes.InvalidArgument, "invalid request")