multi: add script key type filter to queries

Author: guggero

rpcserver.go

@@ -1010,7 +1010,7 @@ func (r *rpcServer) checkBalanceOverflow(ctx context.Context,
 	case assetID != nil:
 		// Retrieve the current asset balance.
 		balances, err := r.cfg.AssetStore.QueryBalancesByAsset(
-			ctx, assetID, true,
+			ctx, assetID, true, fn.None[asset.ScriptKeyType](),
 		)
 		if err != nil {
 			return fmt.Errorf("unable to query asset balance: %w",
@@ -1026,7 +1026,7 @@ func (r *rpcServer) checkBalanceOverflow(ctx context.Context,
 	case groupPubKey != nil:
 		// Retrieve the current balance of the group.
 		balances, err := r.cfg.AssetStore.QueryAssetBalancesByGroup(
-			ctx, groupPubKey, true,
+			ctx, groupPubKey, true, fn.None[asset.ScriptKeyType](),
 		)
 		if err != nil {
 			return fmt.Errorf("unable to query group balance: %w",
@@ -1105,9 +1105,18 @@ func (r *rpcServer) ListAssets(ctx context.Context,
 		filters.AnchorPoint = outPoint
 	}
 
+	scriptKeyType, includeSpent, err := taprpc.ParseScriptKeyTypeQuery(
+		req.ScriptKeyType,
+	)
+	if err != nil {
+		return nil, fmt.Errorf("unable to parse script key type "+
+			"query: %w", err)
+	}
+	filters.ScriptKeyType = scriptKeyType
+
 	rpcAssets, err := r.fetchRpcAssets(
-		ctx, req.WithWitness, req.IncludeSpent, req.IncludeLeased,
-		filters,
+		ctx, req.WithWitness, req.IncludeSpent || includeSpent,
+		req.IncludeLeased, filters,
 	)
 
 	if err != nil {
@@ -1218,11 +1227,12 @@ func (r *rpcServer) MarshalChainAsset(ctx context.Context, a asset.ChainAsset,
 }
 
 func (r *rpcServer) listBalancesByAsset(ctx context.Context,
-	assetID *asset.ID, includeLeased bool) (*taprpc.ListBalancesResponse,
+	assetID *asset.ID, includeLeased bool,
+	skt fn.Option[asset.ScriptKeyType]) (*taprpc.ListBalancesResponse,
 	error) {
 
 	balances, err := r.cfg.AssetStore.QueryBalancesByAsset(
-		ctx, assetID, includeLeased,
+		ctx, assetID, includeLeased, skt,
 	)
 	if err != nil {
 		return nil, fmt.Errorf("unable to list balances: %w", err)
@@ -1253,11 +1263,12 @@ func (r *rpcServer) listBalancesByAsset(ctx context.Context,
 }
 
 func (r *rpcServer) listBalancesByGroupKey(ctx context.Context,
-	groupKey *btcec.PublicKey,
-	includeLeased bool) (*taprpc.ListBalancesResponse, error) {
+	groupKey *btcec.PublicKey, includeLeased bool,
+	skt fn.Option[asset.ScriptKeyType]) (*taprpc.ListBalancesResponse,
+	error) {
 
 	balances, err := r.cfg.AssetStore.QueryAssetBalancesByGroup(
-		ctx, groupKey, includeLeased,
+		ctx, groupKey, includeLeased, skt,
 	)
 	if err != nil {
 		return nil, fmt.Errorf("unable to list balances: %w", err)
@@ -1292,8 +1303,22 @@ func (r *rpcServer) listBalancesByGroupKey(ctx context.Context,
 func (r *rpcServer) ListUtxos(ctx context.Context,
 	req *taprpc.ListUtxosRequest) (*taprpc.ListUtxosResponse, error) {
 
+	scriptKeyType, includeSpent, err := taprpc.ParseScriptKeyTypeQuery(
+		req.ScriptKeyType,
+	)
+	if err != nil {
+		return nil, fmt.Errorf("unable to parse script key type "+
+			"query: %w", err)
+	}
+
+	filters := &tapdb.AssetQueryFilters{
+		CommitmentConstraints: tapfreighter.CommitmentConstraints{
+			ScriptKeyType: scriptKeyType,
+		},
+	}
+
 	rpcAssets, err := r.fetchRpcAssets(
-		ctx, false, false, req.IncludeLeased, nil,
+		ctx, false, includeSpent, req.IncludeLeased, filters,
 	)
 	if err != nil {
 		return nil, err
@@ -1395,6 +1420,14 @@ func (r *rpcServer) ListGroups(ctx context.Context,
 func (r *rpcServer) ListBalances(ctx context.Context,
 	req *taprpc.ListBalancesRequest) (*taprpc.ListBalancesResponse, error) {
 
+	scriptKeyType, _, err := taprpc.ParseScriptKeyTypeQuery(
+		req.ScriptKeyType,
+	)
+	if err != nil {
+		return nil, fmt.Errorf("unable to parse script key type "+
+			"query: %w", err)
+	}
+
 	switch groupBy := req.GroupBy.(type) {
 	case *taprpc.ListBalancesRequest_AssetId:
 		if !groupBy.AssetId {
@@ -1411,7 +1444,9 @@ func (r *rpcServer) ListBalances(ctx context.Context,
 			copy(assetID[:], req.AssetFilter)
 		}
 
-		return r.listBalancesByAsset(ctx, assetID, req.IncludeLeased)
+		return r.listBalancesByAsset(
+			ctx, assetID, req.IncludeLeased, scriptKeyType,
+		)
 
 	case *taprpc.ListBalancesRequest_GroupKey:
 		if !groupBy.GroupKey {
@@ -1429,7 +1464,7 @@ func (r *rpcServer) ListBalances(ctx context.Context,
 		}
 
 		return r.listBalancesByGroupKey(
-			ctx, groupKey, req.IncludeLeased,
+			ctx, groupKey, req.IncludeLeased, scriptKeyType,
 		)
 
 	default:

tapchannel/aux_funding_controller.go

@@ -794,7 +794,7 @@ func (f *FundingController) fundVirtualPacket(ctx context.Context,
 	fundDesc := &tapsend.FundingDescriptor{
 		AssetSpecifier:    specifier,
 		Amount:            amt,
-		CoinSelectType:    tapsend.Bip86Only,
+		ScriptKeyType:     fn.Some(asset.ScriptKeyBip86),
 		DistinctSpecifier: true,
 	}
 

tapdb/asset_minting_test.go

@@ -1143,7 +1143,9 @@ func TestCommitBatchChainActions(t *testing.T) {
 
 	// We'll now query for the set of balances to ensure they all line up
 	// with the assets we just created, including the group genesis asset.
-	assetBalances, err := confAssets.QueryBalancesByAsset(ctx, nil, false)
+	assetBalances, err := confAssets.QueryBalancesByAsset(
+		ctx, nil, false, fn.None[asset.ScriptKeyType](),
+	)
 	require.NoError(t, err)
 	require.Equal(t, numSeedlings+1, len(assetBalances))
 
@@ -1165,7 +1167,7 @@ func TestCommitBatchChainActions(t *testing.T) {
 	}
 	numKeyGroups := fn.Reduce(mintedAssets, keyGroupSumReducer)
 	assetBalancesByGroup, err := confAssets.QueryAssetBalancesByGroup(
-		ctx, nil, false,
+		ctx, nil, false, fn.None[asset.ScriptKeyType](),
 	)
 	require.NoError(t, err)
 	require.Equal(t, numKeyGroups, len(assetBalancesByGroup))

tapdb/assets_store.go

@@ -24,7 +24,6 @@ import (
 	"github.com/lightninglabs/taproot-assets/tapfreighter"
 	"github.com/lightninglabs/taproot-assets/tappsbt"
 	"github.com/lightninglabs/taproot-assets/tapscript"
-	"github.com/lightninglabs/taproot-assets/tapsend"
 	"github.com/lightningnetwork/lnd/clock"
 	"github.com/lightningnetwork/lnd/keychain"
 )
@@ -897,13 +896,11 @@ func (a *AssetStore) constraintsToDbFilter(
 			assetFilter.AssetIDFilter = nil
 		}
 
-		switch query.CoinSelectType {
-		case tapsend.ScriptTreesAllowed:
-			assetFilter.Bip86ScriptKeysOnly = false
-
-		default:
-			assetFilter.Bip86ScriptKeysOnly = true
-		}
+		// The fn.None option means we don't restrict on script key type
+		// at all.
+		query.ScriptKeyType.WhenSome(func(t asset.ScriptKeyType) {
+			assetFilter.ScriptKeyType = sqlInt16(t)
+		})
 	}
 
 	return assetFilter, nil
@@ -990,8 +987,8 @@ type AssetQueryFilters struct {
 // QueryBalancesByAsset queries the balances for assets or alternatively
 // for a selected one that matches the passed asset ID filter.
 func (a *AssetStore) QueryBalancesByAsset(ctx context.Context,
-	assetID *asset.ID,
-	includeLeased bool) (map[asset.ID]AssetBalance, error) {
+	assetID *asset.ID, includeLeased bool,
+	skt fn.Option[asset.ScriptKeyType]) (map[asset.ID]AssetBalance, error) {
 
 	// We'll now map the application level filtering to the type of
 	// filtering our database query understands.
@@ -1002,6 +999,11 @@ func (a *AssetStore) QueryBalancesByAsset(ctx context.Context,
 		},
 	}
 
+	// The fn.None option means we don't restrict on script key type at all.
+	skt.WhenSome(func(t asset.ScriptKeyType) {
+		assetBalancesFilter.ScriptKeyType = sqlInt16(t)
+	})
+
 	// We exclude the assets that are specifically used for funding custom
 	// channels. The balance of those assets is reported through lnd channel
 	// balance. Those assets are identified by the funding script tree for a
@@ -1070,9 +1072,9 @@ func (a *AssetStore) QueryBalancesByAsset(ctx context.Context,
 // QueryAssetBalancesByGroup queries the asset balances for asset groups or
 // alternatively for a selected one that matches the passed filter.
 func (a *AssetStore) QueryAssetBalancesByGroup(ctx context.Context,
-	groupKey *btcec.PublicKey,
-	includeLeased bool) (map[asset.SerializedKey]AssetGroupBalance,
-	error) {
+	groupKey *btcec.PublicKey, includeLeased bool,
+	skt fn.Option[asset.ScriptKeyType]) (
+	map[asset.SerializedKey]AssetGroupBalance, error) {
 
 	// We'll now map the application level filtering to the type of
 	// filtering our database query understands.
@@ -1083,6 +1085,11 @@ func (a *AssetStore) QueryAssetBalancesByGroup(ctx context.Context,
 		},
 	}
 
+	// The fn.None option means we don't restrict on script key type at all.
+	skt.WhenSome(func(t asset.ScriptKeyType) {
+		assetBalancesFilter.ScriptKeyType = sqlInt16(t)
+	})
+
 	// We exclude the assets that are specifically used for funding custom
 	// channels. The balance of those assets is reported through lnd channel
 	// balance. Those assets are identified by the funding script tree for a

tapdb/assets_store_test.go

@@ -2937,11 +2937,11 @@ func TestQueryAssetBalances(t *testing.T) {
 	// At first, none of the assets should be leased.
 	includeLeased := false
 	balances, err := assetsStore.QueryBalancesByAsset(
-		ctx, nil, includeLeased,
+		ctx, nil, includeLeased, fn.None[asset.ScriptKeyType](),
 	)
 	require.NoError(t, err)
 	balancesByGroup, err := assetsStore.QueryAssetBalancesByGroup(
-		ctx, nil, includeLeased,
+		ctx, nil, includeLeased, fn.None[asset.ScriptKeyType](),
 	)
 	require.NoError(t, err)
 	require.Len(t, balances, numAssets)
@@ -2977,14 +2977,14 @@ func TestQueryAssetBalances(t *testing.T) {
 
 	// Only two assets should be returned that is not leased.
 	unleasedBalances, err := assetsStore.QueryBalancesByAsset(
-		ctx, nil, includeLeased,
+		ctx, nil, includeLeased, fn.None[asset.ScriptKeyType](),
 	)
 	require.NoError(t, err)
 	require.Len(t, unleasedBalances, numAssets-2)
 
 	// Only one group should be returned that is not leased.
 	unleasedBalancesByGroup, err := assetsStore.QueryAssetBalancesByGroup(
-		ctx, nil, includeLeased,
+		ctx, nil, includeLeased, fn.None[asset.ScriptKeyType](),
 	)
 	require.NoError(t, err)
 	require.Len(t, unleasedBalancesByGroup, numGroups-1)
@@ -3007,12 +3007,12 @@ func TestQueryAssetBalances(t *testing.T) {
 	// the same results as when the assets where unleased.
 	includeLeased = true
 	includeLeasedBalances, err := assetsStore.QueryBalancesByAsset(
-		ctx, nil, includeLeased,
+		ctx, nil, includeLeased, fn.None[asset.ScriptKeyType](),
 	)
 	require.NoError(t, err)
 	require.Len(t, includeLeasedBalances, numAssets)
 	includeLeasedBalByGroup, err := assetsStore.QueryAssetBalancesByGroup(
-		ctx, nil, includeLeased,
+		ctx, nil, includeLeased, fn.None[asset.ScriptKeyType](),
 	)
 	require.NoError(t, err)
 	require.Len(t, includeLeasedBalByGroup, len(assetGen.groupKeys))
@@ -3066,11 +3066,11 @@ func TestQueryAssetBalancesCustomChannelFunding(t *testing.T) {
 	// Hit both balance queries, they should return the same result.
 	includeLeased := false
 	balances, err := assetsStore.QueryBalancesByAsset(
-		ctx, nil, includeLeased,
+		ctx, nil, includeLeased, fn.None[asset.ScriptKeyType](),
 	)
 	require.NoError(t, err)
 	balancesByGroup, err := assetsStore.QueryAssetBalancesByGroup(
-		ctx, nil, includeLeased,
+		ctx, nil, includeLeased, fn.None[asset.ScriptKeyType](),
 	)
 	require.NoError(t, err)
 	require.Len(t, balances, numAssets-1)