From 967a9ca7de312d9ef1de3cdf164274e2d9bab285 Mon Sep 17 00:00:00 2001
From: Olaoluwa Osuntokun <laolu32@gmail.com>
Date: Sat, 20 Jan 2018 19:57:34 -0800
Subject: [PATCH] lnwallet: add new HtlcSecondLevelRevoke witness type

In this commit, we add a new witness type to the set of known types.
This new type will be used when we need to sweep an HTLC that the
remote party has taken to the second level.
---
 lnwallet/script_utils.go | 50 ++++++++++++++++++++--------------------
 lnwallet/witnessgen.go   |  9 ++++++++
 2 files changed, 34 insertions(+), 25 deletions(-)

diff --git a/lnwallet/script_utils.go b/lnwallet/script_utils.go
index 65622d31..081a5448 100644
--- a/lnwallet/script_utils.go
+++ b/lnwallet/script_utils.go
@@ -817,6 +817,31 @@ func htlcSpendSuccess(signer Signer, signDesc *SignDescriptor,
 	return witnessStack, nil
 }
 
+// htlcSpendRevoke spends a second-level HTLC output. This function is to be
+// used by the sender or receiver of an HTLC to claim the HTLC after a revoked
+// commitment transaction was broadcast.
+func htlcSpendRevoke(signer Signer, signDesc *SignDescriptor,
+	revokeTx *wire.MsgTx) (wire.TxWitness, error) {
+
+	// We don't need any spacial modifications to the transaction as this
+	// is just sweeping a revoked HTLC output. So we'll generate a regular
+	// witness signature.
+	sweepSig, err := signer.SignOutputRaw(revokeTx, signDesc)
+	if err != nil {
+		return nil, err
+	}
+
+	// We set a one as the first element the witness stack (ignoring the
+	// witness script), in order to force execution to the revocation
+	// clause in the second level HTLC script.
+	witnessStack := wire.TxWitness(make([][]byte, 3))
+	witnessStack[0] = append(sweepSig, byte(signDesc.HashType))
+	witnessStack[1] = []byte{1}
+	witnessStack[2] = signDesc.WitnessScript
+
+	return witnessStack, nil
+}
+
 // HtlcSecondLevelSpend exposes the public witness generation function for
 // spending an HTLC success transaction, either due to an expiring time lock or
 // having had the payment preimage. This method is able to spend any
@@ -848,31 +873,6 @@ func HtlcSecondLevelSpend(signer Signer, signDesc *SignDescriptor,
 	return witnessStack, nil
 }
 
-// htlcTimeoutRevoke spends a second-level HTLC output. This function is to be
-// used by the sender or receiver of an HTLC to claim the HTLC after a revoked
-// commitment transaction was broadcast.
-func htlcSpendRevoke(signer Signer, signDesc *SignDescriptor,
-	revokeTx *wire.MsgTx) (wire.TxWitness, error) {
-
-	// We don't need any spacial modifications to the transaction as this
-	// is just sweeping a revoked HTLC output. So we'll generate a regular
-	// witness signature.
-	sweepSig, err := signer.SignOutputRaw(revokeTx, signDesc)
-	if err != nil {
-		return nil, err
-	}
-
-	// We set a one as the first element the witness stack (ignoring the
-	// witness script), in order to force execution to the revocation
-	// clause in the second level HTLC script.
-	witnessStack := wire.TxWitness(make([][]byte, 3))
-	witnessStack[0] = append(sweepSig, byte(signDesc.HashType))
-	witnessStack[1] = []byte{1}
-	witnessStack[2] = signDesc.WitnessScript
-
-	return witnessStack, nil
-}
-
 // lockTimeToSequence converts the passed relative locktime to a sequence
 // number in accordance to BIP-68.
 // See: https://github.com/bitcoin/bips/blob/master/bip-0068.mediawiki
diff --git a/lnwallet/witnessgen.go b/lnwallet/witnessgen.go
index 59a7a534..aca3bacc 100644
--- a/lnwallet/witnessgen.go
+++ b/lnwallet/witnessgen.go
@@ -63,6 +63,12 @@ const (
 	// pre-image to the HTLC. We can sweep this without any additional
 	// timeout.
 	HtlcAcceptedRemoteSuccess WitnessType = 8
+
+	// HtlcSecondLevelRevoke is a witness that allows us to sweep an HTLC
+	// from the remote party's commitment transaction in the case that the
+	// broadcast a revoked commitment, but then also immediately attempt to
+	// go to the second level to claim the HTLC.
+	HtlcSecondLevelRevoke WitnessType = 9
 )
 
 // WitnessGenerator represents a function which is able to generate the final
@@ -111,6 +117,9 @@ func (wt WitnessType) GenWitnessFunc(signer Signer,
 			// value.
 			return receiverHtlcSpendTimeout(signer, desc, tx, -1)
 
+		case HtlcSecondLevelRevoke:
+			return htlcSpendRevoke(signer, desc, tx)
+
 		default:
 			return nil, fmt.Errorf("unknown witness type: %v", wt)
 		}