From bca5839929b77920fb5e4489a62cf256c7edd3b5 Mon Sep 17 00:00:00 2001
From: "Johan T. Halseth" <johanth@gmail.com>
Date: Mon, 15 Feb 2021 13:38:41 +0100
Subject: [PATCH 01/12] breacharbiter: extract countRevokedFunds

---
 breacharbiter.go | 59 ++++++++++++++++++++++++++++++++++--------------
 1 file changed, 42 insertions(+), 17 deletions(-)

diff --git a/breacharbiter.go b/breacharbiter.go
index a7ad89a4d..3485b1ea7 100644
--- a/breacharbiter.go
+++ b/breacharbiter.go
@@ -641,22 +641,7 @@ justiceTxBroadcast:
 
 		// Compute both the total value of funds being swept and the
 		// amount of funds that were revoked from the counter party.
-		var totalFunds, revokedFunds btcutil.Amount
-		for _, inp := range breachInfo.breachedOutputs {
-			totalFunds += inp.Amount()
-
-			// If the output being revoked is the remote commitment
-			// output or an offered HTLC output, it's amount
-			// contributes to the value of funds being revoked from
-			// the counter party.
-			switch inp.WitnessType() {
-			case input.CommitmentRevoke:
-				revokedFunds += inp.Amount()
-			case input.HtlcOfferedRevoke:
-				revokedFunds += inp.Amount()
-			default:
-			}
-		}
+		totalFunds, revokedFunds := countRevokedFunds(breachInfo, finalTx)
 
 		brarLog.Infof("Justice for ChannelPoint(%v) has "+
 			"been served, %v revoked funds (%v total) "+
@@ -674,13 +659,53 @@ justiceTxBroadcast:
 
 		// TODO(roasbeef): close other active channels with offending
 		// peer
-
 		return
+
 	case <-b.quit:
 		return
 	}
 }
 
+// countRevokedFunds counts the total and revoked funds swept by our justice
+// TX.
+func countRevokedFunds(breachInfo *retributionInfo,
+	spendTx *wire.MsgTx) (btcutil.Amount, btcutil.Amount) {
+
+	// Compute both the total value of funds being swept and the
+	// amount of funds that were revoked from the counter party.
+	var totalFunds, revokedFunds btcutil.Amount
+	for _, txIn := range spendTx.TxIn {
+		op := txIn.PreviousOutPoint
+
+		// Find the corresponding output in our retribution info.
+		for _, inp := range breachInfo.breachedOutputs {
+			// If the spent outpoint is not among the ouputs that
+			// were breached, we can ignore it.
+			if inp.outpoint != op {
+				continue
+			}
+
+			totalFunds += inp.Amount()
+
+			// If the output being revoked is the remote commitment
+			// output or an offered HTLC output, it's amount
+			// contributes to the value of funds being revoked from
+			// the counter party.
+			switch inp.WitnessType() {
+			case input.CommitmentRevoke:
+				revokedFunds += inp.Amount()
+			case input.HtlcOfferedRevoke:
+				revokedFunds += inp.Amount()
+			default:
+			}
+
+			break
+		}
+	}
+
+	return totalFunds, revokedFunds
+}
+
 // cleanupBreach marks the given channel point as fully resolved and removes the
 // retribution for that the channel from the retribution store.
 func (b *breachArbiter) cleanupBreach(chanPoint *wire.OutPoint) error {

From a6724c1088b2ff445f88c82742330fab8a651c37 Mon Sep 17 00:00:00 2001
From: "Johan T. Halseth" <johanth@gmail.com>
Date: Mon, 15 Feb 2021 13:51:07 +0100
Subject: [PATCH 02/12] breacharbiter: split waitForSpendEvent

We split the method waitForSpendEvent into two, such that we can reuse
it in case the commitment is spent by various transactions.
---
 breacharbiter.go | 126 ++++++++++++++++++++++++++---------------------
 1 file changed, 70 insertions(+), 56 deletions(-)

diff --git a/breacharbiter.go b/breacharbiter.go
index 3485b1ea7..db792aedf 100644
--- a/breacharbiter.go
+++ b/breacharbiter.go
@@ -318,24 +318,22 @@ func convertToSecondLevelRevoke(bo *breachedOutput, breachInfo *retributionInfo,
 		bo.outpoint)
 }
 
+// spend is used to wrap the index of the retributionInfo output that gets
+// spent together with the spend details.
+type spend struct {
+	index  int
+	detail *chainntnfs.SpendDetail
+}
+
 // waitForSpendEvent waits for any of the breached outputs to get spent, and
