Merge pull request #1620 from lightninglabs/detect-vpsbt-script-type

tapfreighter+itest: store change script key to DB

Author: guggero

itest/psbt_test.go

@@ -2878,6 +2878,20 @@ func testPsbtExternalCommit(t *harnessTest) {
 	AssertBalanceByID(
 		t.t, bobTapd, targetAssetGenesis.AssetId, assetsToSend,
 	)
+
+	// We make sure the change and the passive are shown properly in the
+	// sender's wallet.
+	AssertBalances(
+		t.t, aliceTapd, targetAsset.Amount-assetsToSend,
+		WithAssetID(targetAssetGenesis.AssetId), WithNumUtxos(1),
+		WithScriptKeyType(asset.ScriptKeyBip86),
+	)
+	passiveAsset := assets[1]
+	AssertBalances(
+		t.t, aliceTapd, passiveAsset.Amount,
+		WithAssetID(passiveAsset.AssetGenesis.AssetId), WithNumUtxos(1),
+		WithScriptKeyType(asset.ScriptKeyBip86),
+	)
 }
 
 // testPsbtLockTimeSend tests that we can send minted assets into a ScriptKey

tapfreighter/fund.go

@@ -121,7 +121,7 @@ func createFundedPacketWithInputs(ctx context.Context, exporter proof.Exporter,
 			}
 		}
 
-		err = deriveChangeOutputKey(ctx, vPkt, keyRing)
+		err = deriveChangeOutputKey(ctx, vPkt, keyRing, addrBook)
 		if err != nil {
 			return nil, fmt.Errorf("unable to derive change "+
 				"output key: %w", err)
@@ -201,7 +201,7 @@ func annotateLocalScriptKeys(ctx context.Context, vPkt *tappsbt.VPacket,
 // back to the local node, assuming there is a change output and it isn't a
 // zero-value tombstone.
 func deriveChangeOutputKey(ctx context.Context, vPkt *tappsbt.VPacket,
-	keyRing KeyRing) error {
+	keyRing KeyRing, addrBook AddrBook) error {
 
 	// If we don't have a split output then there's no change.
 	if !vPkt.HasSplitRootOutput() {
@@ -238,6 +238,16 @@ func deriveChangeOutputKey(ctx context.Context, vPkt *tappsbt.VPacket,
 		changeOut.ScriptKey = asset.NewScriptKeyBip86(
 			changeScriptKey,
 		)
+
+		// To make sure we recognize the script key again if it is
+		// used over the vPSBT RPCs (where we don't encode the script
+		// key type), we now also store the key in the DB.
+		err = addrBook.InsertScriptKey(
+			ctx, changeOut.ScriptKey, changeOut.ScriptKey.Type,
+		)
+		if err != nil {
+			return fmt.Errorf("cannot insert script key: %w", err)
+		}
 	}
 
 	return nil

tapfreighter/fund_test.go

@@ -66,6 +66,18 @@ type mockAddrBook struct {
 	scriptKeys []asset.TweakedScriptKey
 }
 
+func (m *mockAddrBook) InsertScriptKey(_ context.Context,
+	scriptKey asset.ScriptKey, _ asset.ScriptKeyType) error {
+
+	if scriptKey.TweakedScriptKey == nil {
+		return fmt.Errorf("script key must be tweaked")
+	}
+
+	m.scriptKeys = append(m.scriptKeys, *scriptKey.TweakedScriptKey)
+
+	return nil
+}
+
 func (m *mockAddrBook) FetchScriptKey(_ context.Context,
 	tweakedScriptKey *btcec.PublicKey) (*asset.TweakedScriptKey, error) {
 

tapfreighter/wallet.go

@@ -152,6 +152,12 @@ type AddrBook interface {
 	// ErrInternalKeyNotFound is returned.
 	FetchInternalKeyLocator(ctx context.Context,
 		rawKey *btcec.PublicKey) (keychain.KeyLocator, error)
+
+	// InsertScriptKey inserts an address related script key into the
+	// database, so it can be recognized as belonging to the wallet when a
+	// transfer comes in later on.
+	InsertScriptKey(ctx context.Context,
+		scriptKey asset.ScriptKey, keyType asset.ScriptKeyType) error
 }
 
 // AnchorVTxnsParams holds all the parameters needed to create a BTC level