breacharbiter_test: add table-driven breach spend tests

breacharbiter_test.go

@@ -29,6 +29,7 @@ import (
 	"github.com/lightningnetwork/lnd/htlcswitch"
 	"github.com/lightningnetwork/lnd/input"
 	"github.com/lightningnetwork/lnd/keychain"
+	"github.com/lightningnetwork/lnd/lntest"
 	"github.com/lightningnetwork/lnd/lnwallet"
 	"github.com/lightningnetwork/lnd/lnwire"
 	"github.com/lightningnetwork/lnd/shachain"
@@ -1156,10 +1157,162 @@ func TestBreachHandoffFail(t *testing.T) {
 	assertArbiterBreach(t, brar, chanPoint)
 }
 
-// TestBreachSecondLevelTransfer tests that sweep of a HTLC output on a
-// breached commitment is transferred to a second level spend if the output is
-// already spent.
-func TestBreachSecondLevelTransfer(t *testing.T) {
+type publAssertion func(*testing.T, map[wire.OutPoint]*wire.MsgTx,
+	chan *wire.MsgTx)
+
+type breachTest struct {
+	name string
+
+	// spend2ndLevel requests that second level htlcs be spent *again*, as
+	// if by a remote party or watchtower. The outpoint of the second level
+	// htlc is in effect "readded" to the set of inputs.
+	spend2ndLevel bool
+
+	// sendFinalConf informs the test to send a confirmation for the justice
+	// transaction before asserting the arbiter is cleaned up.
+	sendFinalConf bool
+
+	// whenNonZeroInputs is called after spending an input but there are
+	// further inputs to spend in the test.
+	whenNonZeroInputs publAssertion
+
+	// whenZeroInputs is called after spending an input but there are no
+	// further inputs to spend in the test.
+	whenZeroInputs publAssertion
+}
+
+var (
+	// commitSpendTx is used to spend commitment outputs.
+	commitSpendTx = &wire.MsgTx{
+		TxOut: []*wire.TxOut{
+			{Value: 500000000},
+		},
+	}
+	// htlc2ndLevlTx is used to transition an htlc output on the commitment
+	// transaction to a second level htlc.
+	htlc2ndLevlTx = &wire.MsgTx{
+		TxOut: []*wire.TxOut{
+			{Value: 20000},
+		},
+	}
+	// htlcSpendTx is used to spend from a second level htlc.
+	htlcSpendTx = &wire.MsgTx{
+		TxOut: []*wire.TxOut{
+			{Value: 10000},
+		},
+	}
+)
+
+var breachTests = []breachTest{
+	{
+		name:          "all spends",
+		spend2ndLevel: true,
+		whenNonZeroInputs: func(t *testing.T,
+			inputs map[wire.OutPoint]*wire.MsgTx,
+			publTx chan *wire.MsgTx) {
+
+			var tx *wire.MsgTx
+			select {
+			case tx = <-publTx:
+			case <-time.After(5 * time.Second):
+				t.Fatalf("tx was not published")
+			}
+
+			// The justice transaction should have thee same number
+			// of inputs as we are tracking in the test.
+			if len(tx.TxIn) != len(inputs) {
+				t.Fatalf("expected justice txn to have %d "+
+					"inputs, found %d", len(inputs),
+					len(tx.TxIn))
+			}
+
+			// Ensure that each input exists on the justice
+			// transaction.
+			for in := range inputs {
+				findInputIndex(t, in, tx)
+			}
+
+		},
+		whenZeroInputs: func(t *testing.T,
+			inputs map[wire.OutPoint]*wire.MsgTx,
+			publTx chan *wire.MsgTx) {
+
+			// Sanity check to ensure the brar doesn't try to
+			// broadcast another sweep, since all outputs have been
+			// spent externally.
+			select {
+			case <-publTx:
+				t.Fatalf("tx published unexpectedly")
+			case <-time.After(50 * time.Millisecond):
+			}
+		},
+	},
+	{
+		name:          "commit spends, second level sweep",
+		spend2ndLevel: false,
+		sendFinalConf: true,
+		whenNonZeroInputs: func(t *testing.T,
+			inputs map[wire.OutPoint]*wire.MsgTx,
+			publTx chan *wire.MsgTx) {
+
+			select {
+			case <-publTx:
+			case <-time.After(5 * time.Second):
+				t.Fatalf("tx was not published")
+			}
+		},
+		whenZeroInputs: func(t *testing.T,
+			inputs map[wire.OutPoint]*wire.MsgTx,
+			publTx chan *wire.MsgTx) {
+
+			// Now a transaction attempting to spend from the second
+			// level tx should be published instead. Let this
+			// publish succeed by setting the publishing error to
+			// nil.
+			var tx *wire.MsgTx
+			select {
+			case tx = <-publTx:
+			case <-time.After(5 * time.Second):
+				t.Fatalf("tx was not published")
+			}
+
+			// The commitment outputs should be gone, and there
+			// should only be a single htlc spend.
+			if len(tx.TxIn) != 1 {
+				t.Fatalf("expect 1 htlc output, found %d "+
+					"outputs", len(tx.TxIn))
+			}
+
+			// The remaining TxIn  previously attempting to spend
+			// the HTLC outpoint should now be spending from the
+			// second level tx.
+			//
+			// NOTE: Commitment outputs and htlc sweeps are spent
+			// with a different transactions (and thus txids),
+			// ensuring we aren't mistaking this for a different
+			// output type.
+			onlyInput := tx.TxIn[0].PreviousOutPoint.Hash
+			if onlyInput != htlc2ndLevlTx.TxHash() {
+				t.Fatalf("tx not attempting to spend second "+
+					"level tx, %v", tx.TxIn[0])
+			}
+		},
+	},
+}
+
+// TestBreachSpends checks the behavior of the breach arbiter in response to
+// spend events on a channels outputs by asserting that it properly removes or
+// modifies the inputs from the justice txn.
+func TestBreachSpends(t *testing.T) {
+	for _, test := range breachTests {
+		tc := test
+		t.Run(tc.name, func(t *testing.T) {
+			testBreachSpends(t, tc)
+		})
+	}
+}
+
+func testBreachSpends(t *testing.T, test breachTest) {
 	brar, alice, _, bobClose, contractBreaches,
 		cleanUpChans, cleanUpArb := initBreachedState(t)
 	defer cleanUpChans()
@@ -1171,12 +1324,16 @@ func TestBreachSecondLevelTransfer(t *testing.T) {
 		chanPoint    = alice.ChanPoint
 		publTx       = make(chan *wire.MsgTx)
 		publErr      error
+		publMtx      sync.Mutex
 	)
 
 	// Make PublishTransaction always return ErrDoubleSpend to begin with.
 	publErr = lnwallet.ErrDoubleSpend
 	brar.cfg.PublishTransaction = func(tx *wire.MsgTx) error {
 		publTx <- tx
+
+		publMtx.Lock()
+		defer publMtx.Unlock()
 		return publErr
 	}
 
@@ -1204,11 +1361,32 @@ func TestBreachSecondLevelTransfer(t *testing.T) {
 		t.Fatalf("breach arbiter didn't send ack back")
 	}
 
+	state := alice.State()
+	err = state.CloseChannel(&channeldb.ChannelCloseSummary{
+		ChanPoint:               state.FundingOutpoint,
+		ChainHash:               state.ChainHash,
+		RemotePub:               state.IdentityPub,
+		CloseType:               channeldb.BreachClose,
+		Capacity:                state.Capacity,
+		IsPending:               true,
+		ShortChanID:             state.ShortChanID(),
+		RemoteCurrentRevocation: state.RemoteCurrentRevocation,
+		RemoteNextRevocation:    state.RemoteNextRevocation,
+		LocalChanConfig:         state.LocalChanCfg,
+	})
+	if err != nil {
+		t.Fatalf("unable to close channel: %v", err)
+	}
+
 	// After exiting, the breach arbiter should have persisted the
 	// retribution information and the channel should be shown as pending
 	// force closed.
 	assertArbiterBreach(t, brar, chanPoint)
 
+	// Assert that the database sees the channel as pending close, otherwise
+	// the breach arbiter won't be able to fully close it.
+	assertPendingClosed(t, alice)
+
 	// Notify that the breaching transaction is confirmed, to trigger the
 	// retribution logic.
 	notifier := brar.cfg.Notifier.(*mockSpendNotifier)
@@ -1224,47 +1402,93 @@ func TestBreachSecondLevelTransfer(t *testing.T) {
 		t.Fatalf("tx was not published")
 	}
 
-	if tx.TxIn[0].PreviousOutPoint.Hash != forceCloseTx.TxHash() {
-		t.Fatalf("tx not attempting to spend commitment")
+	// All outputs should initially spend from the force closed txn.
+	forceTxID := forceCloseTx.TxHash()
+	for _, txIn := range tx.TxIn {
+		if txIn.PreviousOutPoint.Hash != forceTxID {
+			t.Fatalf("og justice tx not spending commitment")
+		}
 	}
 
-	// Find the index of the TxIn spending the HTLC output.
-	htlcOutpoint := &retribution.HtlcRetributions[0].OutPoint
-	htlcIn := -1
-	for i, txIn := range tx.TxIn {
-		if txIn.PreviousOutPoint == *htlcOutpoint {
-			htlcIn = i
-		}
-	}
-	if htlcIn == -1 {
-		t.Fatalf("htlc in not found")
+	localOutpoint := retribution.LocalOutpoint
+	remoteOutpoint := retribution.RemoteOutpoint
+	htlcOutpoint := retribution.HtlcRetributions[0].OutPoint
+
+	// Construct a map from outpoint on the force close to the transaction
+	// we want it to be spent by. As the test progresses, this map will be
+	// updated to contain only the set of commitment or second level
+	// outpoints that remain to be spent.
+	inputs := map[wire.OutPoint]*wire.MsgTx{
+		htlcOutpoint:   htlc2ndLevlTx,
+		localOutpoint:  commitSpendTx,
+		remoteOutpoint: commitSpendTx,
 	}
 
-	// Since publishing the transaction failed above, the breach arbiter
-	// will attempt another second level check. Now notify that the htlc
-	// output is spent by a second level tx.
-	secondLvlTx := &wire.MsgTx{
-		TxOut: []*wire.TxOut{
-			{Value: 1},
-		},
-	}
-	notifier.Spend(htlcOutpoint, 2, secondLvlTx)
+	// Until no more inputs to spend remain, deliver the spend events and
+	// process the assertions prescribed by the test case.
+	for len(inputs) > 0 {
+		var (
+			op      wire.OutPoint
+			spendTx *wire.MsgTx
+		)
 
-	// Now a transaction attempting to spend from the second level tx
-	// should be published instead. Let this publish succeed by setting the
-	// publishing error to nil.
+		// Pick an outpoint at random from the set of inputs.
+		for op, spendTx = range inputs {
+			delete(inputs, op)
+			break
+		}
+
+		// Deliver the spend notification for the chosen transaction.
+		notifier.Spend(&op, 2, spendTx)
+
+		// When the second layer transfer is detected, add back the
+		// outpoint of the second layer tx so that we can spend it
+		// again. Only do so if the test requests this behavior.
+		spendTxID := spendTx.TxHash()
+		if test.spend2ndLevel && spendTxID == htlc2ndLevlTx.TxHash() {
+			// Create the second level outpoint that will be spent,
+			// the index is always zero for these 1-in-1-out txns.
+			spendOp := wire.OutPoint{Hash: spendTxID}
+			inputs[spendOp] = htlcSpendTx
+		}
+
+		if len(inputs) > 0 {
+			test.whenNonZeroInputs(t, inputs, publTx)
+		} else {
+			// Reset the publishing error so that any publication,
+			// made by the breach arbiter, if any, will succeed.
+			publMtx.Lock()
 			publErr = nil
-	select {
-	case tx = <-publTx:
-	case <-time.After(5 * time.Second):
-		t.Fatalf("tx was not published")
+			publMtx.Unlock()
+			test.whenZeroInputs(t, inputs, publTx)
+		}
 	}
 
-	// The TxIn previously attempting to spend the HTLC outpoint should now
-	// be spending from the second level tx.
-	if tx.TxIn[htlcIn].PreviousOutPoint.Hash != secondLvlTx.TxHash() {
-		t.Fatalf("tx not attempting to spend second level tx, %v", tx.TxIn[0])
+	// Deliver confirmation of sweep if the test expects it.
+	if test.sendFinalConf {
+		notifier.confChannel <- &chainntnfs.TxConfirmation{}
 	}
+
+	// Assert that the channel is fully resolved.
+	assertBrarCleanup(t, brar, alice.ChanPoint, alice.State().Db)
+}
+
+// findInputIndex returns the index of the input that spends from the given
+// outpoint. This method fails if the outpoint is not found.
+func findInputIndex(t *testing.T, op wire.OutPoint, tx *wire.MsgTx) int {
+	t.Helper()
+
+	inputIdx := -1
+	for i, txIn := range tx.TxIn {
+		if txIn.PreviousOutPoint == op {
+			inputIdx = i
+		}
+	}
+	if inputIdx == -1 {
+		t.Fatalf("input %v in not found", op)
+	}
+
+	return inputIdx
 }
 
 // assertArbiterBreach checks that the breach arbiter has persisted the breach
@@ -1272,6 +1496,8 @@ func TestBreachSecondLevelTransfer(t *testing.T) {
 func assertArbiterBreach(t *testing.T, brar *breachArbiter,
 	chanPoint *wire.OutPoint) {
 
+	t.Helper()
+
 	isBreached, err := brar.IsBreached(chanPoint)
 	if err != nil {
 		t.Fatalf("unable to determine if channel is "+
@@ -1290,6 +1516,8 @@ func assertArbiterBreach(t *testing.T, brar *breachArbiter,
 func assertNoArbiterBreach(t *testing.T, brar *breachArbiter,
 	chanPoint *wire.OutPoint) {
 
+	t.Helper()
+
 	isBreached, err := brar.IsBreached(chanPoint)
 	if err != nil {
 		t.Fatalf("unable to determine if channel is "+
@@ -1302,9 +1530,77 @@ func assertNoArbiterBreach(t *testing.T, brar *breachArbiter,
 	}
 }
 
+// assertBrarCleanup blocks until the given channel point has been removed the
+// retribution store and the channel is fully closed in the database.
+func assertBrarCleanup(t *testing.T, brar *breachArbiter,
+	chanPoint *wire.OutPoint, db *channeldb.DB) {
+
+	t.Helper()
+
+	err := lntest.WaitNoError(func() error {
+		isBreached, err := brar.IsBreached(chanPoint)
+		if err != nil {
+			return err
+		}
+
+		if isBreached {
+			return fmt.Errorf("channel %v still breached",
+				chanPoint)
+		}
+
+		closedChans, err := db.FetchClosedChannels(false)
+		if err != nil {
+			return err
+		}
+
+		for _, channel := range closedChans {
+			switch {
+			// Wrong channel.
+			case channel.ChanPoint != *chanPoint:
+				continue
+
+			// Right channel, fully closed!
+			case !channel.IsPending:
+				return nil
+			}
+
+			// Still pending.
+			return fmt.Errorf("channel %v still pending "+
+				"close", chanPoint)
+		}
+
+		return fmt.Errorf("channel %v not closed", chanPoint)
+
+	}, time.Second)
+	if err != nil {
+		t.Fatalf(err.Error())
+	}
+}
+
+// assertPendingClosed checks that the channel has been marked pending closed in
+// the channel database.
+func assertPendingClosed(t *testing.T, c *lnwallet.LightningChannel) {
+	t.Helper()
+
+	closedChans, err := c.State().Db.FetchClosedChannels(true)
+	if err != nil {
+		t.Fatalf("unable to load pending closed channels: %v", err)
+	}
+
+	for _, chanSummary := range closedChans {
+		if chanSummary.ChanPoint == *c.ChanPoint {
+			return
+		}
+	}
+
+	t.Fatalf("channel %v was not marked pending closed", c.ChanPoint)
+}
+
 // assertNotPendingClosed checks that the channel has not been marked pending
 // closed in the channel database.
 func assertNotPendingClosed(t *testing.T, c *lnwallet.LightningChannel) {
+	t.Helper()
+
 	closedChans, err := c.State().Db.FetchClosedChannels(true)
 	if err != nil {
 		t.Fatalf("unable to load pending closed channels: %v", err)