From beba62b62b2391373176ec50d371257e2abaf7f1 Mon Sep 17 00:00:00 2001
From: yyforyongyu <yong2452@gmail.com>
Date: Thu, 8 Aug 2024 15:13:23 +0800
Subject: [PATCH] wallet: add new method `FetchDerivationInfo`

---
 wallet/utxos.go      | 54 ++++++++++++++++++++++++++++----------------
 wallet/utxos_test.go | 35 ++++++++++++++++++++++++++++
 2 files changed, 69 insertions(+), 20 deletions(-)

diff --git a/wallet/utxos.go b/wallet/utxos.go
index 34da168ec9..819b638cfd 100644
--- a/wallet/utxos.go
+++ b/wallet/utxos.go
@@ -107,6 +107,8 @@ func (w *Wallet) UnspentOutputs(policy OutputSelectionPolicy) ([]*TransactionOut
 // full transaction, the target txout, the derivation info and the number of
 // confirmations are returned. Otherwise, a non-nil error value of ErrNotMine
 // is returned instead.
+//
+// NOTE: This method is kept for compatibility.
 func (w *Wallet) FetchInputInfo(prevOut *wire.OutPoint) (*wire.MsgTx,
 	*wire.TxOut, *psbt.Bip32Derivation, int64, error) {
 
@@ -115,30 +117,11 @@ func (w *Wallet) FetchInputInfo(prevOut *wire.OutPoint) (*wire.MsgTx,
 		return nil, nil, nil, 0, err
 	}
 
-	pkScript := txOut.PkScript
-	addr, err := w.fetchOutputAddr(pkScript)
+	derivation, err := w.FetchDerivationInfo(txOut.PkScript)
 	if err != nil {
 		return nil, nil, nil, 0, err
 	}
 
-	pubKeyAddr, ok := addr.(waddrmgr.ManagedPubKeyAddress)
-	if !ok {
-		return nil, nil, nil, 0, ErrNotMine
-	}
-	keyScope, derivationPath, _ := pubKeyAddr.DerivationInfo()
-
-	derivation := &psbt.Bip32Derivation{
-		PubKey:               pubKeyAddr.PubKey().SerializeCompressed(),
-		MasterKeyFingerprint: derivationPath.MasterKeyFingerprint,
-		Bip32Path: []uint32{
-			keyScope.Purpose + hdkeychain.HardenedKeyStart,
-			keyScope.Coin + hdkeychain.HardenedKeyStart,
-			derivationPath.Account,
-			derivationPath.Branch,
-			derivationPath.Index,
-		},
-	}
-
 	return tx, txOut, derivation, value, nil
 }
 
@@ -210,3 +193,34 @@ func (w *Wallet) FetchOutpointInfo(prevOut *wire.OutPoint) (*wire.MsgTx,
 		PkScript: pkScript,
 	}, confs, nil
 }
+
+// FetchDerivationInfo queries for the wallet's knowledge of the passed
+// pkScript and constructs the derivation info and returns it.
+func (w *Wallet) FetchDerivationInfo(pkScript []byte) (*psbt.Bip32Derivation,
+	error) {
+
+	addr, err := w.fetchOutputAddr(pkScript)
+	if err != nil {
+		return nil, err
+	}
+
+	pubKeyAddr, ok := addr.(waddrmgr.ManagedPubKeyAddress)
+	if !ok {
+		return nil, ErrNotMine
+	}
+	keyScope, derivationPath, _ := pubKeyAddr.DerivationInfo()
+
+	derivation := &psbt.Bip32Derivation{
+		PubKey:               pubKeyAddr.PubKey().SerializeCompressed(),
+		MasterKeyFingerprint: derivationPath.MasterKeyFingerprint,
+		Bip32Path: []uint32{
+			keyScope.Purpose + hdkeychain.HardenedKeyStart,
+			keyScope.Coin + hdkeychain.HardenedKeyStart,
+			derivationPath.Account,
+			derivationPath.Branch,
+			derivationPath.Index,
+		},
+	}
+
+	return derivation, nil
+}
diff --git a/wallet/utxos_test.go b/wallet/utxos_test.go
index 85cde530f7..28797f4da1 100644
--- a/wallet/utxos_test.go
+++ b/wallet/utxos_test.go
@@ -128,3 +128,38 @@ func TestFetchOutpointInfo(t *testing.T) {
 	require.Equal(t, utxOut.PkScript, tx.TxOut[prevOut.Index].PkScript)
 	require.Equal(t, int64(0-testBlockHeight), confirmations)
 }
+
+// TestFetchDerivationInfo checks that the wallet can gather the derivation
+// info about an output based on the pkScript.
+func TestFetchDerivationInfo(t *testing.T) {
+	t.Parallel()
+
+	w, cleanup := testWallet(t)
+	defer cleanup()
+
+	// Create an address we can use to send some coins to.
+	addr, err := w.CurrentAddress(0, waddrmgr.KeyScopeBIP0084)
+	require.NoError(t, err)
+	p2shAddr, err := txscript.PayToAddrScript(addr)
+	require.NoError(t, err)
+
+	// Add an output paying to the wallet's address to the database.
+	utxOut := wire.NewTxOut(100000, p2shAddr)
+	incomingTx := &wire.MsgTx{
+		TxIn:  []*wire.TxIn{{}},
+		TxOut: []*wire.TxOut{utxOut},
+	}
+	addUtxo(t, w, incomingTx)
+
+	info, err := w.FetchDerivationInfo(utxOut.PkScript)
+	require.NoError(t, err)
+
+	require.Len(t, info.Bip32Path, 5)
+	require.Equal(t, waddrmgr.KeyScopeBIP0084.Purpose+
+		hdkeychain.HardenedKeyStart, info.Bip32Path[0])
+	require.Equal(t, waddrmgr.KeyScopeBIP0084.Coin+
+		hdkeychain.HardenedKeyStart, info.Bip32Path[1])
+	require.EqualValues(t, hdkeychain.HardenedKeyStart, info.Bip32Path[2])
+	require.Equal(t, uint32(0), info.Bip32Path[3])
+	require.Equal(t, uint32(0), info.Bip32Path[4])
+}