address+tapdb: add AddrByScriptKeyAndVersion

Because we'll use the address' script key as the bare/raw public key
that will be tweaked for each individual output of an address v2 send,
we'll need to be able to find addresses by that script key.
We'll also make the script key unique for v2 addresses to avoid multiple
records being returned here.

Author: guggero

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,
@@ -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) {

tapdb/addrs.go

@@ -38,9 +38,15 @@ type (
 	// information.
 	Addresses = sqlc.FetchAddrsRow
 
-	// AddrByTaprootOutput is a type alias for returning an address by its
-	// Taproot output key.
-	AddrByTaprootOutput = sqlc.FetchAddrByTaprootOutputKeyRow
+	// SingleAddrQuery is a type alias for returning an address by its
+	// Taproot output key, x-only script key or version (or a combination of
+	// those).
+	SingleAddrQuery = sqlc.QueryAddrParams
+
+	// SingleAddrRow is a type alias for returning an address by either its
+	// Taproot output key, x-only script key or version (or a combination
+	// of those).
+	SingleAddrRow = sqlc.QueryAddrRow
 
 	// AddrManaged is a type alias for setting an address as managed.
 	AddrManaged = sqlc.SetAddrManagedParams
@@ -111,11 +117,11 @@ type AddrBook interface {
 	// passed AddrQuery.
 	FetchAddrs(ctx context.Context, arg AddrQuery) ([]Addresses, error)
 
-	// FetchAddrByTaprootOutputKey returns a single address based on its
-	// Taproot output key or a sql.ErrNoRows error if no such address
-	// exists.
-	FetchAddrByTaprootOutputKey(ctx context.Context,
-		arg []byte) (AddrByTaprootOutput, error)
+	// QueryAddr returns a single address based on its Taproot output key,
+	// its x-only script key or version, or a sql.ErrNoRows error if no such
+	// address exists.
+	QueryAddr(ctx context.Context, arg SingleAddrQuery) (SingleAddrRow,
+		error)
 
 	// UpsertAddr upserts a new address into the database returning the
 	// primary key.
@@ -509,8 +515,21 @@ func (t *TapAddressBook) AddrByTaprootOutput(ctx context.Context,
 		readOpts = NewAddrBookReadTx()
 	)
 	err := t.db.ExecTx(ctx, &readOpts, func(db AddrBook) error {
-		var err error
-		addr, err = fetchAddr(ctx, db, t.params, key)
+		row, err := db.QueryAddr(ctx, SingleAddrQuery{
+			TaprootOutputKey: schnorr.SerializePubKey(key),
+		})
+		switch {
+		case errors.Is(err, sql.ErrNoRows):
+			return address.ErrNoAddr
+
+		case err != nil:
+			return err
+		}
+
+		addr, err = parseAddr(
+			ctx, db, t.params, row.Addr, row.ScriptKey,
+			row.InternalKey, row.InternalKey_2,
+		)
 		return err
 	})
 	if err != nil {
@@ -520,22 +539,48 @@ func (t *TapAddressBook) AddrByTaprootOutput(ctx context.Context,
 	return addr, nil
 }
 
-// fetchAddr fetches a single address identified by its taproot output key from
-// the database and populates all its fields.
-func fetchAddr(ctx context.Context, db AddrBook, params *address.ChainParams,
-	taprootOutputKey *btcec.PublicKey) (*address.AddrWithKeyInfo, error) {
+// AddrByScriptKeyAndVersion returns a single address based on its script key
+// and version or a sql.ErrNoRows error if no such address exists.
+func (t *TapAddressBook) AddrByScriptKeyAndVersion(ctx context.Context,
+	scriptKey *btcec.PublicKey,
+	version address.Version) (*address.AddrWithKeyInfo, error) {
 
-	dbAddr, err := db.FetchAddrByTaprootOutputKey(
-		ctx, schnorr.SerializePubKey(taprootOutputKey),
+	var (
+		addr     *address.AddrWithKeyInfo
+		readOpts = NewAddrBookReadTx()
 	)
-	switch {
-	case errors.Is(err, sql.ErrNoRows):
-		return nil, address.ErrNoAddr
+	err := t.db.ExecTx(ctx, &readOpts, func(db AddrBook) error {
+		row, err := db.QueryAddr(ctx, SingleAddrQuery{
+			XOnlyScriptKey: schnorr.SerializePubKey(scriptKey),
+			Version:        sqlInt16(version),
+		})
+		switch {
+		case errors.Is(err, sql.ErrNoRows):
+			return address.ErrNoAddr
 
-	case err != nil:
+		case err != nil:
+			return err
+		}
+
+		addr, err = parseAddr(
+			ctx, db, t.params, row.Addr, row.ScriptKey,
+			row.InternalKey, row.InternalKey_2,
+		)
+		return err
+	})
+	if err != nil {
 		return nil, err
 	}
 
+	return addr, nil
+}
+
+// fetchAddr fetches a single address identified by its taproot output key from
+// the database and populates all its fields.
+func parseAddr(ctx context.Context, db AddrBook, params *address.ChainParams,
+	dbAddr sqlc.Addr, dbScriptKey sqlc.ScriptKey, dbInternalKey,
+	dbTaprootKey sqlc.InternalKey) (*address.AddrWithKeyInfo, error) {
+
 	genesis, err := fetchGenesis(ctx, db, dbAddr.GenesisAssetID)
 	if err != nil {
 		return nil, fmt.Errorf("error fetching genesis: %w", err)
@@ -566,21 +611,21 @@ func fetchAddr(ctx context.Context, db AddrBook, params *address.ChainParams,
 		}
 	}
 
-	scriptKey, err := parseScriptKey(dbAddr.InternalKey, dbAddr.ScriptKey)
+	scriptKey, err := parseScriptKey(dbInternalKey, dbScriptKey)
 	if err != nil {
 		return nil, fmt.Errorf("unable to decode script key: %w", err)
 	}
 
-	internalKey, err := btcec.ParsePubKey(dbAddr.RawTaprootKey)
+	internalKey, err := btcec.ParsePubKey(dbTaprootKey.RawKey)
 	if err != nil {
 		return nil, fmt.Errorf("unable to decode taproot key: %w", err)
 	}
 	internalKeyDesc := keychain.KeyDescriptor{
 		KeyLocator: keychain.KeyLocator{
 			Family: keychain.KeyFamily(
-				dbAddr.TaprootKeyFamily,
+				dbTaprootKey.KeyFamily,
 			),
-			Index: uint32(dbAddr.TaprootKeyIndex),
+			Index: uint32(dbTaprootKey.KeyIndex),
 		},
 		PubKey: internalKey,
 	}
@@ -612,6 +657,12 @@ func fetchAddr(ctx context.Context, db AddrBook, params *address.ChainParams,
 		return nil, fmt.Errorf("unable to make addr: %w", err)
 	}
 
+	taprootOutputKey, err := tapAddr.TaprootOutputKey()
+	if err != nil {
+		return nil, fmt.Errorf("unable to get taproot output key: %w",
+			err)
+	}
+
 	return &address.AddrWithKeyInfo{
 		Tap:              tapAddr,
 		ScriptKeyTweak:   *scriptKey.TweakedScriptKey,
@@ -866,11 +917,25 @@ func (t *TapAddressBook) QueryAddrEvents(
 					"output key: %w", err)
 			}
 
-			addr, err := fetchAddr(
-				ctx, db, t.params, taprootOutputKey,
+			row, err := db.QueryAddr(ctx, SingleAddrQuery{
+				TaprootOutputKey: schnorr.SerializePubKey(
+					taprootOutputKey,
+				),
+			})
+			switch {
+			case errors.Is(err, sql.ErrNoRows):
+				return address.ErrNoAddr
+
+			case err != nil:
+				return err
+			}
+
+			addr, err := parseAddr(
+				ctx, db, t.params, row.Addr, row.ScriptKey,
+				row.InternalKey, row.InternalKey_2,
 			)
 			if err != nil {
-				return fmt.Errorf("error fetching address: %w",
+				return fmt.Errorf("error parsing address: %w",
 					err)
 			}
 

tapdb/addrs_test.go

@@ -785,3 +785,87 @@ func TestScriptKeyTypeUpsert(t *testing.T) {
 		)
 	})
 }
+
+// TestQueryAddrEvents tests that we can query address events by their taproot
+// output key.
+func TestQueryAddrEvents(t *testing.T) {
+	t.Parallel()
+
+	// First, make a new addr book instance we'll use in the test below.
+	testClock := clock.NewTestClock(time.Now())
+	addrBook, _ := newAddrBook(t, testClock)
+
+	ctx := context.Background()
+
+	// Insert a test address and event into the database.
+	proofCourierAddr := address.RandProofCourierAddr(t)
+	addr, assetGen, assetGroup := address.RandAddr(
+		t, chainParams, proofCourierAddr,
+	)
+	err := addrBook.db.ExecTx(
+		ctx, WriteTxOption(),
+		insertFullAssetGen(ctx, assetGen, assetGroup),
+	)
+	require.NoError(t, err)
+
+	err = addrBook.InsertAddrs(ctx, *addr)
+	require.NoError(t, err)
+
+	tx := randWalletTx()
+	event, err := addrBook.GetOrCreateEvent(
+		ctx, address.StatusTransactionDetected, addr, tx, 0,
+	)
+	require.NoError(t, err)
+
+	// Query events for the address.
+	queryParams := address.EventQueryParams{
+		AddrTaprootOutputKey: schnorr.SerializePubKey(
+			&addr.TaprootOutputKey,
+		),
+	}
+	events, err := addrBook.QueryAddrEvents(ctx, queryParams)
+	require.NoError(t, err)
+	require.Len(t, events, 1)
+
+	// Verify the returned event matches the inserted event.
+	require.Equal(t, event.ID, events[0].ID)
+	require.Equal(t, event.Status, events[0].Status)
+	require.Equal(t, event.Outpoint, events[0].Outpoint)
+}
+
+// TestAddrByScriptKeyAndVersion tests that we can retrieve an address by its
+// script key and version.
+func TestAddrByScriptKeyAndVersion(t *testing.T) {
+	t.Parallel()
+
+	// First, make a new addr book instance we'll use in the test below.
+	testClock := clock.NewTestClock(time.Now())
+	addrBook, _ := newAddrBook(t, testClock)
+
+	ctx := context.Background()
+
+	// Insert a test address into the database.
+	proofCourierAddr := address.RandProofCourierAddr(t)
+	addr, assetGen, assetGroup := address.RandAddr(
+		t, chainParams, proofCourierAddr,
+	)
+	err := addrBook.db.ExecTx(
+		ctx, WriteTxOption(),
+		insertFullAssetGen(ctx, assetGen, assetGroup),
+	)
+	require.NoError(t, err)
+
+	err = addrBook.InsertAddrs(ctx, *addr)
+	require.NoError(t, err)
+
+	// Query the address by script key and version.
+	result, err := addrBook.AddrByScriptKeyAndVersion(
+		ctx, &addr.ScriptKey, addr.Version,
+	)
+	require.NoError(t, err)
+	require.NotNil(t, result)
+
+	// Verify the returned address matches the inserted address.
+	require.Equal(t, addr.ScriptKey, result.ScriptKey)
+	require.Equal(t, addr.Version, result.Version)
+}