contractcourt: simplify htlcTimeoutResolver, unify with HTLC contest logic

In this commit, we simplify the existing `htlcTImeoutResolver` with some
newly refactored out methods from the `htlcTimeoutContestResolver`. The
resulting logic is easier to follow as it's more linear, and only deals
with spend notifications rather than both spend _and_ confirmation
notifications.

contractcourt/channel_arbitrator_test.go

@@ -511,8 +511,10 @@ func TestChannelArbitratorLocalForceClosePendingHtlc(t *testing.T) {
 
 	// The force close request should trigger broadcast of the commitment
 	// transaction.
-	assertStateTransitions(t, arbLog.newStates,
+	assertStateTransitions(
-		StateBroadcastCommit, StateCommitmentBroadcasted)
+		t, arbLog.newStates, StateBroadcastCommit,
+		StateCommitmentBroadcasted,
+	)
 	select {
 	case <-respChan:
 	case <-time.After(5 * time.Second):
@@ -529,7 +531,17 @@ func TestChannelArbitratorLocalForceClosePendingHtlc(t *testing.T) {
 	}
 
 	// Now notify about the local force close getting confirmed.
-	closeTx := &wire.MsgTx{}
+	closeTx := &wire.MsgTx{
+		TxIn: []*wire.TxIn{
+			{
+				PreviousOutPoint: wire.OutPoint{},
+				Witness: [][]byte{
+					{0x1},
+					{0x2},
+				},
+			},
+		},
+	}
 
 	htlcOp := wire.OutPoint{
 		Hash:  closeTx.TxHash(),
@@ -569,8 +581,10 @@ func TestChannelArbitratorLocalForceClosePendingHtlc(t *testing.T) {
 		&channeldb.ChannelCloseSummary{},
 	}
 
-	assertStateTransitions(t, arbLog.newStates, StateContractClosed,
+	assertStateTransitions(
-		StateWaitingFullResolution)
+		t, arbLog.newStates, StateContractClosed,
+		StateWaitingFullResolution,
+	)
 
 	// htlcOutgoingContestResolver is now active and waiting for the HTLC to
 	// expire. It should not yet have passed it on for incubation.
@@ -592,12 +606,13 @@ func TestChannelArbitratorLocalForceClosePendingHtlc(t *testing.T) {
 		t.Fatalf("no response received")
 	}
 
-	// Notify resolver that output of the commitment has been spent.
+	// Notify resolver that the HTLC output of the commitment has been
-	notifier.confChan <- &chainntnfs.TxConfirmation{}
+	// spent.
+	notifier.spendChan <- &chainntnfs.SpendDetail{SpendingTx: closeTx}
 
 	// As this is our own commitment transaction, the HTLC will go through
-	// to the second level. Channel arbitrator should still not be marked as
+	// to the second level. Channel arbitrator should still not be marked
-	// resolved.
+	// as resolved.
 	select {
 	case <-resolved:
 		t.Fatalf("channel resolved prematurely")
@@ -605,7 +620,7 @@ func TestChannelArbitratorLocalForceClosePendingHtlc(t *testing.T) {
 	}
 
 	// Notify resolver that the second level transaction is spent.
-	notifier.spendChan <- &chainntnfs.SpendDetail{}
+	notifier.spendChan <- &chainntnfs.SpendDetail{SpendingTx: closeTx}
 
 	// At this point channel should be marked as resolved.
 	assertStateTransitions(t, arbLog.newStates, StateFullyResolved)

contractcourt/htlc_timeout_resolver.go

@@ -91,7 +91,9 @@ const (
 // claimCleanUp is a helper method that's called once the HTLC output is spent
 // by the remote party. It'll extract the preimage, add it to the global cache,
 // and finally send the appropriate clean up message.
-func (h *htlcTimeoutResolver) claimCleanUp(commitSpend *chainntnfs.SpendDetail) (ContractResolver, error) {
+func (h *htlcTimeoutResolver) claimCleanUp(
+	commitSpend *chainntnfs.SpendDetail) (ContractResolver, error) {
+
 	// Depending on if this is our commitment or not, then we'll be looking
 	// for a different witness pattern.
 	spenderIndex := commitSpend.SpenderInputIndex
@@ -172,15 +174,51 @@ func (h *htlcTimeoutResolver) chainDetailsToWatch() (*wire.OutPoint, []byte, err
 	// we need to watch.
 	outPointToWatch := h.htlcResolution.SignedTimeoutTx.TxIn[0].PreviousOutPoint
 	witness := h.htlcResolution.SignedTimeoutTx.TxIn[0].Witness
-	scriptToWatch, err := input.WitnessScriptHash(
+	scriptToWatch, err := input.WitnessScriptHash(witness[len(witness)-1])
-		witness[len(witness)-1],
-	)
 	if err != nil {
 		return nil, nil, err
 	}
 
 	return &outPointToWatch, scriptToWatch, nil
 }
+
+// isSuccessSpend returns true if the passed spend on the specified commitment
+// is a success spend that reveals the pre-image or not.
+func isSuccessSpend(spend *chainntnfs.SpendDetail, localCommit bool) bool {
+	// Based on the spending input index and transaction, obtain the
+	// witness that tells us what type of spend this is.
+	spenderIndex := spend.SpenderInputIndex
+	spendingInput := spend.SpendingTx.TxIn[spenderIndex]
+	spendingWitness := spendingInput.Witness
+
+	// If this is the remote commitment then the only possible spends for
+	// outgoing HTLCs are:
+	//
+	//  RECVR: <0> <sender sig> <recvr sig> <preimage> (2nd level success spend)
+	//  REVOK: <sig> <key>
+	//  SENDR: <sig> 0
+	//
+	// In this case, if 5 witness elements are present (factoring the
+	// witness script), and the 3rd element is the size of the pre-image,
+	// then this is a remote spend. If not, then we swept it ourselves, or
+	// revoked their output.
+	if !localCommit {
+		return len(spendingWitness) == expectedRemoteWitnessSuccessSize &&
+			len(spendingWitness[remotePreimageIndex]) == lntypes.HashSize
+	}
+
+	// Otherwise, for our commitment, the only possible spends for an
+	// outgoing HTLC are:
+	//
+	//  SENDR: <0> <sendr sig>  <recvr sig> <0> (2nd level timeout)
+	//  RECVR: <recvr sig>  <preimage>
+	//  REVOK: <revoke sig> <revoke key>
+	//
+	// So the only success case has the pre-image as the 2nd (index 1)
+	// element in the witness.
+	return len(spendingWitness[localPreimageIndex]) == lntypes.HashSize
+}
+
 // Resolve kicks off full resolution of an outgoing HTLC output. If it's our
 // commitment, it isn't resolved until we see the second level HTLC txn
 // confirmed. If it's the remote party's commitment, we don't resolve until we
@@ -243,60 +281,59 @@ func (h *htlcTimeoutResolver) Resolve() (ContractResolver, error) {
 		return nil
 	}
 
-	// With the output sent to the nursery, we'll now wait until the output
+	// Now that we've handed off the HTLC to the nursery, we'll watch for a
-	// has been fully resolved before sending the clean up message.
+	// spend of the output, and make our next move off of that. Depending
-	//
+	// on if this is our commitment, or the remote party's commitment,
-	// TODO(roasbeef): need to be able to cancel nursery?
+	// we'll be watching a different outpoint and script.
-	//  * if they pull on-chain while we're waiting
+	outpointToWatch, scriptToWatch, err := h.chainDetailsToWatch()
-
+	if err != nil {
-	// If we don't have a second layer transaction, then this is a remote
+		return nil, err
-	// party's commitment, so we'll watch for a direct spend.
+	}
-	if h.htlcResolution.SignedTimeoutTx == nil {
+	spendNtfn, err := h.Notifier.RegisterSpendNtfn(
-		// We'll block until: the HTLC output has been spent, and the
+		outpointToWatch, scriptToWatch, h.broadcastHeight,
-		// transaction spending that output is sufficiently confirmed.
+	)
-		log.Infof("%T(%v): waiting for nursery to spend CLTV-locked "+
+	if err != nil {
-			"output", h, h.htlcResolution.ClaimOutpoint)
+		return nil, err
-		if err := waitForOutputResolution(); err != nil {
-			return nil, err
-		}
-	} else {
-		// Otherwise, this is our commitment, so we'll watch for the
-		// second-level transaction to be sufficiently confirmed.
-		secondLevelTXID := h.htlcResolution.SignedTimeoutTx.TxHash()
-		sweepScript := h.htlcResolution.SignedTimeoutTx.TxOut[0].PkScript
-		confNtfn, err := h.Notifier.RegisterConfirmationsNtfn(
-			&secondLevelTXID, sweepScript, 1, h.broadcastHeight,
-		)
-		if err != nil {
-			return nil, err
-		}
-
-		log.Infof("%T(%v): waiting second-level tx (txid=%v) to be "+
-			"fully confirmed", h, h.htlcResolution.ClaimOutpoint,
-			secondLevelTXID)
-
-		select {
-		case _, ok := <-confNtfn.Confirmed:
-			if !ok {
-				return nil, fmt.Errorf("quitting")
-			}
-
-		case <-h.Quit:
-			return nil, fmt.Errorf("quitting")
-		}
 	}
 
-	// TODO(roasbeef): need to watch for remote party sweeping with pre-image?
+	log.Infof("%T(%v): waiting for HTLC output %v to be spent"+
-	//  * have another waiting on spend above, will check the type, if it's
+		"fully confirmed", h, h.htlcResolution.ClaimOutpoint,
-	//    pre-image, then we'll cancel, and send a clean up back with
+		outpointToWatch)
-	//    pre-image, also add to preimage cache
+
+	// We'll block here until either we exit, or the HTLC output on the
+	// commitment transaction has been spent.
+	var (
+		spend *chainntnfs.SpendDetail
+		ok    bool
+	)
+	select {
+	case spend, ok = <-spendNtfn.Spend:
+		if !ok {
+			return nil, fmt.Errorf("quitting")
+		}
+
+	case <-h.Quit:
+		return nil, fmt.Errorf("quitting")
+	}
+
+	// If the spend reveals the pre-image, then we'll enter the clean up
+	// workflow to pass the pre-image back to the incoming link, add it to
+	// the witness cache, and exit.
+	if isSuccessSpend(spend, h.htlcResolution.SignedTimeoutTx != nil) {
+		log.Infof("%T(%v): HTLC has been swept with pre-image by "+
+			"remote party during timeout flow! Adding pre-image to "+
+			"witness cache", h.htlcResolution.ClaimOutpoint)
+
+		return h.claimCleanUp(spend)
+	}
 
 	log.Infof("%T(%v): resolving htlc with incoming fail msg, fully "+
 		"confirmed", h, h.htlcResolution.ClaimOutpoint)
 
 	// At this point, the second-level transaction is sufficiently
 	// confirmed, or a transaction directly spending the output is.
-	// Therefore, we can now send back our clean up message.
+	// Therefore, we can now send back our clean up message, failing the
+	// HTLC on the incoming link.
 	failureMsg := &lnwire.FailPermanentChannelFailure{}
 	if err := h.DeliverResolutionMsg(ResolutionMsg{
 		SourceChan: h.ShortChanID,
@@ -307,7 +344,7 @@ func (h *htlcTimeoutResolver) Resolve() (ContractResolver, error) {
 	}
 
 	// Finally, if this was an output on our commitment transaction, we'll
-	// for the second-level HTLC output to be spent, and for that
+	// wait for the second-level HTLC output to be spent, and for that
 	// transaction itself to confirm.
 	if h.htlcResolution.SignedTimeoutTx != nil {
 		log.Infof("%T(%v): waiting for nursery to spend CSV delayed "+