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",
@@ -1059,8 +1059,17 @@ func (r *rpcServer) ListAssets(ctx context.Context,
 			"and include_leased")
 	}
 
+	scriptKeyType, includeSpent, err := taprpc.ParseScriptKeyTypeQuery(
+		req.ScriptKeyType,
+	)
+	if err != nil {
+		return nil, fmt.Errorf("unable to parse script key type "+
+			"query: %w", err)
+	}
+
 	rpcAssets, err := r.fetchRpcAssets(
-		ctx, req.WithWitness, req.IncludeSpent, req.IncludeLeased,
+		ctx, req.WithWitness, req.IncludeSpent || includeSpent,
+		req.IncludeLeased, scriptKeyType,
 	)
 	if err != nil {
 		return nil, err
@@ -1114,10 +1123,16 @@ func (r *rpcServer) ListAssets(ctx context.Context,
 }
 
 func (r *rpcServer) fetchRpcAssets(ctx context.Context, withWitness,
-	includeSpent, includeLeased bool) ([]*taprpc.Asset, error) {
+	includeSpent, includeLeased bool,
+	scriptKeyType fn.Option[asset.ScriptKeyType]) ([]*taprpc.Asset, error) {
 
+	filter := &tapdb.AssetQueryFilters{
+		CommitmentConstraints: tapfreighter.CommitmentConstraints{
+			ScriptKeyType: scriptKeyType,
+		},
+	}
 	assets, err := r.cfg.AssetStore.FetchAllAssets(
-		ctx, includeSpent, includeLeased, nil,
+		ctx, includeSpent, includeLeased, filter,
 	)
 	if err != nil {
 		return nil, fmt.Errorf("unable to read chain assets: %w", err)
@@ -1169,11 +1184,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)
@@ -1204,11 +1220,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)
@@ -1243,7 +1260,17 @@ func (r *rpcServer) listBalancesByGroupKey(ctx context.Context,
 func (r *rpcServer) ListUtxos(ctx context.Context,
 	req *taprpc.ListUtxosRequest) (*taprpc.ListUtxosResponse, error) {
 
-	rpcAssets, err := r.fetchRpcAssets(ctx, false, false, req.IncludeLeased)
+	scriptKeyType, includeSpent, err := taprpc.ParseScriptKeyTypeQuery(
+		req.ScriptKeyType,
+	)
+	if err != nil {
+		return nil, fmt.Errorf("unable to parse script key type "+
+			"query: %w", err)
+	}
+
+	rpcAssets, err := r.fetchRpcAssets(
+		ctx, false, includeSpent, req.IncludeLeased, scriptKeyType,
+	)
 	if err != nil {
 		return nil, err
 	}
@@ -1344,6 +1371,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 {
@@ -1360,7 +1395,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 {
@@ -1378,7 +1415,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

@@ -806,7 +806,7 @@ func (f *FundingController) fundVirtualPacket(ctx context.Context,
 	fundDesc := &tapsend.FundingDescriptor{
 		AssetSpecifier: specifier,
 		Amount:         amt,
-		CoinSelectType: tapsend.Bip86Only,
+		ScriptKeyType:  fn.Some(asset.ScriptKeyBip86),
 	}
 
 	// Fund the packet. This will derive an anchor internal key for us, but

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"
 )
@@ -873,13 +872,11 @@ func (a *AssetStore) constraintsToDbFilter(
 		// TODO(roasbeef): only want to allow asset ID or other and not
 		// both?
 
-		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
@@ -959,8 +956,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.
@@ -971,6 +968,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
@@ -1039,9 +1041,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.
@@ -1052,6 +1054,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

@@ -2628,11 +2628,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)
@@ -2668,14 +2668,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)
@@ -2698,12 +2698,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))
@@ -2757,11 +2757,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)