Merge pull request #1621 from lightninglabs/pedersen-keys-refactor

[group key addrs 5/6]: Add Pedersen unique script key type

Author: Roasbeef

address/book.go

@@ -127,6 +127,12 @@ type Storage interface {
 	AddrByTaprootOutput(ctx context.Context,
 		key *btcec.PublicKey) (*AddrWithKeyInfo, error)
 
+	// AddrByScriptKeyAndVersion returns a single address based on its
+	// script key and version or a sql.ErrNoRows error if no such address
+	// exists.
+	AddrByScriptKeyAndVersion(context.Context, *btcec.PublicKey,
+		Version) (*AddrWithKeyInfo, error)
+
 	// SetAddrManaged sets an address as being managed by the internal
 	// wallet.
 	SetAddrManaged(ctx context.Context, addr *AddrWithKeyInfo,
@@ -504,7 +510,7 @@ func (b *Book) NewAddressWithKeys(ctx context.Context, addrVersion Version,
 
 	// We might not know the type of script key, if it was given to us
 	// through an RPC call. So we make a guess here.
-	keyType := scriptKey.DetermineType()
+	keyType := scriptKey.DetermineType(fn.Ptr(assetGroup.Genesis.ID()))
 
 	err = b.cfg.Store.InsertScriptKey(ctx, scriptKey, keyType)
 	if err != nil {
@@ -603,6 +609,14 @@ func (b *Book) AddrByTaprootOutput(ctx context.Context,
 	return b.cfg.Store.AddrByTaprootOutput(ctx, key)
 }
 
+// AddrByScriptKeyAndVersion returns a single address based on its script key
+// and version or a sql.ErrNoRows error if no such address exists.
+func (b *Book) AddrByScriptKeyAndVersion(ctx context.Context,
+	scriptKey *btcec.PublicKey, version Version) (*AddrWithKeyInfo, error) {
+
+	return b.cfg.Store.AddrByScriptKeyAndVersion(ctx, scriptKey, version)
+}
+
 // SetAddrManaged sets an address as being managed by the internal
 // wallet.
 func (b *Book) SetAddrManaged(ctx context.Context, addr *AddrWithKeyInfo,

address/book_test.go

@@ -175,6 +175,13 @@ type MockStorage struct {
 	mock.Mock
 }
 
+func (m *MockStorage) AddrByScriptKeyAndVersion(ctx context.Context,
+	key *btcec.PublicKey, version Version) (*AddrWithKeyInfo, error) {
+
+	args := m.Called(ctx, key, version)
+	return args.Get(0).(*AddrWithKeyInfo), args.Error(1)
+}
+
 func (m *MockStorage) GetOrCreateEvent(ctx context.Context, status Status,
 	addr *AddrWithKeyInfo, walletTx *lndclient.Transaction,
 	outputIdx uint32) (*Event, error) {

asset/asset.go

@@ -150,8 +150,85 @@ const (
 	// Keys related to channels are not shown in asset balances (unless
 	// specifically requested) and are _never_ used for coin selection.
 	ScriptKeyScriptPathChannel ScriptKeyType = 5
+
+	// ScriptKeyUniquePedersen is the script key type used for assets that
+	// use a unique script key, tweaked with a Pedersen commitment key in a
+	// single Tapscript leaf. This is used to avoid collisions in the
+	// universe when there are multiple grouped asset UTXOs within the same
+	// on-chain output.
+	ScriptKeyUniquePedersen ScriptKeyType = 6
+)
+
+var (
+	// AllScriptKeyTypes is a slice of all known script key types.
+	AllScriptKeyTypes = []ScriptKeyType{
+		ScriptKeyUnknown,
+		ScriptKeyBip86,
+		ScriptKeyScriptPathExternal,
+		ScriptKeyBurn,
+		ScriptKeyTombstone,
+		ScriptKeyScriptPathChannel,
+		ScriptKeyUniquePedersen,
+	}
+
+	// ScriptKeyTypesNoChannel is a slice of all known script key types
+	// that are not related to channels. This is used to filter out channel
+	// related script keys when querying for assets that are not related to
+	// channels.
+	ScriptKeyTypesNoChannel = []ScriptKeyType{
+		ScriptKeyUnknown,
+		ScriptKeyBip86,
+		ScriptKeyScriptPathExternal,
+		ScriptKeyBurn,
+		ScriptKeyTombstone,
+		ScriptKeyUniquePedersen,
+	}
 )
 
+// ScriptKeyTypeForDatabaseQuery returns a slice of script key types that should
+// be used when querying the database for assets. The returned slice will either
+// contain all script key types or only those that are not related to channels,
+// depending on the `excludeChannelRelated` parameter. Unless the user specifies
+// a specific script key type, in which case the returned slice will only
+// contain that specific script key type.
+func ScriptKeyTypeForDatabaseQuery(excludeChannelRelated bool,
+	userSpecified fn.Option[ScriptKeyType]) []ScriptKeyType {
+
+	// If the user specified a script key type, we use that directly to
+	// filter the results.
+	if userSpecified.IsSome() {
+		specifiedType := userSpecified.UnwrapOr(ScriptKeyUnknown)
+		dbTypes := []ScriptKeyType{
+			specifiedType,
+		}
+
+		// If the user specifically requested BIP-86 script keys, we
+		// also include the Pedersen unique script key type, because
+		// those can be spent the same way as BIP-86 script keys, and
+		// they should be treated the same way as BIP-86 script keys.
+		if specifiedType == ScriptKeyBip86 {
+			dbTypes = append(dbTypes, ScriptKeyUniquePedersen)
+		}
+
+		return dbTypes
+	}
+
+	// For some queries, we want to get all the assets with all possible
+	// script key types. For those, we use the full set of script key types.
+	dbTypes := fn.CopySlice(AllScriptKeyTypes)
+
+	// For some RPCs (mostly balance related), we exclude the assets that
+	// are specifically used for funding custom channels by default. 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.
+	if excludeChannelRelated {
+		dbTypes = fn.CopySlice(ScriptKeyTypesNoChannel)
+	}
+
+	return dbTypes
+}
+
 var (
 	// ZeroPrevID is the blank prev ID used for genesis assets and also
 	// asset split leaves.
@@ -392,7 +469,7 @@ func NewSpecifierFromGroupKey(groupPubKey btcec.PublicKey) Specifier {
 	}
 }
 
-// NewExlusiveSpecifier creates a specifier that may only include one of asset
+// NewExclusiveSpecifier creates a specifier that may only include one of asset
 // ID or group key. If both are set then a specifier over the group key is
 // created.
 func NewExclusiveSpecifier(id *ID,
@@ -1043,6 +1120,65 @@ func EqualKeyDescriptors(a, o keychain.KeyDescriptor) bool {
 	return a.PubKey.IsEqual(o.PubKey)
 }
 
+// ScriptKeyDerivationMethod is the method used to derive the script key of an
+// asset send output from the recipient's internal key and the asset ID of
+// the output. This is used to ensure that the script keys are unique for each
+// asset ID, so that proofs can be fetched from the universe without collisions.
+type ScriptKeyDerivationMethod uint8
+
+const (
+	// ScriptKeyDerivationUniquePedersen means the script key is derived
+	// using the address's recipient ID key and a single leaf that contains
+	// an un-spendable Pedersen commitment key
+	// (OP_CHECKSIG <NUMS_key + asset_id * G>). This can be used to
+	// create unique script keys for each virtual packet in the fragment,
+	// to avoid proof collisions in the universe, where the script keys
+	// should be spendable by a hardware wallet that only supports
+	// miniscript policies for signing P2TR outputs.
+	ScriptKeyDerivationUniquePedersen ScriptKeyDerivationMethod = 0
+)
+
+// DeriveUniqueScriptKey derives a unique script key for the given asset ID
+// using the recipient's internal key and the specified derivation method.
+func DeriveUniqueScriptKey(internalKey btcec.PublicKey, assetID ID,
+	method ScriptKeyDerivationMethod) (ScriptKey, error) {
+
+	switch method {
+	// For the unique Pedersen method, we derive the script key using the
+	// internal key and the asset ID using a Pedersen commitment key in a
+	// single OP_CHECKSIG leaf.
+	case ScriptKeyDerivationUniquePedersen:
+		leaf, err := NewNonSpendableScriptLeaf(
+			PedersenVersion, assetID[:],
+		)
+		if err != nil {
+			return ScriptKey{}, fmt.Errorf("unable to create "+
+				"non-spendable leaf: %w", err)
+		}
+
+		rootHash := leaf.TapHash()
+		scriptPubKey, _ := schnorr.ParsePubKey(schnorr.SerializePubKey(
+			txscript.ComputeTaprootOutputKey(
+				&internalKey, rootHash[:],
+			),
+		))
+		return ScriptKey{
+			PubKey: scriptPubKey,
+			TweakedScriptKey: &TweakedScriptKey{
+				RawKey: keychain.KeyDescriptor{
+					PubKey: &internalKey,
+				},
+				Tweak: rootHash[:],
+				Type:  ScriptKeyUniquePedersen,
+			},
+		}, nil
+
+	default:
+		return ScriptKey{}, fmt.Errorf("unknown script key derivation "+
+			"method: %d", method)
+	}
+}
+
 // TweakedScriptKey is an embedded struct which is primarily used by wallets to
 // be able to keep track of the tweak of a script key alongside the raw key
 // derivation information.
@@ -1142,14 +1278,16 @@ func (s *ScriptKey) HasScriptPath() bool {
 }
 
 // DetermineType attempts to determine the type of the script key based on the
-// information available. This method will only return ScriptKeyUnknown if the
-// following condition is met:
+// information available. This method will only return ScriptKeyUnknown if one
+// of the following conditions is met:
 //   - The script key doesn't have a script path, but the final Taproot output
 //     key doesn't match a BIP-0086 key derived from the internal key. This will
 //     be the case for "foreign" script keys we import from proofs, where we set
 //     the internal key to the same key as the tweaked script key (because we
 //     don't know the internal key, as it's not part of the proof encoding).
-func (s *ScriptKey) DetermineType() ScriptKeyType {
+//   - No asset ID was provided (because it is unavailable in the given
+//     context), and the script key is a unique Pedersen-based key.
+func (s *ScriptKey) DetermineType(id *ID) ScriptKeyType {
 	// If we have an explicit script key type set, we can return that.
 	if s.TweakedScriptKey != nil &&
 		s.TweakedScriptKey.Type != ScriptKeyUnknown {
@@ -1178,6 +1316,24 @@ func (s *ScriptKey) DetermineType() ScriptKeyType {
 		if bip86.PubKey.IsEqual(s.PubKey) {
 			return ScriptKeyBip86
 		}
+
+		// If we have the asset's ID, we can check whether this is a
+		// Pedersen-based key. If we don't have the ID, then we can't
+		// determine the type, so we'll end up in the default return
+		// below.
+		if id != nil {
+			scriptKey, err := DeriveUniqueScriptKey(
+				*s.TweakedScriptKey.RawKey.PubKey, *id,
+				ScriptKeyDerivationUniquePedersen,
+			)
+			if err != nil {
+				return ScriptKeyUnknown
+			}
+
+			if scriptKey.PubKey.IsEqual(s.PubKey) {
+				return ScriptKeyUniquePedersen
+			}
+		}
 	}
 
 	return ScriptKeyUnknown

authmailbox/client_test.go

@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"context"
 	"fmt"
+	"net/url"
 	"os"
 	"testing"
 	"time"
@@ -236,6 +237,40 @@ func TestServerClientAuthAndRestart(t *testing.T) {
 		client2.stop(t)
 	})
 
+	// We also add a multi-subscription to the same two keys, so we can make
+	// sure we can receive messages from multiple clients at once.
+	multiSub := NewMultiSubscription(*clientCfg)
+	err := multiSub.Subscribe(
+		ctx, url.URL{Host: clientCfg.ServerAddress}, clientKey1, filter,
+	)
+	require.NoError(t, err)
+	err = multiSub.Subscribe(
+		ctx, url.URL{Host: clientCfg.ServerAddress}, clientKey2, filter,
+	)
+	require.NoError(t, err)
+	t.Cleanup(func() {
+		require.NoError(t, multiSub.Stop())
+	})
+	msgChan := multiSub.MessageChan()
+	readMultiSub := func(targetID ...uint64) {
+		t.Helper()
+		select {
+		case inboundMsgs := <-msgChan:
+			receivedIDs := fn.Map(
+				inboundMsgs.Messages,
+				func(msg *mboxrpc.MailboxMessage) uint64 {
+					return msg.MessageId
+				},
+			)
+			for _, target := range targetID {
+				require.Contains(t, receivedIDs, target)
+			}
+		case <-time.After(testTimeout):
+			t.Fatalf("timeout waiting for message with ID %v",
+				targetID)
+		}
+	}
+
 	// Send a message to all clients.
 	msg1 := &Message{
 		ID:               1000,
@@ -244,14 +279,15 @@ func TestServerClientAuthAndRestart(t *testing.T) {
 	}
 
 	// We also store the message in the store, so we can retrieve it later.
-	_, err := harness.mockMsgStore.StoreMessage(ctx, randOp, msg1)
+	_, err = harness.mockMsgStore.StoreMessage(ctx, randOp, msg1)
 	require.NoError(t, err)
 
 	harness.srv.publishMessage(msg1)
 
 	// We should be able to receive that message.
 	client1.readMessages(t, msg1.ID)
 	client2.readMessages(t, msg1.ID)
+	readMultiSub(msg1.ID)
 
 	// We now stop the server and assert that the subscription is no longer
 	// active.
@@ -282,6 +318,7 @@ func TestServerClientAuthAndRestart(t *testing.T) {
 	// We should be able to receive that message.
 	client1.readMessages(t, msg2.ID)
 	client2.readMessages(t, msg2.ID)
+	readMultiSub(msg2.ID)
 
 	// If we now start a third client, we should be able to receive all
 	// three messages, given we are using the same key and specify the
@@ -314,6 +351,23 @@ func TestServerClientAuthAndRestart(t *testing.T) {
 	harness.srv.publishMessage(msg3)
 	client4.expectNoMessage(t)
 	client1.readMessages(t, msg3.ID)
+	client2.readMessages(t, msg3.ID)
+	client3.readMessages(t, msg3.ID)
+	readMultiSub(msg3.ID)
+
+	// Let's make sure that a message sent to the second key is only
+	// received by the fourth client and the multi-subscription.
+	msg4 := &Message{
+		ID:               1001,
+		ReceiverKey:      *clientKey2.PubKey,
+		ArrivalTimestamp: time.Now(),
+	}
+	harness.srv.publishMessage(msg4)
+	client1.expectNoMessage(t)
+	client2.expectNoMessage(t)
+	client3.expectNoMessage(t)
+	client4.readMessages(t, msg4.ID)
+	readMultiSub(msg4.ID)
 }
 
 // TestSendMessage tests the SendMessage RPC of the server and its ability to

authmailbox/mock.go

@@ -87,5 +87,8 @@ func (s *MockMsgStore) QueryMessages(_ context.Context,
 }
 
 func (s *MockMsgStore) NumMessages(context.Context) uint64 {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
 	return uint64(len(s.messages))
 }