-// mutates the breachInfo to be able to sweep it. This method should be used
-// when we fail to publish the justice tx because of a double spend, indicating
-// that the counter party has taken one of the breached outputs to the second
-// level. The spendNtfns map is a cache used to store registered spend
-// subscriptions, in case we must call this method multiple times.
+// returns the spend details for those outputs. The spendNtfns map is a cache
+// used to store registered spend subscriptions, in case we must call this
+// method multiple times.
 func (b *breachArbiter) waitForSpendEvent(breachInfo *retributionInfo,
-	spendNtfns map[wire.OutPoint]*chainntnfs.SpendEvent) error {
+	spendNtfns map[wire.OutPoint]*chainntnfs.SpendEvent) ([]spend, error) {
 
 	inputs := breachInfo.breachedOutputs
 
-	// spend is used to wrap the index of the output that gets spent
-	// together with the spend details.
-	type spend struct {
-		index  int
-		detail *chainntnfs.SpendDetail
-	}
-
 	// We create a channel the first goroutine that gets a spend event can
 	// signal. We make it buffered in case multiple spend events come in at
 	// the same time.
@@ -378,7 +376,7 @@ func (b *breachArbiter) waitForSpendEvent(breachInfo *retributionInfo,
 				// to avoid entering an infinite loop.
 				select {
 				case <-b.quit:
-					return errBrarShuttingDown
+					return nil, errBrarShuttingDown
 				default:
 					continue
 				}
@@ -438,62 +436,75 @@ func (b *breachArbiter) waitForSpendEvent(breachInfo *retributionInfo,
 		// channel before ranging over its content.
 		close(allSpends)
 
-		doneOutputs := make(map[int]struct{})
+		// Gather all detected spends and return them.
+		var spends []spend
 		for s := range allSpends {
 			breachedOutput := &inputs[s.index]
 			delete(spendNtfns, breachedOutput.outpoint)
 
-			switch breachedOutput.witnessType {
-			case input.HtlcAcceptedRevoke:
-				fallthrough
-			case input.HtlcOfferedRevoke:
-				brarLog.Infof("Spend on second-level"+
-					"%s(%v) for ChannelPoint(%v) "+
-					"transitions to second-level output",
-					breachedOutput.witnessType,
-					breachedOutput.outpoint,
-					breachInfo.chanPoint)
+			spends = append(spends, s)
+		}
 
-				// In this case we'll morph our initial revoke
-				// spend to instead point to the second level
-				// output, and update the sign descriptor in the
-				// process.
-				convertToSecondLevelRevoke(
-					breachedOutput, breachInfo, s.detail,
-				)
+		return spends, nil
 
-				continue
-			}
+	case <-b.quit:
+		return nil, errBrarShuttingDown
+	}
+}
 
-			brarLog.Infof("Spend on %s(%v) for ChannelPoint(%v) "+
-				"transitions output to terminal state, "+
-				"removing input from justice transaction",
+// updateBreachInfo mutates the passed breachInfo by removing or converting any
+// outputs among the spends.
+func updateBreachInfo(breachInfo *retributionInfo, spends []spend) {
+	inputs := breachInfo.breachedOutputs
+	doneOutputs := make(map[int]struct{})
+
+	for _, s := range spends {
+		breachedOutput := &inputs[s.index]
+
+		switch breachedOutput.witnessType {
+		case input.HtlcAcceptedRevoke:
+			fallthrough
+		case input.HtlcOfferedRevoke:
+			brarLog.Infof("Spend on second-level "+
+				"%s(%v) for ChannelPoint(%v) "+
+				"transitions to second-level output",
 				breachedOutput.witnessType,
 				breachedOutput.outpoint, breachInfo.chanPoint)
 
-			doneOutputs[s.index] = struct{}{}
+			// In this case we'll morph our initial revoke
+			// spend to instead point to the second level
+			// output, and update the sign descriptor in the
+			// process.
+			convertToSecondLevelRevoke(
+				breachedOutput, breachInfo, s.detail,
+			)
+
+			continue
 		}
 
-		// Filter the inputs for which we can no longer proceed.
-		var nextIndex int
-		for i := range inputs {
-			if _, ok := doneOutputs[i]; ok {
-				continue
-			}
+		brarLog.Infof("Spend on %s(%v) for ChannelPoint(%v) "+
+			"transitions output to terminal state, "+
+			"removing input from justice transaction",
+			breachedOutput.witnessType,
+			breachedOutput.outpoint, breachInfo.chanPoint)
 
-			inputs[nextIndex] = inputs[i]
-			nextIndex++
-		}
-
-		// Update our remaining set of outputs before continuing with
-		// another attempt at publication.
-		breachInfo.breachedOutputs = inputs[:nextIndex]
-
-	case <-b.quit:
-		return errBrarShuttingDown
+		doneOutputs[s.index] = struct{}{}
 	}
 
-	return nil
+	// Filter the inputs for which we can no longer proceed.
+	var nextIndex int
+	for i := range inputs {
+		if _, ok := doneOutputs[i]; ok {
+			continue
+		}
+
+		inputs[nextIndex] = inputs[i]
+		nextIndex++
+	}
+
+	// Update our remaining set of outputs before continuing with
+	// another attempt at publication.
+	breachInfo.breachedOutputs = inputs[:nextIndex]
 }
 
 // exactRetribution is a goroutine which is executed once a contract breach has
@@ -587,7 +598,9 @@ justiceTxBroadcast:
 				"attempting to craft new justice tx.")
 			finalTx = nil
 
-			err := b.waitForSpendEvent(breachInfo, spendNtfns)
+			spends, err := b.waitForSpendEvent(
+				breachInfo, spendNtfns,
+			)
 			if err != nil {
 				if err != errBrarShuttingDown {
 					brarLog.Errorf("error waiting for "+
@@ -596,6 +609,7 @@ justiceTxBroadcast:
 				return
 			}
 
+			updateBreachInfo(breachInfo, spends)
 			if len(breachInfo.breachedOutputs) == 0 {
 				brarLog.Debugf("No more outputs to sweep for "+
 					"breach, marking ChannelPoint(%v) "+

From 3aa5e650fb3b2b71c20d27c88bcfbb1393b9d830 Mon Sep 17 00:00:00 2001
From: "Johan T. Halseth" <johanth@gmail.com>
Date: Tue, 20 Apr 2021 08:50:55 +0200
Subject: [PATCH 03/12] lntest/mock: set input index on spend event

---
 lntest/mock/spendnotifier.go | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/lntest/mock/spendnotifier.go b/lntest/mock/spendnotifier.go
index 7d51b458d..ced93f251 100644
--- a/lntest/mock/spendnotifier.go
+++ b/lntest/mock/spendnotifier.go
@@ -67,13 +67,20 @@ func (s *SpendNotifier) Spend(outpoint *wire.OutPoint, height int32,
 	s.mtx.Lock()
 	defer s.mtx.Unlock()
 
+	var inputIndex uint32
+	for i, in := range txn.TxIn {
+		if in.PreviousOutPoint == *outpoint {
+			inputIndex = uint32(i)
+		}
+	}
+
 	txnHash := txn.TxHash()
 	details := &chainntnfs.SpendDetail{
 		SpentOutPoint:     outpoint,
 		SpendingHeight:    height,
 		SpendingTx:        txn,
 		SpenderTxHash:     &txnHash,
-		SpenderInputIndex: outpoint.Index,
+		SpenderInputIndex: inputIndex,
 	}
 
 	// Cache details in case of late registration.

From a192718807ae489f2fab858509e56fbe4b4c6e58 Mon Sep 17 00:00:00 2001
From: "Johan T. Halseth" <johanth@gmail.com>
Date: Tue, 20 Apr 2021 14:54:31 +0200
Subject: [PATCH 04/12] breacharbiter_test: distinguish spending transactions
 from justice tx inputs

Since we want to test more complex combinations of spends of the
breached outputs, we use two maps tracking

1. which transaction will spend the outpoint
2. which outpoints we expect the breacharbiter to include in the justice
   tx

This let us trigger spends of the individual outputs, and depending on
what we want to test check whether the breacharbiter sweeps the expected
outpoints.
---
 breacharbiter_test.go | 78 +++++++++++++++++++++++++++++++------------
 1 file changed, 57 insertions(+), 21 deletions(-)

diff --git a/breacharbiter_test.go b/breacharbiter_test.go
index 0abdab2e2..bfadab85a 100644
--- a/breacharbiter_test.go
+++ b/breacharbiter_test.go
@@ -1220,7 +1220,7 @@ func TestBreachHandoffFail(t *testing.T) {
 	assertArbiterBreach(t, brar, chanPoint)
 }
 
-type publAssertion func(*testing.T, map[wire.OutPoint]*wire.MsgTx,
+type publAssertion func(*testing.T, map[wire.OutPoint]struct{},
 	chan *wire.MsgTx)
 
 type breachTest struct {
@@ -1271,7 +1271,7 @@ var breachTests = []breachTest{
 		name:          "all spends",
 		spend2ndLevel: true,
 		whenNonZeroInputs: func(t *testing.T,
-			inputs map[wire.OutPoint]*wire.MsgTx,
+			inputs map[wire.OutPoint]struct{},
 			publTx chan *wire.MsgTx) {
 
 			var tx *wire.MsgTx
@@ -1281,7 +1281,7 @@ var breachTests = []breachTest{
 				t.Fatalf("tx was not published")
 			}
 
-			// The justice transaction should have thee same number
+			// The justice transaction should have the 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 "+
@@ -1297,7 +1297,7 @@ var breachTests = []breachTest{
 
 		},
 		whenZeroInputs: func(t *testing.T,
-			inputs map[wire.OutPoint]*wire.MsgTx,
+			inputs map[wire.OutPoint]struct{},
 			publTx chan *wire.MsgTx) {
 
 			// Sanity check to ensure the brar doesn't try to
@@ -1315,17 +1315,33 @@ var breachTests = []breachTest{
 		spend2ndLevel: false,
 		sendFinalConf: true,
 		whenNonZeroInputs: func(t *testing.T,
-			inputs map[wire.OutPoint]*wire.MsgTx,
+			inputs map[wire.OutPoint]struct{},
 			publTx chan *wire.MsgTx) {
 
+			var tx *wire.MsgTx
 			select {
-			case <-publTx:
+			case tx = <-publTx:
 			case <-time.After(5 * time.Second):
 				t.Fatalf("tx was not published")
 			}
+
+			// The justice transaction should have the 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,
+			inputs map[wire.OutPoint]struct{},
 			publTx chan *wire.MsgTx) {
 
 			// Now a transaction attempting to spend from the second
@@ -1486,49 +1502,69 @@ func testBreachSpends(t *testing.T, test breachTest) {
 	// 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{
+	spentBy := map[wire.OutPoint]*wire.MsgTx{
 		htlcOutpoint:   htlc2ndLevlTx,
 		localOutpoint:  commitSpendTx,
 		remoteOutpoint: commitSpendTx,
 	}
 
+	// We also keep a map of those remaining outputs we expect the
+	// breacharbiter to try and sweep.
+	inputsToSweep := map[wire.OutPoint]struct{}{
+		htlcOutpoint:   {},
+		localOutpoint:  {},
+		remoteOutpoint: {},
+	}
+
 	// Until no more inputs to spend remain, deliver the spend events and
 	// process the assertions prescribed by the test case.
-	for len(inputs) > 0 {
+	for len(spentBy) > 0 {
 		var (
 			op      wire.OutPoint
 			spendTx *wire.MsgTx
 		)
 
 		// Pick an outpoint at random from the set of inputs.
-		for op, spendTx = range inputs {
-			delete(inputs, op)
+		for op, spendTx = range spentBy {
+			delete(spentBy, 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.
+		// Since the remote just swept this input, we expect our next
+		// justice transaction to not include them.
+		delete(inputsToSweep, op)
+
+		// If this is the second-level spend, we must add the new
+		// outpoint to our expected sweeps.
 		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.
+		if 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
+			inputsToSweep[spendOp] = struct{}{}
+
+			// 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.
+			if test.spend2ndLevel {
+				spentBy[spendOp] = htlcSpendTx
+			}
 		}
 
-		if len(inputs) > 0 {
-			test.whenNonZeroInputs(t, inputs, publTx)
+		if len(spentBy) > 0 {
+			test.whenNonZeroInputs(t, inputsToSweep, publTx)
 		} else {
 			// Reset the publishing error so that any publication,
 			// made by the breach arbiter, if any, will succeed.
 			publMtx.Lock()
 			publErr = nil
 			publMtx.Unlock()
-			test.whenZeroInputs(t, inputs, publTx)
+			test.whenZeroInputs(t, inputsToSweep, publTx)
 		}
 	}
 

From e7ee5ad51f2298b5c4c796bf4ebf12d55376b58c Mon Sep 17 00:00:00 2001
From: "Johan T. Halseth" <johanth@gmail.com>
Date: Tue, 20 Apr 2021 09:24:10 +0200
Subject: [PATCH 05/12] breacharbiter_test: extract sweep tx creation into
 method

---
 breacharbiter_test.go | 83 ++++++++++++++++++++++++++++++++++---------
 1 file changed, 67 insertions(+), 16 deletions(-)

diff --git a/breacharbiter_test.go b/breacharbiter_test.go
index bfadab85a..7072250a1 100644
--- a/breacharbiter_test.go
+++ b/breacharbiter_test.go
@@ -1221,7 +1221,7 @@ func TestBreachHandoffFail(t *testing.T) {
 }
 
 type publAssertion func(*testing.T, map[wire.OutPoint]struct{},
-	chan *wire.MsgTx)
+	chan *wire.MsgTx, chainhash.Hash)
 
 type breachTest struct {
 	name string
@@ -1244,27 +1244,71 @@ type breachTest struct {
 	whenZeroInputs publAssertion
 }
 
-var (
+type spendTxs struct {
+	commitSpendTx    *wire.MsgTx
+	htlc2ndLevlTx    *wire.MsgTx
+	htlc2ndLevlSpend *wire.MsgTx
+}
+
+func getSpendTransactions(_ input.Signer, _ *wire.OutPoint,
+	retribution *lnwallet.BreachRetribution) *spendTxs {
+
+	localOutpoint := retribution.LocalOutpoint
+	remoteOutpoint := retribution.RemoteOutpoint
+	htlcOutpoint := retribution.HtlcRetributions[0].OutPoint
+
 	// commitSpendTx is used to spend commitment outputs.
-	commitSpendTx = &wire.MsgTx{
+	commitSpendTx := &wire.MsgTx{
+		TxIn: []*wire.TxIn{
+			{
+				PreviousOutPoint: localOutpoint,
+			},
+			{
+				PreviousOutPoint: remoteOutpoint,
+			},
+		},
 		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{
+	htlc2ndLevlTx := &wire.MsgTx{
+		TxIn: []*wire.TxIn{
+			{
+				PreviousOutPoint: htlcOutpoint,
+			},
+		},
 		TxOut: []*wire.TxOut{
 			{Value: 20000},
 		},
 	}
+
+	secondLvlOp := wire.OutPoint{
+		Hash:  htlc2ndLevlTx.TxHash(),
+		Index: 0,
+	}
+
 	// htlcSpendTx is used to spend from a second level htlc.
-	htlcSpendTx = &wire.MsgTx{
+	htlcSpendTx := &wire.MsgTx{
+		TxIn: []*wire.TxIn{
+			{
+				PreviousOutPoint: secondLvlOp,
+			},
+		},
+
 		TxOut: []*wire.TxOut{
 			{Value: 10000},
 		},
 	}
-)
+
+	return &spendTxs{
+		commitSpendTx:    commitSpendTx,
+		htlc2ndLevlTx:    htlc2ndLevlTx,
+		htlc2ndLevlSpend: htlcSpendTx,
+	}
+}
 
 var breachTests = []breachTest{
 	{
@@ -1272,7 +1316,7 @@ var breachTests = []breachTest{
 		spend2ndLevel: true,
 		whenNonZeroInputs: func(t *testing.T,
 			inputs map[wire.OutPoint]struct{},
-			publTx chan *wire.MsgTx) {
+			publTx chan *wire.MsgTx, _ chainhash.Hash) {
 
 			var tx *wire.MsgTx
 			select {
@@ -1298,7 +1342,7 @@ var breachTests = []breachTest{
 		},
 		whenZeroInputs: func(t *testing.T,
 			inputs map[wire.OutPoint]struct{},
-			publTx chan *wire.MsgTx) {
+			publTx chan *wire.MsgTx, _ chainhash.Hash) {
 
 			// Sanity check to ensure the brar doesn't try to
 			// broadcast another sweep, since all outputs have been
@@ -1316,7 +1360,7 @@ var breachTests = []breachTest{
 		sendFinalConf: true,
 		whenNonZeroInputs: func(t *testing.T,
 			inputs map[wire.OutPoint]struct{},
-			publTx chan *wire.MsgTx) {
+			publTx chan *wire.MsgTx, _ chainhash.Hash) {
 
 			var tx *wire.MsgTx
 			select {
@@ -1342,7 +1386,8 @@ var breachTests = []breachTest{
 		},
 		whenZeroInputs: func(t *testing.T,
 			inputs map[wire.OutPoint]struct{},
-			publTx chan *wire.MsgTx) {
+			publTx chan *wire.MsgTx,
+			htlc2ndLevlTxHash chainhash.Hash) {
 
 			// Now a transaction attempting to spend from the second
 			// level tx should be published instead. Let this
@@ -1371,7 +1416,7 @@ var breachTests = []breachTest{
 			// ensuring we aren't mistaking this for a different
 			// output type.
 			onlyInput := tx.TxIn[0].PreviousOutPoint.Hash
-			if onlyInput != htlc2ndLevlTx.TxHash() {
+			if onlyInput != htlc2ndLevlTxHash {
 				t.Fatalf("tx not attempting to spend second "+
 					"level tx, %v", tx.TxIn[0])
 			}
@@ -1498,14 +1543,18 @@ func testBreachSpends(t *testing.T, test breachTest) {
 	remoteOutpoint := retribution.RemoteOutpoint
 	htlcOutpoint := retribution.HtlcRetributions[0].OutPoint
 
+	spendTxs := getSpendTransactions(
+		brar.cfg.Signer, chanPoint, retribution,
+	)
+
 	// 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.
 	spentBy := map[wire.OutPoint]*wire.MsgTx{
-		htlcOutpoint:   htlc2ndLevlTx,
-		localOutpoint:  commitSpendTx,
-		remoteOutpoint: commitSpendTx,
+		htlcOutpoint:   spendTxs.htlc2ndLevlTx,
+		localOutpoint:  spendTxs.commitSpendTx,
+		remoteOutpoint: spendTxs.commitSpendTx,
 	}
 
 	// We also keep a map of those remaining outputs we expect the
@@ -1516,6 +1565,8 @@ func testBreachSpends(t *testing.T, test breachTest) {
 		remoteOutpoint: {},
 	}
 
+	htlc2ndLevlTx := spendTxs.htlc2ndLevlTx
+	htlcSpendTx := spendTxs.htlc2ndLevlSpend
 	// Until no more inputs to spend remain, deliver the spend events and
 	// process the assertions prescribed by the test case.
 	for len(spentBy) > 0 {
@@ -1557,14 +1608,14 @@ func testBreachSpends(t *testing.T, test breachTest) {
 		}
 
 		if len(spentBy) > 0 {
-			test.whenNonZeroInputs(t, inputsToSweep, publTx)
+			test.whenNonZeroInputs(t, inputsToSweep, publTx, htlc2ndLevlTx.TxHash())
 		} else {
 			// Reset the publishing error so that any publication,
 			// made by the breach arbiter, if any, will succeed.
 			publMtx.Lock()
 			publErr = nil
 			publMtx.Unlock()
-			test.whenZeroInputs(t, inputsToSweep, publTx)
+			test.whenZeroInputs(t, inputsToSweep, publTx, htlc2ndLevlTx.TxHash())
 		}
 	}
 

From 0a0b5f89c95e6b1382122322316710bc62a6403f Mon Sep 17 00:00:00 2001
From: "Johan T. Halseth" <johanth@gmail.com>
Date: Mon, 26 Apr 2021 15:29:59 +0200
Subject: [PATCH 06/12] input: create IsHtlcSpendRevoke

---
 input/script_utils.go | 60 ++++++++++++++++++++++++++++---------------
 1 file changed, 40 insertions(+), 20 deletions(-)

diff --git a/input/script_utils.go b/input/script_utils.go
index dc3c4ccf9..7ce6fff53 100644
--- a/input/script_utils.go
+++ b/input/script_utils.go
@@ -313,21 +313,32 @@ func SenderHtlcSpendRevokeWithKey(signer Signer, signDesc *SignDescriptor,
 func SenderHtlcSpendRevoke(signer Signer, signDesc *SignDescriptor,
 	sweepTx *wire.MsgTx) (wire.TxWitness, error) {
 
-	if signDesc.KeyDesc.PubKey == nil {
-		return nil, fmt.Errorf("cannot generate witness with nil " +
-			"KeyDesc pubkey")
+	revokeKey, err := deriveRevokePubKey(signDesc)
+	if err != nil {
+		return nil, err
 	}
 
-	// Derive the revocation key using the local revocation base point and
-	// commitment point.
-	revokeKey := DeriveRevocationPubkey(
-		signDesc.KeyDesc.PubKey,
-		signDesc.DoubleTweak.PubKey(),
-	)
-
 	return SenderHtlcSpendRevokeWithKey(signer, signDesc, revokeKey, sweepTx)
 }
 
+// IsHtlcSpendRevoke is used to determine if the passed spend is spending a
+// HTLC output using the revocation key.
+func IsHtlcSpendRevoke(txIn *wire.TxIn, signDesc *SignDescriptor) (
+	bool, error) {
+
+	revokeKey, err := deriveRevokePubKey(signDesc)
+	if err != nil {
+		return false, err
+	}
+
+	if len(txIn.Witness) == 3 &&
+		bytes.Equal(txIn.Witness[1], revokeKey.SerializeCompressed()) {
+		return true, nil
+	}
+
+	return false, nil
+}
+
 // SenderHtlcSpendRedeem constructs a valid witness allowing the receiver of an
 // HTLC to redeem the pending output in the scenario that the sender broadcasts
 // their version of the commitment transaction. A valid spend requires
@@ -575,16 +586,7 @@ func ReceiverHtlcSpendRevokeWithKey(signer Signer, signDesc *SignDescriptor,
 	return witnessStack, nil
 }
 
-// ReceiverHtlcSpendRevoke constructs a valid witness allowing the sender of an
-// HTLC within a previously revoked commitment transaction to re-claim the
-// pending funds in the case that the receiver broadcasts this revoked
-// commitment transaction. This method first derives the appropriate revocation
-// key, and requires that the provided SignDescriptor has a local revocation
-// basepoint and commitment secret in the PubKey and DoubleTweak fields,
-// respectively.
-func ReceiverHtlcSpendRevoke(signer Signer, signDesc *SignDescriptor,
-	sweepTx *wire.MsgTx) (wire.TxWitness, error) {
-
+func deriveRevokePubKey(signDesc *SignDescriptor) (*btcec.PublicKey, error) {
 	if signDesc.KeyDesc.PubKey == nil {
 		return nil, fmt.Errorf("cannot generate witness with nil " +
 			"KeyDesc pubkey")
@@ -597,6 +599,24 @@ func ReceiverHtlcSpendRevoke(signer Signer, signDesc *SignDescriptor,
 		signDesc.DoubleTweak.PubKey(),
 	)
 
+	return revokeKey, nil
+}
+
+// ReceiverHtlcSpendRevoke constructs a valid witness allowing the sender of an
+// HTLC within a previously revoked commitment transaction to re-claim the
+// pending funds in the case that the receiver broadcasts this revoked
+// commitment transaction. This method first derives the appropriate revocation
+// key, and requires that the provided SignDescriptor has a local revocation
+// basepoint and commitment secret in the PubKey and DoubleTweak fields,
+// respectively.
+func ReceiverHtlcSpendRevoke(signer Signer, signDesc *SignDescriptor,
+	sweepTx *wire.MsgTx) (wire.TxWitness, error) {
+
+	revokeKey, err := deriveRevokePubKey(signDesc)
+	if err != nil {
+		return nil, err
+	}
+
 	return ReceiverHtlcSpendRevokeWithKey(signer, signDesc, revokeKey, sweepTx)
 }
 

From c3b2791158de6c188d0fa1f82888d0b60fb15fac Mon Sep 17 00:00:00 2001
From: "Johan T. Halseth" <johanth@gmail.com>
Date: Tue, 20 Apr 2021 08:46:23 +0200
Subject: [PATCH 07/12] breacharbiter: don't transition to second level if own
 spend

---
 breacharbiter.go      |  19 ++++++++
 breacharbiter_test.go | 105 ++++++++++++++++++++++++++++++++++++++++--
 2 files changed, 120 insertions(+), 4 deletions(-)

diff --git a/breacharbiter.go b/breacharbiter.go
index db792aedf..fdd5cd163 100644
--- a/breacharbiter.go
+++ b/breacharbiter.go
@@ -460,11 +460,30 @@ func updateBreachInfo(breachInfo *retributionInfo, spends []spend) {
 
 	for _, s := range spends {
 		breachedOutput := &inputs[s.index]
+		txIn := s.detail.SpendingTx.TxIn[s.detail.SpenderInputIndex]
 
 		switch breachedOutput.witnessType {
 		case input.HtlcAcceptedRevoke:
 			fallthrough
 		case input.HtlcOfferedRevoke:
+			// If the HTLC output was spent using the revocation
+			// key, it is our own spend, and we can forget the
+			// output. Otherwise it has been taken to the second
+			// level.
+			signDesc := &breachedOutput.signDesc
+			ok, err := input.IsHtlcSpendRevoke(txIn, signDesc)
+			if err != nil {
+				brarLog.Errorf("Unable to determine if "+
+					"revoke spend: %v", err)
+				break
+			}
+
+			if ok {
+				brarLog.Debugf("HTLC spend was our own " +
+					"revocation spend")
+				break
+			}
+
 			brarLog.Infof("Spend on second-level "+
 				"%s(%v) for ChannelPoint(%v) "+
 				"transitions to second-level output",
diff --git a/breacharbiter_test.go b/breacharbiter_test.go
index 7072250a1..d7818fe6f 100644
--- a/breacharbiter_test.go
+++ b/breacharbiter_test.go
@@ -36,6 +36,7 @@ import (
 	"github.com/lightningnetwork/lnd/lnwallet/chainfee"
 	"github.com/lightningnetwork/lnd/lnwire"
 	"github.com/lightningnetwork/lnd/shachain"
+	"github.com/stretchr/testify/require"
 )
 
 var (
@@ -1231,6 +1232,10 @@ type breachTest struct {
 	// htlc is in effect "readded" to the set of inputs.
 	spend2ndLevel bool
 
+	// sweepHtlc tests that the HTLC output is swept using the revocation
+	// path in a separate tx.
+	sweepHtlc bool
+
 	// sendFinalConf informs the test to send a confirmation for the justice
 	// transaction before asserting the arbiter is cleaned up.
 	sendFinalConf bool
@@ -1248,10 +1253,11 @@ type spendTxs struct {
 	commitSpendTx    *wire.MsgTx
 	htlc2ndLevlTx    *wire.MsgTx
 	htlc2ndLevlSpend *wire.MsgTx
+	htlcSweep        *wire.MsgTx
 }
 
-func getSpendTransactions(_ input.Signer, _ *wire.OutPoint,
-	retribution *lnwallet.BreachRetribution) *spendTxs {
+func getSpendTransactions(signer input.Signer, chanPoint *wire.OutPoint,
+	retribution *lnwallet.BreachRetribution) (*spendTxs, error) {
 
 	localOutpoint := retribution.LocalOutpoint
 	remoteOutpoint := retribution.RemoteOutpoint
@@ -1303,11 +1309,50 @@ func getSpendTransactions(_ input.Signer, _ *wire.OutPoint,
 		},
 	}
 
+	// htlcSweep is used to spend the HTLC output directly using the
+	// revocation key.
+	htlcSweep := &wire.MsgTx{
+		TxIn: []*wire.TxIn{
+			{
+				PreviousOutPoint: htlcOutpoint,
+			},
+		},
+		TxOut: []*wire.TxOut{
+			{Value: 21000},
+		},
+	}
+
+	// In order for  the breacharbiter to detect that it is being spent
+	// using the revocation key, it will inspect the witness. Therefore
+	// sign and add the witness to the HTLC sweep.
+	retInfo := newRetributionInfo(chanPoint, retribution)
+
+	hashCache := txscript.NewTxSigHashes(htlcSweep)
+	for i := range retInfo.breachedOutputs {
+		inp := &retInfo.breachedOutputs[i]
+
+		// Find the HTLC output. so we can add the witness.
+		switch inp.witnessType {
+		case input.HtlcAcceptedRevoke:
+			fallthrough
+		case input.HtlcOfferedRevoke:
+			inputScript, err := inp.CraftInputScript(
+				signer, htlcSweep, hashCache, 0,
+			)
+			if err != nil {
+				return nil, err
+			}
+
+			htlcSweep.TxIn[0].Witness = inputScript.Witness
+		}
+	}
+
 	return &spendTxs{
 		commitSpendTx:    commitSpendTx,
 		htlc2ndLevlTx:    htlc2ndLevlTx,
 		htlc2ndLevlSpend: htlcSpendTx,
-	}
+		htlcSweep:        htlcSweep,
+	}, nil
 }
 
 var breachTests = []breachTest{
@@ -1422,6 +1467,50 @@ var breachTests = []breachTest{
 			}
 		},
 	},
+	{ // nolint: dupl
+		// Test that if the HTLC output is swept via the revoke path
+		// (by us) in a separate tx, it will be handled correctly.
+		name:      "sweep htlc",
+		sweepHtlc: true,
+		whenNonZeroInputs: func(t *testing.T,
+			inputs map[wire.OutPoint]struct{},
+			publTx chan *wire.MsgTx, _ chainhash.Hash) {
+
+			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 the 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]struct{},
+			publTx chan *wire.MsgTx, _ chainhash.Hash) {
+
+			// 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):
+			}
+		},
+	},
 }
 
 // TestBreachSpends checks the behavior of the breach arbiter in response to
@@ -1543,9 +1632,10 @@ func testBreachSpends(t *testing.T, test breachTest) {
 	remoteOutpoint := retribution.RemoteOutpoint
 	htlcOutpoint := retribution.HtlcRetributions[0].OutPoint
 
-	spendTxs := getSpendTransactions(
+	spendTxs, err := getSpendTransactions(
 		brar.cfg.Signer, chanPoint, retribution,
 	)
+	require.NoError(t, err)
 
 	// 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
@@ -1567,6 +1657,13 @@ func testBreachSpends(t *testing.T, test breachTest) {
 
 	htlc2ndLevlTx := spendTxs.htlc2ndLevlTx
 	htlcSpendTx := spendTxs.htlc2ndLevlSpend
+
+	// If the test is checking sweep of the HTLC directly without the
+	// second level, insert the sweep tx instead.
+	if test.sweepHtlc {
+		spentBy[htlcOutpoint] = spendTxs.htlcSweep
+	}
+
 	// Until no more inputs to spend remain, deliver the spend events and
 	// process the assertions prescribed by the test case.
 	for len(spentBy) > 0 {

From 3be9b74694cc7d78d354c1030000fc68ea06ee01 Mon Sep 17 00:00:00 2001
From: "Johan T. Halseth" <johanth@gmail.com>
Date: Tue, 20 Apr 2021 15:42:23 +0200
Subject: [PATCH 08/12] breacharbiter: replace justice tx conf check with spend
 check

Since we want to potentially broadcast multiple versions of the justice
TX, instead of waiting for confirmation of a specific TXID, we instead
wait for the breached outputs to be spent.
---
 breacharbiter.go      | 133 ++++++++++++++++++++----------------------
 breacharbiter_test.go |  51 +++++++++++-----
 2 files changed, 101 insertions(+), 83 deletions(-)

diff --git a/breacharbiter.go b/breacharbiter.go
index fdd5cd163..78a4b8028 100644
--- a/breacharbiter.go
+++ b/breacharbiter.go
@@ -538,17 +538,14 @@ func (b *breachArbiter) exactRetribution(confChan *chainntnfs.ConfirmationEvent,
 	defer b.wg.Done()
 
 	// TODO(roasbeef): state needs to be checkpointed here
-	var breachConfHeight uint32
 	select {
-	case breachConf, ok := <-confChan.Confirmed:
+	case _, ok := <-confChan.Confirmed:
 		// If the second value is !ok, then the channel has been closed
 		// signifying a daemon shutdown, so we exit.
 		if !ok {
 			return
 		}
 
-		breachConfHeight = breachConf.BlockHeight
-
 		// Otherwise, if this is a real confirmation notification, then
 		// we fall through to complete our duty.
 	case <-b.quit:
@@ -570,6 +567,10 @@ func (b *breachArbiter) exactRetribution(confChan *chainntnfs.ConfirmationEvent,
 		return
 	}
 
+	// Compute both the total value of funds being swept and the
+	// amount of funds that were revoked from the counter party.
+	var totalFunds, revokedFunds btcutil.Amount
+
 	// If this retribution has not been finalized before, we will first
 	// construct a sweep transaction and write it to disk. This will allow
 	// the breach arbiter to re-register for notifications for the justice
@@ -605,30 +606,49 @@ justiceTxBroadcast:
 	err = b.cfg.PublishTransaction(finalTx, label)
 	if err != nil {
 		brarLog.Errorf("Unable to broadcast justice tx: %v", err)
+	}
 
-		if err == lnwallet.ErrDoubleSpend {
-			// Broadcasting the transaction failed because of a
-			// conflict either in the mempool or in chain. We'll
-			// now create spend subscriptions for all HTLC outputs
-			// on the commitment transaction that could possibly
-			// have been spent, and wait for any of them to
-			// trigger.
-			brarLog.Infof("Waiting for a spend event before " +
-				"attempting to craft new justice tx.")
-			finalTx = nil
+	// Regardless of publication succeeded or not, we now wait for any of
+	// the inputs to be spent. If any input got spent by the remote, we
+	// must recreate our justice transaction.
+	var (
+		spendChan = make(chan []spend, 1)
+		errChan   = make(chan error, 1)
+		wg        sync.WaitGroup
+	)
 
-			spends, err := b.waitForSpendEvent(
-				breachInfo, spendNtfns,
-			)
-			if err != nil {
-				if err != errBrarShuttingDown {
-					brarLog.Errorf("error waiting for "+
-						"spend event: %v", err)
-				}
-				return
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+
+		spends, err := b.waitForSpendEvent(breachInfo, spendNtfns)
+		if err != nil {
+			errChan <- err
+			return
+		}
+		spendChan <- spends
+	}()
+
+Loop:
+	for {
+		select {
+		case spends := <-spendChan:
+			// Print the funds swept by the txs.
+			for _, s := range spends {
+				tx := s.detail.SpendingTx
+				t, r := countRevokedFunds(breachInfo, tx)
+				totalFunds += t
+				revokedFunds += r
 			}
 
+			brarLog.Infof("Justice for ChannelPoint(%v) has "+
+				"been served, %v revoked funds (%v total) "+
+				"have been claimed", breachInfo.chanPoint,
+				revokedFunds, totalFunds)
+
+			// Update the breach info with the new spends.
 			updateBreachInfo(breachInfo, spends)
+
 			if len(breachInfo.breachedOutputs) == 0 {
 				brarLog.Debugf("No more outputs to sweep for "+
 					"breach, marking ChannelPoint(%v) "+
@@ -640,63 +660,36 @@ justiceTxBroadcast:
 						"breached ChannelPoint(%v): %v",
 						breachInfo.chanPoint, err)
 				}
-				return
+
+				// TODO(roasbeef): add peer to blacklist?
+
+				// TODO(roasbeef): close other active channels with offending
+				// peer
+				break Loop
 			}
 
+			finalTx = nil
 			brarLog.Infof("Attempting another justice tx "+
 				"with %d inputs",
 				len(breachInfo.breachedOutputs))
 
+			wg.Wait()
 			goto justiceTxBroadcast
+
+		case err := <-errChan:
+			if err != errBrarShuttingDown {
+				brarLog.Errorf("error waiting for "+
+					"spend event: %v", err)
+			}
+			break Loop
+
+		case <-b.quit:
+			break Loop
 		}
 	}
 
-	// As a conclusionary step, we register for a notification to be
-	// dispatched once the justice tx is confirmed. After confirmation we
-	// notify the caller that initiated the retribution workflow that the
-	// deed has been done.
-	justiceTXID := finalTx.TxHash()
-	justiceScript := finalTx.TxOut[0].PkScript
-	confChan, err = b.cfg.Notifier.RegisterConfirmationsNtfn(
-		&justiceTXID, justiceScript, 1, breachConfHeight,
-	)
-	if err != nil {
-		brarLog.Errorf("Unable to register for conf for txid(%v): %v",
-			justiceTXID, err)
-		return
-	}
-
-	select {
-	case _, ok := <-confChan.Confirmed:
-		if !ok {
-			return
-		}
-
-		// Compute both the total value of funds being swept and the
-		// amount of funds that were revoked from the counter party.
-		totalFunds, revokedFunds := countRevokedFunds(breachInfo, finalTx)
-
-		brarLog.Infof("Justice for ChannelPoint(%v) has "+
-			"been served, %v revoked funds (%v total) "+
-			"have been claimed", breachInfo.chanPoint,
-			revokedFunds, totalFunds)
-
-		err = b.cleanupBreach(&breachInfo.chanPoint)
-		if err != nil {
-			brarLog.Errorf("Failed to cleanup breached "+
-				"ChannelPoint(%v): %v", breachInfo.chanPoint,
-				err)
-		}
-
-		// TODO(roasbeef): add peer to blacklist?
-
-		// TODO(roasbeef): close other active channels with offending
-		// peer
-		return
-
-	case <-b.quit:
-		return
-	}
+	// Wait for our go routine to exit.
+	wg.Wait()
 }
 
 // countRevokedFunds counts the total and revoked funds swept by our justice
diff --git a/breacharbiter_test.go b/breacharbiter_test.go
index d7818fe6f..b84775cd2 100644
--- a/breacharbiter_test.go
+++ b/breacharbiter_test.go
@@ -1222,7 +1222,7 @@ func TestBreachHandoffFail(t *testing.T) {
 }
 
 type publAssertion func(*testing.T, map[wire.OutPoint]struct{},
-	chan *wire.MsgTx, chainhash.Hash)
+	chan *wire.MsgTx, chainhash.Hash) *wire.MsgTx
 
 type breachTest struct {
 	name string
@@ -1361,7 +1361,7 @@ var breachTests = []breachTest{
 		spend2ndLevel: true,
 		whenNonZeroInputs: func(t *testing.T,
 			inputs map[wire.OutPoint]struct{},
-			publTx chan *wire.MsgTx, _ chainhash.Hash) {
+			publTx chan *wire.MsgTx, _ chainhash.Hash) *wire.MsgTx {
 
 			var tx *wire.MsgTx
 			select {
@@ -1384,10 +1384,11 @@ var breachTests = []breachTest{
 				findInputIndex(t, in, tx)
 			}
 
+			return tx
 		},
 		whenZeroInputs: func(t *testing.T,
 			inputs map[wire.OutPoint]struct{},
-			publTx chan *wire.MsgTx, _ chainhash.Hash) {
+			publTx chan *wire.MsgTx, _ chainhash.Hash) *wire.MsgTx {
 
 			// Sanity check to ensure the brar doesn't try to
 			// broadcast another sweep, since all outputs have been
@@ -1397,6 +1398,8 @@ var breachTests = []breachTest{
 				t.Fatalf("tx published unexpectedly")
 			case <-time.After(50 * time.Millisecond):
 			}
+
+			return nil
 		},
 	},
 	{
@@ -1405,7 +1408,7 @@ var breachTests = []breachTest{
 		sendFinalConf: true,
 		whenNonZeroInputs: func(t *testing.T,
 			inputs map[wire.OutPoint]struct{},
-			publTx chan *wire.MsgTx, _ chainhash.Hash) {
+			publTx chan *wire.MsgTx, _ chainhash.Hash) *wire.MsgTx {
 
 			var tx *wire.MsgTx
 			select {
@@ -1428,11 +1431,12 @@ var breachTests = []breachTest{
 				findInputIndex(t, in, tx)
 			}
 
+			return tx
 		},
 		whenZeroInputs: func(t *testing.T,
 			inputs map[wire.OutPoint]struct{},
 			publTx chan *wire.MsgTx,
-			htlc2ndLevlTxHash chainhash.Hash) {
+			htlc2ndLevlTxHash chainhash.Hash) *wire.MsgTx {
 
 			// Now a transaction attempting to spend from the second
 			// level tx should be published instead. Let this
@@ -1465,6 +1469,8 @@ var breachTests = []breachTest{
 				t.Fatalf("tx not attempting to spend second "+
 					"level tx, %v", tx.TxIn[0])
 			}
+
+			return tx
 		},
 	},
 	{ // nolint: dupl
@@ -1474,7 +1480,7 @@ var breachTests = []breachTest{
 		sweepHtlc: true,
 		whenNonZeroInputs: func(t *testing.T,
 			inputs map[wire.OutPoint]struct{},
-			publTx chan *wire.MsgTx, _ chainhash.Hash) {
+			publTx chan *wire.MsgTx, _ chainhash.Hash) *wire.MsgTx {
 
 			var tx *wire.MsgTx
 			select {
@@ -1496,10 +1502,12 @@ var breachTests = []breachTest{
 			for in := range inputs {
 				findInputIndex(t, in, tx)
 			}
+
+			return tx
 		},
 		whenZeroInputs: func(t *testing.T,
 			inputs map[wire.OutPoint]struct{},
-			publTx chan *wire.MsgTx, _ chainhash.Hash) {
+			publTx chan *wire.MsgTx, _ chainhash.Hash) *wire.MsgTx {
 
 			// Sanity check to ensure the brar doesn't try to
 			// broadcast another sweep, since all outputs have been
@@ -1509,6 +1517,8 @@ var breachTests = []breachTest{
 				t.Fatalf("tx published unexpectedly")
 			case <-time.After(50 * time.Millisecond):
 			}
+
+			return nil
 		},
 	},
 }
@@ -1567,7 +1577,11 @@ func testBreachSpends(t *testing.T, test breachTest) {
 		},
 		BreachRetribution: retribution,
 	}
-	contractBreaches <- breach
+	select {
+	case contractBreaches <- breach:
+	case <-time.After(15 * time.Second):
+		t.Fatalf("breach not delivered")
+	}
 
 	// We'll also wait to consume the ACK back from the breach arbiter.
 	select {
@@ -1608,7 +1622,12 @@ func testBreachSpends(t *testing.T, test breachTest) {
 	// Notify that the breaching transaction is confirmed, to trigger the
 	// retribution logic.
 	notifier := brar.cfg.Notifier.(*mock.SpendNotifier)
-	notifier.ConfChan <- &chainntnfs.TxConfirmation{}
+
+	select {
+	case notifier.ConfChan <- &chainntnfs.TxConfirmation{}:
+	case <-time.After(15 * time.Second):
+		t.Fatalf("conf not delivered")
+	}
 
 	// The breach arbiter should attempt to sweep all outputs on the
 	// breached commitment. We'll pretend that the HTLC output has been
@@ -1666,6 +1685,7 @@ func testBreachSpends(t *testing.T, test breachTest) {
 
 	// Until no more inputs to spend remain, deliver the spend events and
 	// process the assertions prescribed by the test case.
+	var justiceTx *wire.MsgTx
 	for len(spentBy) > 0 {
 		var (
 			op      wire.OutPoint
@@ -1705,20 +1725,25 @@ func testBreachSpends(t *testing.T, test breachTest) {
 		}
 
 		if len(spentBy) > 0 {
-			test.whenNonZeroInputs(t, inputsToSweep, publTx, htlc2ndLevlTx.TxHash())
+			justiceTx = test.whenNonZeroInputs(t, inputsToSweep, publTx, htlc2ndLevlTx.TxHash())
 		} else {
 			// Reset the publishing error so that any publication,
 			// made by the breach arbiter, if any, will succeed.
 			publMtx.Lock()
 			publErr = nil
 			publMtx.Unlock()
-			test.whenZeroInputs(t, inputsToSweep, publTx, htlc2ndLevlTx.TxHash())
+			justiceTx = test.whenZeroInputs(t, inputsToSweep, publTx, htlc2ndLevlTx.TxHash())
 		}
 	}
 
-	// Deliver confirmation of sweep if the test expects it.
+	// Deliver confirmation of sweep if the test expects it. Since we are
+	// looking for the final justice tx to confirme, we deliver a spend of
+	// all its inputs.
 	if test.sendFinalConf {
-		notifier.ConfChan <- &chainntnfs.TxConfirmation{}
+		for _, txin := range justiceTx.TxIn {
+			op := txin.PreviousOutPoint
+			notifier.Spend(&op, 3, justiceTx)
+		}
 	}
 
 	// Assert that the channel is fully resolved.

From 783d1f9578d1ec298c6017c9ab0fad67d86da6c0 Mon Sep 17 00:00:00 2001
From: "Johan T. Halseth" <johanth@gmail.com>
Date: Wed, 17 Feb 2021 14:22:25 +0100
Subject: [PATCH 09/12] breacharbiter: remove justiceTx finalization

Now that we don't rely on the justice tx TXID anymore, we can remove
finalization of it. Instead we'll recreate the transaction when needed
from the persisted retribution info.
---
 breacharbiter.go      | 108 +++---------------------------------------
 breacharbiter_test.go |  18 -------
 2 files changed, 7 insertions(+), 119 deletions(-)

diff --git a/breacharbiter.go b/breacharbiter.go
index 78a4b8028..090de1f7c 100644
--- a/breacharbiter.go
+++ b/breacharbiter.go
@@ -560,40 +560,18 @@ func (b *breachArbiter) exactRetribution(confChan *chainntnfs.ConfirmationEvent,
 	// SpendEvents between each attempt to not re-register uneccessarily.
 	spendNtfns := make(map[wire.OutPoint]*chainntnfs.SpendEvent)
 
-	finalTx, err := b.cfg.Store.GetFinalizedTxn(&breachInfo.chanPoint)
-	if err != nil {
-		brarLog.Errorf("Unable to get finalized txn for"+
-			"chanid=%v: %v", &breachInfo.chanPoint, err)
-		return
-	}
-
 	// Compute both the total value of funds being swept and the
 	// amount of funds that were revoked from the counter party.
 	var totalFunds, revokedFunds btcutil.Amount
 
-	// If this retribution has not been finalized before, we will first
-	// construct a sweep transaction and write it to disk. This will allow
-	// the breach arbiter to re-register for notifications for the justice
-	// txid.
 justiceTxBroadcast:
-	if finalTx == nil {
-		// With the breach transaction confirmed, we now create the
-		// justice tx which will claim ALL the funds within the
-		// channel.
-		finalTx, err = b.createJusticeTx(breachInfo)
-		if err != nil {
-			brarLog.Errorf("Unable to create justice tx: %v", err)
-			return
-		}
-
-		// Persist our finalized justice transaction before making an
-		// attempt to broadcast.
-		err := b.cfg.Store.Finalize(&breachInfo.chanPoint, finalTx)
-		if err != nil {
-			brarLog.Errorf("Unable to finalize justice tx for "+
-				"chanid=%v: %v", &breachInfo.chanPoint, err)
-			return
-		}
+	// With the breach transaction confirmed, we now create the
+	// justice tx which will claim ALL the funds within the
+	// channel.
+	finalTx, err := b.createJusticeTx(breachInfo)
+	if err != nil {
+		brarLog.Errorf("Unable to create justice tx: %v", err)
+		return
 	}
 
 	brarLog.Debugf("Broadcasting justice tx: %v", newLogClosure(func() string {
@@ -668,7 +646,6 @@ Loop:
 				break Loop
 			}
 
-			finalTx = nil
 			brarLog.Infof("Attempting another justice tx "+
 				"with %d inputs",
 				len(breachInfo.breachedOutputs))
@@ -1265,15 +1242,6 @@ type RetributionStore interface {
 	// is aware of any breaches for the provided channel point.
 	IsBreached(chanPoint *wire.OutPoint) (bool, error)
 
-	// Finalize persists the finalized justice transaction for a particular
-	// channel.
-	Finalize(chanPoint *wire.OutPoint, finalTx *wire.MsgTx) error
-
-	// GetFinalizedTxn loads the finalized justice transaction, if any, from
-	// the retribution store. The finalized transaction will be nil if
-	// Finalize has not yet been called for this channel point.
-	GetFinalizedTxn(chanPoint *wire.OutPoint) (*wire.MsgTx, error)
-
 	// Remove deletes the retributionInfo from disk, if any exists, under
 	// the given key. An error should be re raised if the removal fails.
 	Remove(key *wire.OutPoint) error
@@ -1324,68 +1292,6 @@ func (rs *retributionStore) Add(ret *retributionInfo) error {
 	}, func() {})
 }
 
-// Finalize writes a signed justice transaction to the retribution store. This
-// is done before publishing the transaction, so that we can recover the txid on
-// startup and re-register for confirmation notifications.
-func (rs *retributionStore) Finalize(chanPoint *wire.OutPoint,
-	finalTx *wire.MsgTx) error {
-	return kvdb.Update(rs.db, func(tx kvdb.RwTx) error {
-		justiceBkt, err := tx.CreateTopLevelBucket(justiceTxnBucket)
-		if err != nil {
-			return err
-		}
-
-		var chanBuf bytes.Buffer
-		if err := writeOutpoint(&chanBuf, chanPoint); err != nil {
-			return err
-		}
-
-		var txBuf bytes.Buffer
-		if err := finalTx.Serialize(&txBuf); err != nil {
-			return err
-		}
-
-		return justiceBkt.Put(chanBuf.Bytes(), txBuf.Bytes())
-	}, func() {})
-}
-
-// GetFinalizedTxn loads the finalized justice transaction for the provided
-// channel point. The finalized transaction will be nil if Finalize has yet to
-// be called for this channel point.
-func (rs *retributionStore) GetFinalizedTxn(
-	chanPoint *wire.OutPoint) (*wire.MsgTx, error) {
-
-	var finalTxBytes []byte
-	if err := kvdb.View(rs.db, func(tx kvdb.RTx) error {
-		justiceBkt := tx.ReadBucket(justiceTxnBucket)
-		if justiceBkt == nil {
-			return nil
-		}
-
-		var chanBuf bytes.Buffer
-		if err := writeOutpoint(&chanBuf, chanPoint); err != nil {
-			return err
-		}
-
-		finalTxBytes = justiceBkt.Get(chanBuf.Bytes())
-
-		return nil
-	}, func() {
-		finalTxBytes = nil
-	}); err != nil {
-		return nil, err
-	}
-
-	if finalTxBytes == nil {
-		return nil, nil
-	}
-
-	finalTx := &wire.MsgTx{}
-	err := finalTx.Deserialize(bytes.NewReader(finalTxBytes))
-
-	return finalTx, err
-}
-
 // IsBreached queries the retribution store to discern if this channel was
 // previously breached. This is used when connecting to a peer to determine if
 // it is safe to add a link to the htlcswitch, as we should never add a channel
diff --git a/breacharbiter_test.go b/breacharbiter_test.go
index b84775cd2..3172c0a84 100644
--- a/breacharbiter_test.go
+++ b/breacharbiter_test.go
@@ -408,24 +408,6 @@ func (frs *failingRetributionStore) IsBreached(chanPoint *wire.OutPoint) (bool,
 	return frs.rs.IsBreached(chanPoint)
 }
 
-func (frs *failingRetributionStore) Finalize(chanPoint *wire.OutPoint,
-	finalTx *wire.MsgTx) error {
-
-	frs.mu.Lock()
-	defer frs.mu.Unlock()
-
-	return frs.rs.Finalize(chanPoint, finalTx)
-}
-
-func (frs *failingRetributionStore) GetFinalizedTxn(
-	chanPoint *wire.OutPoint) (*wire.MsgTx, error) {
-
-	frs.mu.Lock()
-	defer frs.mu.Unlock()
-
-	return frs.rs.GetFinalizedTxn(chanPoint)
-}
-
 func (frs *failingRetributionStore) Remove(key *wire.OutPoint) error {
 	frs.mu.Lock()
 	defer frs.mu.Unlock()

From 2d710154c4e1aff9de5fdbc597556d98d9e15a40 Mon Sep 17 00:00:00 2001
From: "Johan T. Halseth" <johanth@gmail.com>
Date: Fri, 12 Feb 2021 12:36:45 +0100
Subject: [PATCH 10/12] breacharbiter: create split variants of justice tx

We define a new struct justiceTxVariants, which holds three different
justice transactions:

1. The "normal" justice tx that spends all breached outputs
2. A tx that spends only the breached to_local output and to_remote output
   (can be nil if none of these exist)
3. A tx that spends all the breached HTLC outputs (can be nil if no HTLC
   outputs exist)

This will later be used to sweep the time sensitive outputs separately,
in case the normal justice tx doesn't confirm in time.
---
 breacharbiter.go      | 98 ++++++++++++++++++++++++++++++++++++++-----
 breacharbiter_test.go | 66 +++++++++++++++++++++++++++++
 2 files changed, 153 insertions(+), 11 deletions(-)

diff --git a/breacharbiter.go b/breacharbiter.go
index 090de1f7c..dfd6452d9 100644
--- a/breacharbiter.go
+++ b/breacharbiter.go
@@ -568,11 +568,12 @@ justiceTxBroadcast:
 	// With the breach transaction confirmed, we now create the
 	// justice tx which will claim ALL the funds within the
 	// channel.
-	finalTx, err := b.createJusticeTx(breachInfo)
+	justiceTxs, err := b.createJusticeTx(breachInfo.breachedOutputs)
 	if err != nil {
 		brarLog.Errorf("Unable to create justice tx: %v", err)
 		return
 	}
+	finalTx := justiceTxs.spendAll
 
 	brarLog.Debugf("Broadcasting justice tx: %v", newLogClosure(func() string {
 		return spew.Sdump(finalTx)
@@ -1073,12 +1074,87 @@ func newRetributionInfo(chanPoint *wire.OutPoint,
 	}
 }
 
-// createJusticeTx creates a transaction which exacts "justice" by sweeping ALL
+// justiceTxVariants is a struct that holds transactions which exacts "justice"
+// by sweeping ALL the funds within the channel which we are now entitled to
+// due to a breach of the channel's contract by the counterparty. There are
+// three variants of the justice transaction:
+//
+// 1. The "normal" justice tx that spends all breached outputs
+// 2. A tx that spends only the breached to_local output and to_remote output
+// (can be nil if none of these exist)
+// 3. A tx that spends all the breached HTLC outputs, and second-level HTLC
+// outputs (can be nil if no HTLC outputs exist).
+//
+// The reason we create these three variants, is that in certain cases (like
+// with the anchor output HTLC malleability), the channel counter party can pin
+// the HTLC outputs with low fee children, hindering our normal justice tx that
+// attempts to spend these outputs from propagating. In this case we want to
+// spend the to_local output separately, before the CSV lock expires.
+type justiceTxVariants struct {
+	spendAll        *wire.MsgTx
+	spendCommitOuts *wire.MsgTx
+	spendHTLCs      *wire.MsgTx
+}
+
+// createJusticeTx creates transactions which exacts "justice" by sweeping ALL
 // the funds within the channel which we are now entitled to due to a breach of
 // the channel's contract by the counterparty. This function returns a *fully*
 // signed transaction with the witness for each input fully in place.
 func (b *breachArbiter) createJusticeTx(
-	r *retributionInfo) (*wire.MsgTx, error) {
+	breachedOutputs []breachedOutput) (*justiceTxVariants, error) {
+
+	var (
+		allInputs    []input.Input
+		commitInputs []input.Input
+		htlcInputs   []input.Input
+	)
+
+	for i := range breachedOutputs {
+		// Grab locally scoped reference to breached output.
+		inp := &breachedOutputs[i]
+		allInputs = append(allInputs, inp)
+
+		// Check if the input is from an HTLC or a commitment output.
+		if inp.WitnessType() == input.HtlcAcceptedRevoke ||
+			inp.WitnessType() == input.HtlcOfferedRevoke ||
+			inp.WitnessType() == input.HtlcSecondLevelRevoke {
+
+			htlcInputs = append(htlcInputs, inp)
+		} else {
+			commitInputs = append(commitInputs, inp)
+		}
+	}
+
+	var (
+		txs = &justiceTxVariants{}
+		err error
+	)
+
+	// For each group of inputs, create a tx that spends them.
+	txs.spendAll, err = b.createSweepTx(allInputs)
+	if err != nil {
+		return nil, err
+	}
+
+	txs.spendCommitOuts, err = b.createSweepTx(commitInputs)
+	if err != nil {
+		return nil, err
+	}
+
+	txs.spendHTLCs, err = b.createSweepTx(htlcInputs)
+	if err != nil {
+		return nil, err
+	}
+
+	return txs, nil
+}
+
+// createSweepTx creates a tx that sweeps the passed inputs back to our wallet.
+func (b *breachArbiter) createSweepTx(inputs []input.Input) (*wire.MsgTx,
+	error) {
+	if len(inputs) == 0 {
+		return nil, nil
+	}
 
 	// We will assemble the breached outputs into a slice of spendable
 	// outputs, while simultaneously computing the estimated weight of the
@@ -1090,7 +1166,7 @@ func (b *breachArbiter) createJusticeTx(
 
 	// Allocate enough space to potentially hold each of the breached
 	// outputs in the retribution info.
-	spendableOutputs = make([]input.Input, 0, len(r.breachedOutputs))
+	spendableOutputs = make([]input.Input, 0, len(inputs))
 
 	// The justice transaction we construct will be a segwit transaction
 	// that pays to a p2wkh output. Components such as the version,
@@ -1099,15 +1175,15 @@ func (b *breachArbiter) createJusticeTx(
 
 	// Next, we iterate over the breached outputs contained in the
 	// retribution info.  For each, we switch over the witness type such
-	// that we contribute the appropriate weight for each input and witness,
-	// finally adding to our list of spendable outputs.
-	for i := range r.breachedOutputs {
+	// that we contribute the appropriate weight for each input and
+	// witness, finally adding to our list of spendable outputs.
+	for i := range inputs {
 		// Grab locally scoped reference to breached output.
-		inp := &r.breachedOutputs[i]
+		inp := inputs[i]
 
-		// First, determine the appropriate estimated witness weight for
-		// the give witness type of this breached output. If the witness
-		// weight cannot be estimated, we will omit it from the
+		// First, determine the appropriate estimated witness weight
+		// for the give witness type of this breached output. If the
+		// witness weight cannot be estimated, we will omit it from the
 		// transaction.
 		witnessWeight, _, err := inp.WitnessType().SizeUpperBound()
 		if err != nil {
diff --git a/breacharbiter_test.go b/breacharbiter_test.go
index 3172c0a84..e8deb82ac 100644
--- a/breacharbiter_test.go
+++ b/breacharbiter_test.go
@@ -1203,6 +1203,72 @@ func TestBreachHandoffFail(t *testing.T) {
 	assertArbiterBreach(t, brar, chanPoint)
 }
 
+// TestBreachCreateJusticeTx tests that we create three different variants of
+// the justice tx.
+func TestBreachCreateJusticeTx(t *testing.T) {
+	brar, _, _, _, _, cleanUpChans, cleanUpArb := initBreachedState(t)
+	defer cleanUpChans()
+	defer cleanUpArb()
+
+	// In this test we just want to check that the correct inputs are added
+	// to the justice tx, not that we create a valid spend, so we just set
+	// some params making the script generation succeed.
+	aliceKeyPriv, _ := btcec.PrivKeyFromBytes(
+		btcec.S256(), channels.AlicesPrivKey,
+	)
+	alicePubKey := aliceKeyPriv.PubKey()
+
+	signDesc := &breachedOutputs[0].signDesc
+	signDesc.KeyDesc.PubKey = alicePubKey
+	signDesc.DoubleTweak = aliceKeyPriv
+
+	// We'll test all the different types of outputs we'll sweep with the
+	// justice tx.
+	outputTypes := []input.StandardWitnessType{
+		input.CommitmentNoDelay,
+		input.CommitSpendNoDelayTweakless,
+		input.CommitmentToRemoteConfirmed,
+		input.CommitmentRevoke,
+		input.HtlcAcceptedRevoke,
+		input.HtlcOfferedRevoke,
+		input.HtlcSecondLevelRevoke,
+	}
+
+	breachedOutputs := make([]breachedOutput, len(outputTypes))
+	for i, wt := range outputTypes {
+		// Create a fake breached output for each type, ensuring they
+		// have different outpoints for our logic to accept them.
+		op := breachedOutputs[0].outpoint
+		op.Index = uint32(i)
+		breachedOutputs[i] = makeBreachedOutput(
+			&op,
+			wt,
+			// Second level scripts doesn't matter in this test.
+			nil,
+			signDesc,
+			1,
+		)
+	}
+
+	// Create the justice transactions.
+	justiceTxs, err := brar.createJusticeTx(breachedOutputs)
+	require.NoError(t, err)
+	require.NotNil(t, justiceTxs)
+
+	// The spendAll tx should be spending all the outputs. This is the
+	// "regular" justice transaction type.
+	require.Len(t, justiceTxs.spendAll.TxIn, len(breachedOutputs))
+
+	// The spendCommitOuts tx should be spending the 4 typed of commit outs
+	// (note that in practice there will be at most two commit outputs per
+	// commmit, but we test all 4 types here).
+	require.Len(t, justiceTxs.spendCommitOuts.TxIn, 4)
+
+	// Finally check that the spendHTLCs tx are spending the two revoked
+	// HTLC types, and the second level type.
+	require.Len(t, justiceTxs.spendHTLCs.TxIn, 3)
+}
+
 type publAssertion func(*testing.T, map[wire.OutPoint]struct{},
 	chan *wire.MsgTx, chainhash.Hash) *wire.MsgTx
 

From db0ec1241299c6213ad6d63cc3f4b9c0551785e5 Mon Sep 17 00:00:00 2001
From: "Johan T. Halseth" <johanth@gmail.com>
Date: Mon, 15 Feb 2021 13:31:08 +0100
Subject: [PATCH 11/12] breacharbiter: broadcast "splitted" justice tx if spend
 all not confirming

In case 4 block passes without our justice tx confirming, we'll "split"
it up, and separately sweep the commitment outs, and HTLC outs.
---
 breacharbiter.go      |  95 +++++++++++++++++++++-
 breacharbiter_test.go | 184 ++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 278 insertions(+), 1 deletion(-)

diff --git a/breacharbiter.go b/breacharbiter.go
index dfd6452d9..0766f9f7d 100644
--- a/breacharbiter.go
+++ b/breacharbiter.go
@@ -25,6 +25,21 @@ import (
 	"github.com/lightningnetwork/lnd/lnwallet/chainfee"
 )
 
+const (
+	// justiceTxConfTarget is the number of blocks we'll use as a
+	// confirmation target when creating the justice transaction. We'll
+	// choose an aggressive target, since we want to be sure it confirms
+	// quickly.
+	justiceTxConfTarget = 2
+
+	// blocksPassedSplitPublish is the number of blocks without
+	// confirmation of the justice tx we'll wait before starting to publish
+	// smaller variants of the justice tx. We do this to mitigate an attack
+	// the channel peer can do by pinning the HTLC outputs of the
+	// commitment with low-fee HTLC transactions.
+	blocksPassedSplitPublish = 4
+)
+
 var (
 	// retributionBucket stores retribution state on disk between detecting
 	// a contract breach, broadcasting a justice transaction that sweeps the
@@ -608,8 +623,20 @@ justiceTxBroadcast:
 		spendChan <- spends
 	}()
 
+	// We'll also register for block notifications, such that in case our
+	// justice tx doesn't confirm within a reasonable timeframe, we can
+	// start to more aggressively sweep the time sensitive outputs.
+	newBlockChan, err := b.cfg.Notifier.RegisterBlockEpochNtfn(nil)
+	if err != nil {
+		brarLog.Errorf("Unable to register for block notifications: %v",
+			err)
+		return
+	}
+	defer newBlockChan.Cancel()
+
 Loop:
 	for {
+
 		select {
 		case spends := <-spendChan:
 			// Print the funds swept by the txs.
@@ -654,6 +681,72 @@ Loop:
 			wg.Wait()
 			goto justiceTxBroadcast
 
+		// On every new block, we check whether we should republish the
+		// transactions.
+		case epoch, ok := <-newBlockChan.Epochs:
+			if !ok {
+				return
+			}
+
+			// If less than four blocks have passed since the
+			// breach confirmed, we'll continue waiting. It was
+			// published with a 2-block fee estimate, so it's not
+			// unexpected that four blocks without confirmation can
+			// pass.
+			splitHeight := breachInfo.breachHeight +
+				blocksPassedSplitPublish
+			if uint32(epoch.Height) < splitHeight {
+				continue Loop
+			}
+
+			brarLog.Warnf("Block height %v arrived without "+
+				"justice tx confirming (breached at "+
+				"height %v), splitting justice tx.",
+				epoch.Height, breachInfo.breachHeight)
+
+			// Otherwise we'll attempt to publish the two separate
+			// justice transactions that sweeps the commitment
+			// outputs and the HTLC outputs separately. This is to
+			// mitigate the case where our "spend all" justice TX
+			// doesn't propagate because the HTLC outputs have been
+			// pinned by low fee HTLC txs.
+			label := labels.MakeLabel(
+				labels.LabelTypeJusticeTransaction, nil,
+			)
+			if justiceTxs.spendCommitOuts != nil {
+				tx := justiceTxs.spendCommitOuts
+
+				brarLog.Debugf("Broadcasting justice tx "+
+					"spending commitment outs: %v",
+					newLogClosure(func() string {
+						return spew.Sdump(tx)
+					}))
+
+				err = b.cfg.PublishTransaction(tx, label)
+				if err != nil {
+					brarLog.Warnf("Unable to broadcast "+
+						"commit out spending justice "+
+						"tx: %v", err)
+				}
+			}
+
+			if justiceTxs.spendHTLCs != nil {
+				tx := justiceTxs.spendHTLCs
+
+				brarLog.Debugf("Broadcasting justice tx "+
+					"spending HTLC outs: %v",
+					newLogClosure(func() string {
+						return spew.Sdump(tx)
+					}))
+
+				err = b.cfg.PublishTransaction(tx, label)
+				if err != nil {
+					brarLog.Warnf("Unable to broadcast "+
+						"HTLC out spending justice "+
+						"tx: %v", err)
+				}
+			}
+
 		case err := <-errChan:
 			if err != errBrarShuttingDown {
 				brarLog.Errorf("error waiting for "+
@@ -1224,7 +1317,7 @@ func (b *breachArbiter) sweepSpendableOutputsTxn(txWeight int64,
 
 	// We'll actually attempt to target inclusion within the next two
 	// blocks as we'd like to sweep these funds back into our wallet ASAP.
-	feePerKw, err := b.cfg.Estimator.EstimateFeePerKW(2)
+	feePerKw, err := b.cfg.Estimator.EstimateFeePerKW(justiceTxConfTarget)
 	if err != nil {
 		return nil, err
 	}
diff --git a/breacharbiter_test.go b/breacharbiter_test.go
index e8deb82ac..358a826a5 100644
--- a/breacharbiter_test.go
+++ b/breacharbiter_test.go
@@ -1798,6 +1798,190 @@ func testBreachSpends(t *testing.T, test breachTest) {
 	assertBrarCleanup(t, brar, alice.ChanPoint, alice.State().Db)
 }
 
+// TestBreachDelayedJusticeConfirmation tests that the breach arbiter will
+// "split" the justice tx in case the first justice tx doesn't confirm within
+// a reasonable time.
+func TestBreachDelayedJusticeConfirmation(t *testing.T) {
+	brar, alice, _, bobClose, contractBreaches,
+		cleanUpChans, cleanUpArb := initBreachedState(t)
+	defer cleanUpChans()
+	defer cleanUpArb()
+
+	var (
+		height       = bobClose.ChanSnapshot.CommitHeight
+		blockHeight  = int32(10)
+		forceCloseTx = bobClose.CloseTx
+		chanPoint    = alice.ChanPoint
+		publTx       = make(chan *wire.MsgTx)
+	)
+
+	// Make PublishTransaction always return succeed.
+	brar.cfg.PublishTransaction = func(tx *wire.MsgTx, _ string) error {
+		publTx <- tx
+		return nil
+	}
+
+	// Notify the breach arbiter about the breach.
+	retribution, err := lnwallet.NewBreachRetribution(
+		alice.State(), height, uint32(blockHeight),
+	)
+	if err != nil {
+		t.Fatalf("unable to create breach retribution: %v", err)
+	}
+
+	processACK := make(chan error, 1)
+	breach := &ContractBreachEvent{
+		ChanPoint: *chanPoint,
+		ProcessACK: func(brarErr error) {
+			processACK <- brarErr
+		},
+		BreachRetribution: retribution,
+	}
+
+	select {
+	case contractBreaches <- breach:
+	case <-time.After(15 * time.Second):
+		t.Fatalf("breach not delivered")
+	}
+
+	// We'll also wait to consume the ACK back from the breach arbiter.
+	select {
+	case err := <-processACK:
+		if err != nil {
+			t.Fatalf("handoff failed: %v", err)
+		}
+	case <-time.After(time.Second * 15):
+		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.(*mock.SpendNotifier)
+
+	select {
+	case notifier.ConfChan <- &chainntnfs.TxConfirmation{}:
+	case <-time.After(15 * time.Second):
+		t.Fatalf("conf not delivered")
+	}
+
+	// The breach arbiter should attempt to sweep all outputs on the
+	// breached commitment.
+	var justiceTx *wire.MsgTx
+	select {
+	case justiceTx = <-publTx:
+	case <-time.After(5 * time.Second):
+		t.Fatalf("tx was not published")
+	}
+
+	require.Len(t, justiceTx.TxIn, 3)
+
+	// All outputs should initially spend from the force closed txn.
+	forceTxID := forceCloseTx.TxHash()
+	for _, txIn := range justiceTx.TxIn {
+		if txIn.PreviousOutPoint.Hash != forceTxID {
+			t.Fatalf("og justice tx not spending commitment")
+		}
+	}
+
+	// Now we'll pretend some blocks pass without the justice tx
+	// confirming.
+	for i := int32(0); i <= 3; i++ {
+		notifier.EpochChan <- &chainntnfs.BlockEpoch{
+			Height: blockHeight + i,
+		}
+
+		// On every epoch, check that no new tx is published.
+		select {
+		case <-publTx:
+			t.Fatalf("tx was published")
+		case <-time.After(20 * time.Millisecond):
+		}
+	}
+
+	// Now mine another block without the justice tx confirming. This
+	// should lead to the breacharbiter publishing the split justice tx
+	// variants.
+	notifier.EpochChan <- &chainntnfs.BlockEpoch{
+		Height: blockHeight + 4,
+	}
+
+	var (
+		splits   []*wire.MsgTx
+		spending = make(map[wire.OutPoint]struct{})
+		maxIndex = uint32(len(forceCloseTx.TxOut)) - 1
+	)
+	for i := 0; i < 2; i++ {
+
+		var tx *wire.MsgTx
+		select {
+		case tx = <-publTx:
+			splits = append(splits, tx)
+
+		case <-time.After(5 * time.Second):
+			t.Fatalf("tx not published")
+		}
+
+		// Check that every input is from the breached tx and that
+		// there are no duplicates.
+		for _, in := range tx.TxIn {
+			op := in.PreviousOutPoint
+			_, ok := spending[op]
+			if ok {
+				t.Fatal("already spent")
+			}
+
+			if op.Hash != forceTxID || op.Index > maxIndex {
+				t.Fatalf("not spending breach")
+			}
+
+			spending[op] = struct{}{}
+		}
+	}
+
+	// All the inputs from the original justice transaction should have
+	// been spent by the 2 splits.
+	require.Len(t, spending, len(justiceTx.TxIn))
+	require.Len(t, splits, 2)
+
+	// Finally notify that they confirm, making the breach arbiter clean
+	// up.
+	for _, tx := range splits {
+		for _, in := range tx.TxIn {
+			op := &in.PreviousOutPoint
+			notifier.Spend(op, blockHeight+5, tx)
+		}
+	}
+
+	// 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 {

From 02268b8912795c7430bdadc93ea45aea0b0fcc0f Mon Sep 17 00:00:00 2001
From: "Johan T. Halseth" <johanth@gmail.com>
Date: Fri, 23 Apr 2021 10:44:49 +0200
Subject: [PATCH 12/12] breacharbiter: fix revoked funds calculation

Since we also must count revoked funds swept from second level revoked
outputs, we move the funds counting into the updateBreachInfo method,
where we already are checking whether the spend is by us or the remote.

We also clean up the logs a bit, to log the incremental sweep of funds
that now can happen.
---
 breacharbiter.go | 99 +++++++++++++++++++-----------------------------
 1 file changed, 39 insertions(+), 60 deletions(-)

diff --git a/breacharbiter.go b/breacharbiter.go
index 0766f9f7d..6aaf00fba 100644
--- a/breacharbiter.go
+++ b/breacharbiter.go
@@ -468,11 +468,15 @@ func (b *breachArbiter) waitForSpendEvent(breachInfo *retributionInfo,
 }
 
 // updateBreachInfo mutates the passed breachInfo by removing or converting any
-// outputs among the spends.
-func updateBreachInfo(breachInfo *retributionInfo, spends []spend) {
+// outputs among the spends. It also counts the total and revoked funds swept
+// by our justice spends.
+func updateBreachInfo(breachInfo *retributionInfo, spends []spend) (
+	btcutil.Amount, btcutil.Amount) {
+
 	inputs := breachInfo.breachedOutputs
 	doneOutputs := make(map[int]struct{})
 
+	var totalFunds, revokedFunds btcutil.Amount
 	for _, s := range spends {
 		breachedOutput := &inputs[s.index]
 		txIn := s.detail.SpendingTx.TxIn[s.detail.SpenderInputIndex]
@@ -516,6 +520,22 @@ func updateBreachInfo(breachInfo *retributionInfo, spends []spend) {
 			continue
 		}
 
+		// Now that we have determined the spend is done by us, we
+		// count the total and revoked funds swept depending on the
+		// input type.
+		switch breachedOutput.witnessType {
+
+		// If the output being revoked is the remote commitment
+		// output or an offered HTLC output, it's amount
+		// contributes to the value of funds being revoked from
+		// the counter party.
+		case input.CommitmentRevoke, input.HtlcSecondLevelRevoke,
+			input.HtlcOfferedRevoke:
+
+			revokedFunds += breachedOutput.Amount()
+		}
+
+		totalFunds += breachedOutput.Amount()
 		brarLog.Infof("Spend on %s(%v) for ChannelPoint(%v) "+
 			"transitions output to terminal state, "+
 			"removing input from justice transaction",
@@ -539,6 +559,7 @@ func updateBreachInfo(breachInfo *retributionInfo, spends []spend) {
 	// Update our remaining set of outputs before continuing with
 	// another attempt at publication.
 	breachInfo.breachedOutputs = inputs[:nextIndex]
+	return totalFunds, revokedFunds
 }
 
 // exactRetribution is a goroutine which is executed once a contract breach has
@@ -639,26 +660,24 @@ Loop:
 
 		select {
 		case spends := <-spendChan:
-			// Print the funds swept by the txs.
-			for _, s := range spends {
-				tx := s.detail.SpendingTx
-				t, r := countRevokedFunds(breachInfo, tx)
-				totalFunds += t
-				revokedFunds += r
-			}
+			// Update the breach info with the new spends.
+			t, r := updateBreachInfo(breachInfo, spends)
+			totalFunds += t
+			revokedFunds += r
 
-			brarLog.Infof("Justice for ChannelPoint(%v) has "+
-				"been served, %v revoked funds (%v total) "+
-				"have been claimed", breachInfo.chanPoint,
+			brarLog.Infof("%v spends from breach tx for "+
+				"ChannelPoint(%v) has been detected, %v "+
+				"revoked funds (%v total) have been claimed",
+				len(spends), breachInfo.chanPoint,
 				revokedFunds, totalFunds)
 
-			// Update the breach info with the new spends.
-			updateBreachInfo(breachInfo, spends)
-
 			if len(breachInfo.breachedOutputs) == 0 {
-				brarLog.Debugf("No more outputs to sweep for "+
-					"breach, marking ChannelPoint(%v) "+
-					"fully resolved", breachInfo.chanPoint)
+				brarLog.Infof("Justice for ChannelPoint(%v) "+
+					"has been served, %v revoked funds "+
+					"(%v total) have been claimed. No "+
+					"more outputs to sweep, marking fully "+
+					"resolved", breachInfo.chanPoint,
+					revokedFunds, totalFunds)
 
 				err = b.cleanupBreach(&breachInfo.chanPoint)
 				if err != nil {
@@ -669,8 +688,8 @@ Loop:
 
 				// TODO(roasbeef): add peer to blacklist?
 
-				// TODO(roasbeef): close other active channels with offending
-				// peer
+				// TODO(roasbeef): close other active channels
+				// with offending peer
 				break Loop
 			}
 
@@ -763,46 +782,6 @@ Loop:
 	wg.Wait()
 }
 
-// countRevokedFunds counts the total and revoked funds swept by our justice
-// TX.
-func countRevokedFunds(breachInfo *retributionInfo,
-	spendTx *wire.MsgTx) (btcutil.Amount, btcutil.Amount) {
-
-	// Compute both the total value of funds being swept and the
-	// amount of funds that were revoked from the counter party.
-	var totalFunds, revokedFunds btcutil.Amount
-	for _, txIn := range spendTx.TxIn {
-		op := txIn.PreviousOutPoint
-
-		// Find the corresponding output in our retribution info.
-		for _, inp := range breachInfo.breachedOutputs {
-			// If the spent outpoint is not among the ouputs that
-			// were breached, we can ignore it.
-			if inp.outpoint != op {
-				continue
-			}
-
-			totalFunds += inp.Amount()
-
-			// If the output being revoked is the remote commitment
-			// output or an offered HTLC output, it's amount
-			// contributes to the value of funds being revoked from
-			// the counter party.
-			switch inp.WitnessType() {
-			case input.CommitmentRevoke:
-				revokedFunds += inp.Amount()
-			case input.HtlcOfferedRevoke:
-				revokedFunds += inp.Amount()
-			default:
-			}
-
-			break
-		}
-	}
-
-	return totalFunds, revokedFunds
-}
-
 // cleanupBreach marks the given channel point as fully resolved and removes the
 // retribution for that the channel from the retribution store.
 func (b *breachArbiter) cleanupBreach(chanPoint *wire.OutPoint) error {