From 2fbee31c6e490e806189bb8e7f834eca2c0054a7 Mon Sep 17 00:00:00 2001
From: eugene <elzeigel@gmail.com>
Date: Fri, 2 Jul 2021 12:43:35 -0400
Subject: [PATCH] chainntnfs: populate spendsByHeight during historical
 dispatch

This commit fixes a buggy scenario where:
- a spend of the desired outpoint occurs
- RegisterSpend is called, not immediately notifying
- caller performs a historical dispatch, calling UpdateSpendDetails
- caller is notified on the Spend channel
- re-org occurs
- caller is not notified of the re-org
We fix this by correctly populating the spendsByHeight map when
dispatchSpendDetails is called. This mirrors the confirmation case.
---
 chainntnfs/txnotifier.go      | 13 +++++
 chainntnfs/txnotifier_test.go | 92 +++++++++++++++++++++++++++++++++++
 2 files changed, 105 insertions(+)

diff --git a/chainntnfs/txnotifier.go b/chainntnfs/txnotifier.go
index 477d45286..7a328e77c 100644
--- a/chainntnfs/txnotifier.go
+++ b/chainntnfs/txnotifier.go
@@ -1325,6 +1325,19 @@ func (n *TxNotifier) dispatchSpendDetails(ntfn *SpendNtfn, details *SpendDetail)
 		return ErrTxNotifierExiting
 	}
 
+	spendHeight := uint32(details.SpendingHeight)
+
+	// We also add to spendsByHeight to notify on chain reorgs.
+	reorgSafeHeight := spendHeight + n.reorgSafetyLimit
+	if reorgSafeHeight > n.currentHeight {
+		txSet, exists := n.spendsByHeight[spendHeight]
+		if !exists {
+			txSet = make(map[SpendRequest]struct{})
+			n.spendsByHeight[spendHeight] = txSet
+		}
+		txSet[ntfn.SpendRequest] = struct{}{}
+	}
+
 	return nil
 }
 
diff --git a/chainntnfs/txnotifier_test.go b/chainntnfs/txnotifier_test.go
index 40fc052d8..28c6e73a3 100644
--- a/chainntnfs/txnotifier_test.go
+++ b/chainntnfs/txnotifier_test.go
@@ -10,6 +10,7 @@ import (
 	"github.com/btcsuite/btcd/wire"
 	"github.com/btcsuite/btcutil"
 	"github.com/lightningnetwork/lnd/chainntnfs"
+	"github.com/stretchr/testify/require"
 )
 
 var (
@@ -1843,6 +1844,97 @@ func TestTxNotifierSpendReorg(t *testing.T) {
 	}
 }
 
+// TestTxNotifierUpdateSpendReorg tests that a call to RegisterSpend after the
+// spend has been confirmed, and then UpdateSpendDetails (called by historical
+// dispatch), followed by a chain re-org will notify on the Reorg channel. This
+// was not always the case and has since been fixed.
+func TestTxNotifierSpendReorgMissed(t *testing.T) {
+	t.Parallel()
+
+	const startingHeight = 10
+	hintCache := newMockHintCache()
+	n := chainntnfs.NewTxNotifier(
+		startingHeight, chainntnfs.ReorgSafetyLimit, hintCache,
+		hintCache,
+	)
+
+	// We'll create a spending transaction that spends the outpoint we'll
+	// watch.
+	op := wire.OutPoint{Index: 1}
+	spendTx := wire.NewMsgTx(2)
+	spendTx.AddTxIn(&wire.TxIn{
+		PreviousOutPoint: op,
+		SignatureScript:  testSigScript,
+	})
+	spendTxHash := spendTx.TxHash()
+
+	// Create the spend details that we'll call UpdateSpendDetails with.
+	spendDetails := &chainntnfs.SpendDetail{
+		SpentOutPoint:     &op,
+		SpenderTxHash:     &spendTxHash,
+		SpendingTx:        spendTx,
+		SpenderInputIndex: 0,
+		SpendingHeight:    startingHeight + 1,
+	}
+
+	// Now confirm the spending transaction.
+	block := btcutil.NewBlock(&wire.MsgBlock{
+		Transactions: []*wire.MsgTx{spendTx},
+	})
+	err := n.ConnectTip(block.Hash(), startingHeight+1, block.Transactions())
+	if err != nil {
+		t.Fatalf("unable to connect block: %v", err)
+	}
+	if err := n.NotifyHeight(startingHeight + 1); err != nil {
+		t.Fatalf("unable to dispatch notifications: %v", err)
+	}
+
+	// We register for the spend now and will not get a spend notification
+	// until we call UpdateSpendDetails.
+	ntfn, err := n.RegisterSpend(&op, testRawScript, 1)
+	if err != nil {
+		t.Fatalf("unable to register spend: %v", err)
+	}
+
+	// Assert that the HistoricalDispatch variable is non-nil. We'll use
+	// the SpendRequest member to update the spend details.
+	require.NotEmpty(t, ntfn.HistoricalDispatch)
+
+	select {
+	case <-ntfn.Event.Spend:
+		t.Fatalf("did not expect to receive spend ntfn")
+	default:
+	}
+
+	// We now call UpdateSpendDetails with our generated spend details to
+	// simulate a historical spend dispatch being performed. This should
+	// result in a notification being received on the Spend channel.
+	err = n.UpdateSpendDetails(
+		ntfn.HistoricalDispatch.SpendRequest, spendDetails,
+	)
+	require.Empty(t, err)
+
+	// Assert that we receive a Spend notification.
+	select {
+	case <-ntfn.Event.Spend:
+	default:
+		t.Fatalf("expected to receive spend ntfn")
+	}
+
+	// We will now re-org the spending transaction out of the chain, and we
+	// should receive a notification on the Reorg channel.
+	err = n.DisconnectTip(startingHeight + 1)
+	require.Empty(t, err)
+
+	select {
+	case <-ntfn.Event.Spend:
+		t.Fatalf("received unexpected spend ntfn")
+	case <-ntfn.Event.Reorg:
+	default:
+		t.Fatalf("expected spend reorg ntfn")
+	}
+}
+
 // TestTxNotifierConfirmHintCache ensures that the height hints for transactions
 // are kept track of correctly with each new block connected/disconnected. This
 // test also asserts that the height hints are not updated until the simulated