From e1a07b68e84e79148dd1f263c3cf38cd42cbaf32 Mon Sep 17 00:00:00 2001
From: Olaoluwa Osuntokun <laolu32@gmail.com>
Date: Mon, 18 Mar 2019 17:00:43 -0700
Subject: [PATCH] contractcourt: extract which timeout HTLC output to watch
 into new method

---
 .../htlc_outgoing_contest_resolver.go         | 32 +++---------------
 contractcourt/htlc_timeout_resolver.go        | 33 +++++++++++++++++++
 2 files changed, 38 insertions(+), 27 deletions(-)

diff --git a/contractcourt/htlc_outgoing_contest_resolver.go b/contractcourt/htlc_outgoing_contest_resolver.go
index 3aa47b8d..222fd1f3 100644
--- a/contractcourt/htlc_outgoing_contest_resolver.go
+++ b/contractcourt/htlc_outgoing_contest_resolver.go
@@ -40,42 +40,20 @@ func (h *htlcOutgoingContestResolver) Resolve() (ContractResolver, error) {
 
 	// Otherwise, we'll watch for two external signals to decide if we'll
 	// morph into another resolver, or fully resolve the contract.
-
+	//
 	// The output we'll be watching for is the *direct* spend from the HTLC
 	// output. If this isn't our commitment transaction, it'll be right on
 	// the resolution. Otherwise, we fetch this pointer from the input of
 	// the time out transaction.
-	var (
-		outPointToWatch wire.OutPoint
-		scriptToWatch   []byte
-		err             error
-	)
-
-	// TODO(joostjager): output already set properly in
-	// lnwallet.newOutgoingHtlcResolution? And script too?
-	if h.htlcResolution.SignedTimeoutTx == nil {
-		outPointToWatch = h.htlcResolution.ClaimOutpoint
-		scriptToWatch = h.htlcResolution.SweepSignDesc.Output.PkScript
-	} else {
-		// If this is the remote party's commitment, then we'll need to
-		// grab watch the output that our timeout transaction points
-		// to. We can directly grab the outpoint, then also extract the
-		// witness script (the last element of the witness stack) to
-		// re-construct the pkScipt we need to watch.
-		outPointToWatch = h.htlcResolution.SignedTimeoutTx.TxIn[0].PreviousOutPoint
-		witness := h.htlcResolution.SignedTimeoutTx.TxIn[0].Witness
-		scriptToWatch, err = input.WitnessScriptHash(
-			witness[len(witness)-1],
-		)
-		if err != nil {
-			return nil, err
-		}
+	outPointToWatch, scriptToWatch, err := h.chainDetailsToWatch()
+	if err != nil {
+		return nil, err
 	}
 
 	// First, we'll register for a spend notification for this output. If
 	// the remote party sweeps with the pre-image, we'll be notified.
 	spendNtfn, err := h.Notifier.RegisterSpendNtfn(
-		&outPointToWatch, scriptToWatch, h.broadcastHeight,
+		outPointToWatch, scriptToWatch, h.broadcastHeight,
 	)
 	if err != nil {
 		return nil, err
diff --git a/contractcourt/htlc_timeout_resolver.go b/contractcourt/htlc_timeout_resolver.go
index fa51463a..d0fab2fb 100644
--- a/contractcourt/htlc_timeout_resolver.go
+++ b/contractcourt/htlc_timeout_resolver.go
@@ -148,6 +148,39 @@ func (h *htlcTimeoutResolver) claimCleanUp(commitSpend *chainntnfs.SpendDetail)
 	h.resolved = true
 	return nil, h.Checkpoint(h)
 }
+
+// chainDetailsToWatch returns the output and script which we use to watch for
+// spends from the direct HTLC output on the commitment transaction.
+//
+// TODO(joostjager): output already set properly in
+// lnwallet.newOutgoingHtlcResolution? And script too?
+func (h *htlcTimeoutResolver) chainDetailsToWatch() (*wire.OutPoint, []byte, error) {
+	// If there's no timeout transaction, then the claim output is the
+	// output directly on the commitment transaction, so we'll just use
+	// that.
+	if h.htlcResolution.SignedTimeoutTx == nil {
+		outPointToWatch := h.htlcResolution.ClaimOutpoint
+		scriptToWatch := h.htlcResolution.SweepSignDesc.Output.PkScript
+
+		return &outPointToWatch, scriptToWatch, nil
+	}
+
+	// If this is the remote party's commitment, then we'll need to grab
+	// watch the output that our timeout transaction points to. We can
+	// directly grab the outpoint, then also extract the witness script
+	// (the last element of the witness stack) to re-construct the pkScript
+	// we need to watch.
+	outPointToWatch := h.htlcResolution.SignedTimeoutTx.TxIn[0].PreviousOutPoint
+	witness := h.htlcResolution.SignedTimeoutTx.TxIn[0].Witness
+	scriptToWatch, err := input.WitnessScriptHash(
+		witness[len(witness)-1],
+	)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	return &outPointToWatch, scriptToWatch, nil
+}
 // 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