contractcourt: make sure isPreimageSpend won't panic

This commit adds a length check in `isPreimageSpend` to make sure it
won't ever panic and adds a unit test to verify it.

contractcourt/htlc_timeout_resolver.go

@@ -113,6 +113,11 @@ const (
 	// swept on-chain by them with pre-image.
 	expectedRemoteWitnessSuccessSize = 5
 
+	// expectedLocalWitnessSuccessSize is the expected size of the witness
+	// on the local commitment transaction for an outgoing HTLC that is
+	// swept on-chain by them with pre-image.
+	expectedLocalWitnessSuccessSize = 3
+
 	// remotePreimageIndex index within the witness on the remote
 	// commitment transaction that will hold they pre-image if they go to
 	// sweep it on chain.
@@ -130,6 +135,12 @@ const (
 	//     <control_block>
 	remoteTaprootWitnessSuccessSize = 5
 
+	// localTaprootWitnessSuccessSize is the expected size of the witness
+	// on the local commitment for taproot channels. The spend path will
+	// look like
+	//  - <receiver sig> <preimage> <success_script> <control_block>
+	localTaprootWitnessSuccessSize = 4
+
 	// taprootRemotePreimageIndex is the index within the witness on the
 	// taproot remote commitment spend that'll hold the pre-image if the
 	// remote party sweeps it.
@@ -329,10 +340,10 @@ func isPreimageSpend(isTaproot bool, spend *chainntnfs.SpendDetail,
 	//   - <sender sig> <receiver sig> <preimage> <success_script>
 	//     <control_block>
 	case isTaproot && !localCommit:
-		preImageIdx := taprootRemotePreimageIndex
+		return checkSizeAndIndex(
-		//nolint:lll
+			spendingWitness, remoteTaprootWitnessSuccessSize,
-		return len(spendingWitness) == remoteTaprootWitnessSuccessSize &&
+			taprootRemotePreimageIndex,
-			len(spendingWitness[preImageIdx]) == lntypes.HashSize
+		)
 
 	// Otherwise, then if this is our local commitment transaction, then if
 	// they're sweeping the transaction, it'll be directly from the output,
@@ -343,8 +354,10 @@ func isPreimageSpend(isTaproot bool, spend *chainntnfs.SpendDetail,
 	//
 	//  - <receiver sig> <preimage> <success_script> <control_block>
 	case isTaproot && localCommit:
-		return len(spendingWitness[localPreimageIndex]) ==
+		return checkSizeAndIndex(
-			lntypes.HashSize
+			spendingWitness, localTaprootWitnessSuccessSize,
+			localPreimageIndex,
+		)
 
 	// If this is the non-taproot, remote commitment then the only possible
 	// spends for outgoing HTLCs are:
@@ -358,9 +371,10 @@ func isPreimageSpend(isTaproot bool, spend *chainntnfs.SpendDetail,
 	// then this is a remote spend. If not, then we swept it ourselves, or
 	// revoked their output.
 	case !isTaproot && !localCommit:
-		return len(spendingWitness) == expectedRemoteWitnessSuccessSize &&
+		return checkSizeAndIndex(
-			len(spendingWitness[remotePreimageIndex]) ==
+			spendingWitness, expectedRemoteWitnessSuccessSize,
-				lntypes.HashSize
+			remotePreimageIndex,
+		)
 
 	// Otherwise, for our non-taproot commitment, the only possible spends
 	// for an outgoing HTLC are:
@@ -373,12 +387,25 @@ func isPreimageSpend(isTaproot bool, spend *chainntnfs.SpendDetail,
 	// element in the witness.
 	case !isTaproot:
 		fallthrough
+
 	default:
-		return len(spendingWitness[localPreimageIndex]) ==
+		return checkSizeAndIndex(
-			lntypes.HashSize
+			spendingWitness, expectedLocalWitnessSuccessSize,
+			localPreimageIndex,
+		)
 	}
 }
 
+// checkSizeAndIndex checks that the witness is of the expected size and that
+// the witness element at the specified index is of the expected size.
+func checkSizeAndIndex(witness wire.TxWitness, size, index int) bool {
+	if len(witness) != size {
+		return false
+	}
+
+	return len(witness[index]) == 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

contractcourt/htlc_timeout_resolver_test.go

@@ -24,6 +24,11 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
+var (
+	dummyBytes    = []byte{0}
+	preimageBytes = bytes.Repeat([]byte{1}, lntypes.HashSize)
+)
+
 type mockWitnessBeacon struct {
 	preImageUpdates chan lntypes.Preimage
 	newPreimages    chan []lntypes.Preimage
@@ -1324,3 +1329,146 @@ func testHtlcTimeout(t *testing.T, resolution lnwallet.OutgoingHtlcResolution,
 		_ = runFromCheckpoint(t, ctx, checkpoints[i+1:])
 	}
 }
+
+// TestCheckSizeAndIndex checks that the `checkSizeAndIndex` behaves as
+// expected.
+func TestCheckSizeAndIndex(t *testing.T) {
+	t.Parallel()
+
+	testCases := []struct {
+		name     string
+		witness  wire.TxWitness
+		size     int
+		index    int
+		expected bool
+	}{
+		{
+			// Test that a witness with the correct size and index
+			// for the preimage.
+			name: "valid preimage",
+			witness: wire.TxWitness{
+				dummyBytes, preimageBytes,
+			},
+			size:     2,
+			index:    1,
+			expected: true,
+		},
+		{
+			// Test that a witness with the wrong size.
+			name: "wrong witness size",
+			witness: wire.TxWitness{
+				dummyBytes, preimageBytes,
+			},
+			size:     3,
+			index:    1,
+			expected: false,
+		},
+		{
+			// Test that a witness with the right size but wrong
+			// preimage index.
+			name: "wrong preimage index",
+			witness: wire.TxWitness{
+				dummyBytes, preimageBytes,
+			},
+			size:     2,
+			index:    0,
+			expected: false,
+		},
+	}
+
+	for _, tc := range testCases {
+		tc := tc
+
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+
+			result := checkSizeAndIndex(
+				tc.witness, tc.size, tc.index,
+			)
+			require.Equal(t, tc.expected, result)
+		})
+	}
+}
+
+// TestIsPreimageSpend tests `isPreimageSpend` can successfully detect a
+// preimage spend based on whether the commitment is local or remote.
+func TestIsPreimageSpend(t *testing.T) {
+	t.Parallel()
+
+	testCases := []struct {
+		name        string
+		witness     wire.TxWitness
+		isTaproot   bool
+		localCommit bool
+	}{
+		{
+			// Test a preimage spend on the remote commitment for
+			// taproot channels.
+			name: "tap preimage spend on remote",
+			witness: wire.TxWitness{
+				dummyBytes, dummyBytes, preimageBytes,
+				dummyBytes, dummyBytes,
+			},
+			isTaproot:   true,
+			localCommit: false,
+		},
+		{
+			// Test a preimage spend on the local commitment for
+			// taproot channels.
+			name: "tap preimage spend on local",
+			witness: wire.TxWitness{
+				dummyBytes, preimageBytes,
+				dummyBytes, dummyBytes,
+			},
+			isTaproot:   true,
+			localCommit: true,
+		},
+		{
+			// Test a preimage spend on the remote commitment for
+			// non-taproot channels.
+			name: "preimage spend on remote",
+			witness: wire.TxWitness{
+				dummyBytes, dummyBytes, dummyBytes,
+				preimageBytes, dummyBytes,
+			},
+			isTaproot:   false,
+			localCommit: false,
+		},
+		{
+			// Test a preimage spend on the local commitment for
+			// non-taproot channels.
+			name: "preimage spend on local",
+			witness: wire.TxWitness{
+				dummyBytes, preimageBytes, dummyBytes,
+			},
+			isTaproot:   false,
+			localCommit: true,
+		},
+	}
+
+	for _, tc := range testCases {
+		tc := tc
+
+		// Run the test.
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+
+			// Create a test spend detail that spends the HTLC
+			// output.
+			spend := &chainntnfs.SpendDetail{
+				SpendingTx:        &wire.MsgTx{},
+				SpenderInputIndex: 0,
+			}
+
+			// Attach the testing witness.
+			spend.SpendingTx.TxIn = []*wire.TxIn{{
+				Witness: tc.witness,
+			}}
+
+			result := isPreimageSpend(
+				tc.isTaproot, spend, tc.localCommit,
+			)
+			require.True(t, result)
+		})
+	}
+}