From 1e7c5415ae231e570cbf97625996ddaac04b2524 Mon Sep 17 00:00:00 2001
From: Olaoluwa Osuntokun <laolu32@gmail.com>
Date: Tue, 15 Oct 2024 19:17:06 -0700
Subject: [PATCH] input: add new Preimage method to input.Input

In this commit, we add a new method to obtain an option of a preimage to
the input.Input struct. This is useful for callers that have an Input,
and want to optionally obtain the preimage.
---
 contractcourt/breach_arbitrator.go |  5 +++++
 input/input.go                     | 31 +++++++++++++++++++++++++++++-
 input/mocks.go                     | 11 +++++++++++
 3 files changed, 46 insertions(+), 1 deletion(-)

diff --git a/contractcourt/breach_arbitrator.go b/contractcourt/breach_arbitrator.go
index dc690e85c..82b997746 100644
--- a/contractcourt/breach_arbitrator.go
+++ b/contractcourt/breach_arbitrator.go
@@ -1160,6 +1160,11 @@ func (bo *breachedOutput) SignDesc() *input.SignDescriptor {
 	return &bo.signDesc
 }
 
+// Preimage returns the preimage that was used to create the breached output.
+func (bo *breachedOutput) Preimage() fn.Option[lntypes.Preimage] {
+	return fn.None[lntypes.Preimage]()
+}
+
 // CraftInputScript computes a valid witness that allows us to spend from the
 // breached output. It does so by first generating and memoizing the witness
 // generation function, which parameterized primarily by the witness type and
diff --git a/input/input.go b/input/input.go
index 6693e9fa8..6835f82c8 100644
--- a/input/input.go
+++ b/input/input.go
@@ -69,6 +69,9 @@ type Input interface {
 	// ResolutionBlob returns a special opaque blob to be used to
 	// sweep/resolve this input.
 	ResolutionBlob() fn.Option[tlv.Blob]
+
+	// Preimage returns the preimage for the input if it is an HTLC input.
+	Preimage() fn.Option[lntypes.Preimage]
 }
 
 // TxInfo describes properties of a parent tx that are relevant for CPFP.
@@ -285,6 +288,11 @@ func (bi *BaseInput) CraftInputScript(signer Signer, txn *wire.MsgTx,
 	return witnessFunc(txn, hashCache, txinIdx)
 }
 
+// Preimage returns the preimage for the input if it is an HTLC input.
+func (bi *BaseInput) Preimage() fn.Option[lntypes.Preimage] {
+	return fn.None[lntypes.Preimage]()
+}
+
 // HtlcSucceedInput constitutes a sweep input that needs a pre-image. The input
 // is expected to reside on the commitment tx of the remote party and should
 // not be a second level tx output.
@@ -357,7 +365,6 @@ func (h *HtlcSucceedInput) CraftInputScript(signer Signer, txn *wire.MsgTx,
 		}
 
 		desc.SignMethod = TaprootScriptSpendSignMethod
-
 		witness, err = SenderHTLCScriptTaprootRedeem(
 			signer, &desc, txn, h.preimage, nil, nil,
 		)
@@ -375,6 +382,15 @@ func (h *HtlcSucceedInput) CraftInputScript(signer Signer, txn *wire.MsgTx,
 	}, nil
 }
 
+// Preimage returns the preimage for the input if it is an HTLC input.
+func (h *HtlcSucceedInput) Preimage() fn.Option[lntypes.Preimage] {
+	if len(h.preimage) == 0 {
+		return fn.None[lntypes.Preimage]()
+	}
+
+	return fn.Some(lntypes.Preimage(h.preimage))
+}
+
 // HtlcSecondLevelAnchorInput is an input type used to spend HTLC outputs
 // using a re-signed second level transaction, either via the timeout or success
 // paths.
@@ -391,6 +407,8 @@ type HtlcSecondLevelAnchorInput struct {
 		hashCache *txscript.TxSigHashes,
 		prevOutputFetcher txscript.PrevOutputFetcher,
 		txinIdx int) (wire.TxWitness, error)
+
+	preimage []byte
 }
 
 // RequiredTxOut returns the tx out needed to be present on the sweep tx for
@@ -427,6 +445,15 @@ func (i *HtlcSecondLevelAnchorInput) CraftInputScript(signer Signer,
 	}, nil
 }
 
+// Preimage returns the preimage for the input if it is an HTLC input.
+func (i *HtlcSecondLevelAnchorInput) Preimage() fn.Option[lntypes.Preimage] {
+	if len(i.preimage) == 0 {
+		return fn.None[lntypes.Preimage]()
+	}
+
+	return fn.Some(lntypes.Preimage(i.preimage))
+}
+
 // MakeHtlcSecondLevelTimeoutAnchorInput creates an input allowing the sweeper
 // to spend the HTLC output on our commit using the second level timeout
 // transaction.
@@ -545,6 +572,7 @@ func MakeHtlcSecondLevelSuccessAnchorInput(signedTx *wire.MsgTx,
 		SignedTx:      signedTx,
 		inputKit:      input.inputKit,
 		createWitness: createWitness,
+		preimage:      preimage[:],
 	}
 }
 
@@ -588,6 +616,7 @@ func MakeHtlcSecondLevelSuccessTaprootInput(signedTx *wire.MsgTx,
 		inputKit:      input.inputKit,
 		SignedTx:      signedTx,
 		createWitness: createWitness,
+		preimage:      preimage[:],
 	}
 }
 
diff --git a/input/mocks.go b/input/mocks.go
index 695525955..c2af637ed 100644
--- a/input/mocks.go
+++ b/input/mocks.go
@@ -140,6 +140,17 @@ func (m *MockInput) ResolutionBlob() fn.Option[tlv.Blob] {
 	return info.(fn.Option[tlv.Blob])
 }
 
+func (m *MockInput) Preimage() fn.Option[lntypes.Preimage] {
+	args := m.Called()
+
+	info := args.Get(0)
+	if info == nil {
+		return fn.None[lntypes.Preimage]()
+	}
+
+	return info.(fn.Option[lntypes.Preimage])
+}
+
 // MockWitnessType implements the `WitnessType` interface and is used by other
 // packages for mock testing.
 type MockWitnessType struct {