tapdb: change script key exclusion mechanism

With an explicit type for a script key now being used, we can turn the
key based matching into a type based matching, which will also be more
generic for unique funding script keys that are required for grouped
asset channel funding.

Author: guggero

tapdb/assets_store.go

@@ -23,7 +23,6 @@ import (
 	"github.com/lightninglabs/taproot-assets/tapdb/sqlc"
 	"github.com/lightninglabs/taproot-assets/tapfreighter"
 	"github.com/lightninglabs/taproot-assets/tappsbt"
-	"github.com/lightninglabs/taproot-assets/tapscript"
 	"github.com/lightningnetwork/lnd/clock"
 	"github.com/lightningnetwork/lnd/keychain"
 )
@@ -999,19 +998,25 @@ func (a *AssetStore) QueryBalancesByAsset(ctx context.Context,
 		},
 	}
 
+	// 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 specific script key type
+	// for channel keys. We exclude them unless explicitly queried for.
+	assetBalancesFilter.ExcludeScriptKeyType = sqlInt16(
+		asset.ScriptKeyScriptPathChannel,
+	)
+
 	// 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
-	// custom channel asset-level script key.
-	excludeKey := asset.NewScriptKey(
-		tapscript.NewChannelFundingScriptTree().TaprootKey,
-	)
-	assetBalancesFilter.ExcludeKey = excludeKey.PubKey.SerializeCompressed()
+		// If the user explicitly wants to see the channel related asset
+		// balances, we need to set the exclude type to NULL.
+		if t == asset.ScriptKeyScriptPathChannel {
+			nullValue := sql.NullInt16{}
+			assetBalancesFilter.ExcludeScriptKeyType = nullValue
+		}
+	})
 
 	// By default, we only show assets that are not leased.
 	if !includeLeased {
@@ -1085,19 +1090,25 @@ func (a *AssetStore) QueryAssetBalancesByGroup(ctx context.Context,
 		},
 	}
 
+	// 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 specific script key type
+	// for channel keys. We exclude them unless explicitly queried for.
+	assetBalancesFilter.ExcludeScriptKeyType = sqlInt16(
+		asset.ScriptKeyScriptPathChannel,
+	)
+
 	// 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
-	// custom channel asset-level script key.
-	excludeKey := asset.NewScriptKey(
-		tapscript.NewChannelFundingScriptTree().TaprootKey,
-	)
-	assetBalancesFilter.ExcludeKey = excludeKey.PubKey.SerializeCompressed()
+		// If the user explicitly wants to see the channel related asset
+		// balances, we need to set the exclude type to NULL.
+		if t == asset.ScriptKeyScriptPathChannel {
+			nullValue := sql.NullInt16{}
+			assetBalancesFilter.ExcludeScriptKeyType = nullValue
+		}
+	})
 
 	// By default, we only show assets that are not leased.
 	if !includeLeased {

tapdb/assets_store_test.go

@@ -11,6 +11,7 @@ import (
 	"time"
 
 	"github.com/btcsuite/btcd/btcec/v2"
+	"github.com/btcsuite/btcd/btcec/v2/schnorr"
 	"github.com/btcsuite/btcd/chaincfg/chainhash"
 	"github.com/btcsuite/btcd/txscript"
 	"github.com/btcsuite/btcd/wire"
@@ -3027,23 +3028,33 @@ func TestQueryAssetBalancesCustomChannelFunding(t *testing.T) {
 	const numGroups = 1
 	assetGen := newAssetGenerator(t, numAssets, numGroups)
 
-	fundingScriptKey := asset.NewScriptKey(
-		tapscript.NewChannelFundingScriptTree().TaprootKey,
+	fundingKey := tapscript.NewChannelFundingScriptTree()
+	xOnlyKey, _ := schnorr.ParsePubKey(
+		schnorr.SerializePubKey(fundingKey.TaprootKey),
 	)
+	fundingScriptKey := asset.ScriptKey{
+		PubKey: xOnlyKey,
+		TweakedScriptKey: &asset.TweakedScriptKey{
+			RawKey: keychain.KeyDescriptor{
+				PubKey: fundingKey.InternalKey,
+			},
+			Type: asset.ScriptKeyScriptPathChannel,
+		},
+	}
 
 	assetDesc := []assetDesc{
 		{
 			assetGen:    assetGen.assetGens[0],
 			anchorPoint: assetGen.anchorPoints[0],
 			keyGroup:    assetGen.groupKeys[0],
-			amt:         4,
+			amt:         8,
 			scriptKey:   &fundingScriptKey,
 		},
 		{
 			assetGen:    assetGen.assetGens[1],
 			anchorPoint: assetGen.anchorPoints[1],
 			keyGroup:    assetGen.groupKeys[0],
-			amt:         4,
+			amt:         12,
 		},
 	}
 	assetGen.genAssets(t, assetsStore, assetDesc)
@@ -3074,4 +3085,31 @@ func TestQueryAssetBalancesCustomChannelFunding(t *testing.T) {
 		balanceByGroupSum += balance.Balance
 	}
 	require.Equal(t, assetDesc[1].amt, balanceByGroupSum)
+
+	// If we explicitly query for channel related script keys, we should get
+	// just those assets.
+	balances, err = assetsStore.QueryBalancesByAsset(
+		ctx, nil, includeLeased,
+		fn.Some(asset.ScriptKeyScriptPathChannel),
+	)
+	require.NoError(t, err)
+	balancesByGroup, err = assetsStore.QueryAssetBalancesByGroup(
+		ctx, nil, includeLeased,
+		fn.Some(asset.ScriptKeyScriptPathChannel),
+	)
+	require.NoError(t, err)
+	require.Len(t, balances, numAssets-1)
+	require.Len(t, balancesByGroup, numAssets-1)
+
+	balanceSum = uint64(0)
+	for _, balance := range balances {
+		balanceSum += balance.Balance
+	}
+	require.Equal(t, assetDesc[0].amt, balanceSum)
+
+	balanceByGroupSum = uint64(0)
+	for _, balance := range balancesByGroup {
+		balanceByGroupSum += balance.Balance
+	}
+	require.Equal(t, assetDesc[0].amt, balanceByGroupSum)
 }

tapdb/sqlc/assets.sql.go

@@ -339,8 +339,8 @@ JOIN managed_utxos utxos
 JOIN script_keys
     ON assets.script_key_id = script_keys.script_key_id
 WHERE spent = FALSE AND 
-      (script_keys.tweaked_script_key != sqlc.narg('exclude_key') OR
-        sqlc.narg('exclude_key') IS NULL) AND
+      (script_keys.key_type != sqlc.narg('exclude_script_key_type') OR
+        sqlc.narg('exclude_script_key_type') IS NULL) AND
       (sqlc.narg('script_key_type') = script_keys.key_type OR 
         sqlc.narg('script_key_type') IS NULL)
 GROUP BY assets.genesis_id, genesis_info_view.asset_id,
@@ -369,9 +369,9 @@ JOIN managed_utxos utxos
        END
 JOIN script_keys
     ON assets.script_key_id = script_keys.script_key_id
-WHERE spent = FALSE AND 
-      (script_keys.tweaked_script_key != sqlc.narg('exclude_key') OR
-        sqlc.narg('exclude_key') IS NULL) AND
+WHERE spent = FALSE AND
+      (script_keys.key_type != sqlc.narg('exclude_script_key_type') OR
+        sqlc.narg('exclude_script_key_type') IS NULL) AND
       (sqlc.narg('script_key_type') = script_keys.key_type OR
         sqlc.narg('script_key_type') IS NULL)
 GROUP BY key_group_info_view.tweaked_group_key;