contractcourt/htlc_timeout_test: expand timeout tests

contractcourt/htlc_success_resolver_test.go

@@ -3,7 +3,6 @@ package contractcourt
 import (
 	"bytes"
 	"fmt"
-	"io"
 	"reflect"
 	"testing"
 
@@ -22,30 +21,41 @@ import (
 
 var testHtlcAmt = lnwire.MilliSatoshi(200000)
 
-type htlcSuccessResolverTestContext struct {
-	resolver           *htlcSuccessResolver
+type htlcResolverTestContext struct {
+	resolver ContractResolver
+
+	checkpoint func(_ ContractResolver,
+		_ ...*channeldb.ResolverReport) error
+
 	notifier           *mock.ChainNotifier
 	resolverResultChan chan resolveResult
-	t                  *testing.T
+	resolutionChan     chan ResolutionMsg
+
+	t *testing.T
 }
 
-func newHtlcSuccessResolverTextContext(t *testing.T, checkpoint io.Reader) *htlcSuccessResolverTestContext {
+func newHtlcResolverTestContext(t *testing.T,
+	newResolver func(htlc channeldb.HTLC,
+		cfg ResolverConfig) ContractResolver) *htlcResolverTestContext {
+
 	notifier := &mock.ChainNotifier{
 		EpochChan: make(chan *chainntnfs.BlockEpoch, 1),
 		SpendChan: make(chan *chainntnfs.SpendDetail, 1),
 		ConfChan:  make(chan *chainntnfs.TxConfirmation, 1),
 	}
 
-	checkPointChan := make(chan struct{}, 1)
-
-	testCtx := &htlcSuccessResolverTestContext{
-		notifier: notifier,
-		t:        t,
+	testCtx := &htlcResolverTestContext{
+		checkpoint:     nil,
+		notifier:       notifier,
+		resolutionChan: make(chan ResolutionMsg, 1),
+		t:              t,
 	}
 
+	witnessBeacon := newMockWitnessBeacon()
 	chainCfg := ChannelArbitratorConfig{
 		ChainArbitratorConfig: ChainArbitratorConfig{
-			Notifier: notifier,
+			Notifier:   notifier,
+			PreimageDB: witnessBeacon,
 			PublishTx: func(_ *wire.MsgTx, _ string) error {
 				return nil
 			},
@@ -54,6 +64,16 @@ func newHtlcSuccessResolverTextContext(t *testing.T, checkpoint io.Reader) *htlc
 				*lnwallet.IncomingHtlcResolution, uint32) error {
 				return nil
 			},
+			DeliverResolutionMsg: func(msgs ...ResolutionMsg) error {
+				if len(msgs) != 1 {
+					return fmt.Errorf("expected 1 "+
+						"resolution msg, instead got %v",
+						len(msgs))
+				}
+
+				testCtx.resolutionChan <- msgs[0]
+				return nil
+			},
 		},
 		PutResolverReport: func(_ kvdb.RwTx,
 			report *channeldb.ResolverReport) error {
@@ -61,43 +81,31 @@ func newHtlcSuccessResolverTextContext(t *testing.T, checkpoint io.Reader) *htlc
 			return nil
 		},
 	}
+	// Since we want to replace this checkpoint method later in the test,
+	// we wrap the call to it in a closure. The linter will complain about
+	// this so set nolint directive.
+	checkpointFunc := func(c ContractResolver, // nolint
+		r ...*channeldb.ResolverReport) error {
+		return testCtx.checkpoint(c, r...)
+	}
 
 	cfg := ResolverConfig{
 		ChannelArbitratorConfig: chainCfg,
-		Checkpoint: func(_ ContractResolver,
-			_ ...*channeldb.ResolverReport) error {
-
-			checkPointChan <- struct{}{}
-			return nil
-		},
+		Checkpoint:              checkpointFunc,
 	}
+
 	htlc := channeldb.HTLC{
 		RHash:     testResHash,
 		OnionBlob: testOnionBlob,
 		Amt:       testHtlcAmt,
 	}
-	if checkpoint != nil {
-		var err error
-		testCtx.resolver, err = newSuccessResolverFromReader(checkpoint, cfg)
-		if err != nil {
-			t.Fatal(err)
-		}
 
-		testCtx.resolver.Supplement(htlc)
-
-	} else {
-
-		testCtx.resolver = &htlcSuccessResolver{
-			contractResolverKit: *newContractResolverKit(cfg),
-			htlcResolution:      lnwallet.IncomingHtlcResolution{},
-			htlc:                htlc,
-		}
-	}
+	testCtx.resolver = newResolver(htlc, cfg)
 
 	return testCtx
 }
 
-func (i *htlcSuccessResolverTestContext) resolve() {
+func (i *htlcResolverTestContext) resolve() {
 	// Start resolver.
 	i.resolverResultChan = make(chan resolveResult, 1)
 	go func() {
@@ -109,7 +117,7 @@ func (i *htlcSuccessResolverTestContext) resolve() {
 	}()
 }
 
-func (i *htlcSuccessResolverTestContext) waitForResult() {
+func (i *htlcResolverTestContext) waitForResult() {
 	i.t.Helper()
 
 	result := <-i.resolverResultChan
@@ -152,11 +160,12 @@ func TestHtlcSuccessSingleStage(t *testing.T) {
 		{
 			// We send a confirmation for our sweep tx to indicate
 			// that our sweep succeeded.
-			preCheckpoint: func(ctx *htlcSuccessResolverTestContext,
+			preCheckpoint: func(ctx *htlcResolverTestContext,
 				_ bool) error {
 				// The resolver will create and publish a sweep
 				// tx.
-				ctx.resolver.Sweeper.(*mockSweeper).
+				resolver := ctx.resolver.(*htlcSuccessResolver)
+				resolver.Sweeper.(*mockSweeper).
 					createSweepTxChan <- sweepTx
 
 				// Confirm the sweep, which should resolve it.
@@ -242,7 +251,7 @@ func TestHtlcSuccessSecondStageResolution(t *testing.T) {
 			// It will then wait for the Nursery to spend the
 			// output. We send a spend notification for our output
 			// to resolve our htlc.
-			preCheckpoint: func(ctx *htlcSuccessResolverTestContext,
+			preCheckpoint: func(ctx *htlcResolverTestContext,
 				_ bool) error {
 				ctx.notifier.SpendChan <- &chainntnfs.SpendDetail{
 					SpendingTx:    sweepTx,
@@ -361,11 +370,11 @@ func TestHtlcSuccessSecondStageResolutionSweeper(t *testing.T) {
 		{
 			// The HTLC output on the commitment should be offered
 			// to the sweeper. We'll notify that it gets spent.
-			preCheckpoint: func(ctx *htlcSuccessResolverTestContext,
+			preCheckpoint: func(ctx *htlcResolverTestContext,
 				_ bool) error {
 
-				inp := <-ctx.resolver.Sweeper.(*mockSweeper).
-					sweptInputs
+				resolver := ctx.resolver.(*htlcSuccessResolver)
+				inp := <-resolver.Sweeper.(*mockSweeper).sweptInputs
 				op := inp.OutPoint()
 				if *op != commitOutpoint {
 					return fmt.Errorf("outpoint %v swept, "+
@@ -389,7 +398,7 @@ func TestHtlcSuccessSecondStageResolutionSweeper(t *testing.T) {
 		{
 			// The resolver will wait for the second-level's CSV
 			// lock to expire.
-			preCheckpoint: func(ctx *htlcSuccessResolverTestContext,
+			preCheckpoint: func(ctx *htlcResolverTestContext,
 				resumed bool) error {
 
 				// If we are resuming from a checkpoint, we
@@ -410,8 +419,8 @@ func TestHtlcSuccessSecondStageResolutionSweeper(t *testing.T) {
 
 				// We expect it to sweep the second-level
 				// transaction we notfied about above.
-				inp := <-ctx.resolver.Sweeper.(*mockSweeper).
-					sweptInputs
+				resolver := ctx.resolver.(*htlcSuccessResolver)
+				inp := <-resolver.Sweeper.(*mockSweeper).sweptInputs
 				op := inp.OutPoint()
 				exp := wire.OutPoint{
 					Hash:  reSignedHash,
@@ -451,7 +460,7 @@ type checkpoint struct {
 	// preCheckpoint is a method that will be called before we reach the
 	// checkpoint, to carry out any needed operations to drive the resolver
 	// in this stage.
-	preCheckpoint func(*htlcSuccessResolverTestContext, bool) error
+	preCheckpoint func(*htlcResolverTestContext, bool) error
 
 	// data we expect the resolver to be checkpointed with next.
 	incubating bool
@@ -471,8 +480,15 @@ func testHtlcSuccess(t *testing.T, resolution lnwallet.IncomingHtlcResolution,
 	// We first run the resolver from start to finish, ensuring it gets
 	// checkpointed at every expected stage. We store the checkpointed data
 	// for the next portion of the test.
-	ctx := newHtlcSuccessResolverTextContext(t, nil)
-	ctx.resolver.htlcResolution = resolution
+	ctx := newHtlcResolverTestContext(t,
+		func(htlc channeldb.HTLC, cfg ResolverConfig) ContractResolver {
+			return &htlcSuccessResolver{
+				contractResolverKit: *newContractResolverKit(cfg),
+				htlc:                htlc,
+				htlcResolution:      resolution,
+			}
+		},
+	)
 
 	checkpointedState := runFromCheckpoint(t, ctx, checkpoints)
 
@@ -480,8 +496,18 @@ func testHtlcSuccess(t *testing.T, resolution lnwallet.IncomingHtlcResolution,
 	// run the test from that checkpoint.
 	for i := range checkpointedState {
 		cp := bytes.NewReader(checkpointedState[i])
-		ctx := newHtlcSuccessResolverTextContext(t, cp)
-		ctx.resolver.htlcResolution = resolution
+		ctx := newHtlcResolverTestContext(t,
+			func(htlc channeldb.HTLC, cfg ResolverConfig) ContractResolver {
+				resolver, err := newSuccessResolverFromReader(cp, cfg)
+				if err != nil {
+					t.Fatal(err)
+				}
+
+				resolver.Supplement(htlc)
+				resolver.htlcResolution = resolution
+				return resolver
+			},
+		)
 
 		// Run from the given checkpoint, ensuring we'll hit the rest.
 		_ = runFromCheckpoint(t, ctx, checkpoints[i+1:])
@@ -490,7 +516,7 @@ func testHtlcSuccess(t *testing.T, resolution lnwallet.IncomingHtlcResolution,
 
 // runFromCheckpoint executes the Resolve method on the success resolver, and
 // asserts that it checkpoints itself according to the expected checkpoints.
-func runFromCheckpoint(t *testing.T, ctx *htlcSuccessResolverTestContext,
+func runFromCheckpoint(t *testing.T, ctx *htlcResolverTestContext,
 	expectedCheckpoints []checkpoint) [][]byte {
 
 	defer timeout(t)()
@@ -501,25 +527,34 @@ func runFromCheckpoint(t *testing.T, ctx *htlcSuccessResolverTestContext,
 	// checkpointed state and reports are equal to what we expect.
 	nextCheckpoint := 0
 	checkpointChan := make(chan struct{})
-	ctx.resolver.Checkpoint = func(resolver ContractResolver,
+	ctx.checkpoint = func(resolver ContractResolver,
 		reports ...*channeldb.ResolverReport) error {
 
 		if nextCheckpoint >= len(expectedCheckpoints) {
 			t.Fatal("did not expect more checkpoints")
 		}
 
-		h := resolver.(*htlcSuccessResolver)
-		cp := expectedCheckpoints[nextCheckpoint]
-
-		if h.resolved != cp.resolved {
-			t.Fatalf("expected checkpoint to be resolve=%v, had %v",
-				cp.resolved, h.resolved)
+		var resolved, incubating bool
+		if h, ok := resolver.(*htlcSuccessResolver); ok {
+			resolved = h.resolved
+			incubating = h.outputIncubating
+		}
+		if h, ok := resolver.(*htlcTimeoutResolver); ok {
+			resolved = h.resolved
+			incubating = h.outputIncubating
 		}
 
-		if !reflect.DeepEqual(h.outputIncubating, cp.incubating) {
+		cp := expectedCheckpoints[nextCheckpoint]
+
+		if resolved != cp.resolved {
+			t.Fatalf("expected checkpoint to be resolve=%v, had %v",
+				cp.resolved, resolved)
+		}
+
+		if !reflect.DeepEqual(incubating, cp.incubating) {
 			t.Fatalf("expected checkpoint to be have "+
 				"incubating=%v, had %v", cp.incubating,
-				h.outputIncubating)
+				incubating)
 
 		}
 

contractcourt/htlc_timeout_resolver_test.go

@@ -8,6 +8,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/btcsuite/btcd/chaincfg/chainhash"
 	"github.com/btcsuite/btcd/txscript"
 	"github.com/btcsuite/btcd/wire"
 	"github.com/btcsuite/btcutil"
@@ -481,3 +482,834 @@ func TestHtlcTimeoutResolver(t *testing.T) {
 		}
 	}
 }
+
+// NOTE: the following tests essentially checks many of the same scenarios as
+// the test above, but they expand on it by checking resuming from checkpoints
+// at every stage.
+
+// TestHtlcTimeoutSingleStage tests a remote commitment confirming, and the
+// local node sweeping the HTLC output directly after timeout.
+func TestHtlcTimeoutSingleStage(t *testing.T) {
+	commitOutpoint := wire.OutPoint{Index: 3}
+
+	sweepTx := &wire.MsgTx{
+		TxIn:  []*wire.TxIn{{}},
+		TxOut: []*wire.TxOut{{}},
+	}
+
+	// singleStageResolution is a resolution for a htlc on the remote
+	// party's commitment.
+	singleStageResolution := lnwallet.OutgoingHtlcResolution{
+		ClaimOutpoint: commitOutpoint,
+		SweepSignDesc: testSignDesc,
+	}
+
+	sweepTxid := sweepTx.TxHash()
+	claim := &channeldb.ResolverReport{
+		OutPoint:        commitOutpoint,
+		Amount:          btcutil.Amount(testSignDesc.Output.Value),
+		ResolverType:    channeldb.ResolverTypeOutgoingHtlc,
+		ResolverOutcome: channeldb.ResolverOutcomeTimeout,
+		SpendTxID:       &sweepTxid,
+	}
+
+	checkpoints := []checkpoint{
+		{
+			// Output should be handed off to the nursery.
+			incubating: true,
+		},
+		{
+			// We send a confirmation the sweep tx from published
+			// by the nursery.
+			preCheckpoint: func(ctx *htlcResolverTestContext,
+				_ bool) error {
+				// The nursery will create and publish a sweep
+				// tx.
+				ctx.notifier.SpendChan <- &chainntnfs.SpendDetail{
+					SpendingTx:    sweepTx,
+					SpenderTxHash: &sweepTxid,
+				}
+
+				// The resolver should deliver a failure
+				// resolition message (indicating we
+				// successfully timed out the HTLC).
+				select {
+				case resolutionMsg := <-ctx.resolutionChan:
+					if resolutionMsg.Failure == nil {
+						t.Fatalf("expected failure resolution msg")
+					}
+				case <-time.After(time.Second * 5):
+					t.Fatalf("resolution not sent")
+				}
+
+				return nil
+			},
+
+			// After the sweep has confirmed, we expect the
+			// checkpoint to be resolved, and with the above
+			// report.
+			incubating: true,
+			resolved:   true,
+			reports: []*channeldb.ResolverReport{
+				claim,
+			},
+		},
+	}
+
+	testHtlcTimeout(
+		t, singleStageResolution, checkpoints,
+	)
+}
+
+// TestHtlcTimeoutSecondStage tests a local commitment being confirmed, and the
+// local node claiming the HTLC output using the second-level timeout tx.
+func TestHtlcTimeoutSecondStage(t *testing.T) {
+	commitOutpoint := wire.OutPoint{Index: 2}
+	htlcOutpoint := wire.OutPoint{Index: 3}
+
+	sweepTx := &wire.MsgTx{
+		TxIn:  []*wire.TxIn{{}},
+		TxOut: []*wire.TxOut{{}},
+	}
+	sweepHash := sweepTx.TxHash()
+
+	timeoutTx := &wire.MsgTx{
+		TxIn: []*wire.TxIn{
+			{
+				PreviousOutPoint: commitOutpoint,
+			},
+		},
+		TxOut: []*wire.TxOut{
+			{
+				Value:    111,
+				PkScript: []byte{0xaa, 0xaa},
+			},
+		},
+	}
+
+	signer := &mock.DummySigner{}
+	witness, err := input.SenderHtlcSpendTimeout(
+		&mock.DummySignature{}, txscript.SigHashAll,
+		signer, &testSignDesc, timeoutTx,
+	)
+	require.NoError(t, err)
+	timeoutTx.TxIn[0].Witness = witness
+
+	timeoutTxid := timeoutTx.TxHash()
+
+	// twoStageResolution is a resolution for a htlc on the local
+	// party's commitment.
+	twoStageResolution := lnwallet.OutgoingHtlcResolution{
+		ClaimOutpoint:   htlcOutpoint,
+		SignedTimeoutTx: timeoutTx,
+		SweepSignDesc:   testSignDesc,
+	}
+
+	firstStage := &channeldb.ResolverReport{
+		OutPoint:        commitOutpoint,
+		Amount:          testHtlcAmt.ToSatoshis(),
+		ResolverType:    channeldb.ResolverTypeOutgoingHtlc,
+		ResolverOutcome: channeldb.ResolverOutcomeFirstStage,
+		SpendTxID:       &timeoutTxid,
+	}
+
+	secondState := &channeldb.ResolverReport{
+		OutPoint:        htlcOutpoint,
+		Amount:          btcutil.Amount(testSignDesc.Output.Value),
+		ResolverType:    channeldb.ResolverTypeOutgoingHtlc,
+		ResolverOutcome: channeldb.ResolverOutcomeTimeout,
+		SpendTxID:       &sweepHash,
+	}
+
+	checkpoints := []checkpoint{
+		{
+			// Output should be handed off to the nursery.
+			incubating: true,
+		},
+		{
+			// We send a confirmation for our sweep tx to indicate
+			// that our sweep succeeded.
+			preCheckpoint: func(ctx *htlcResolverTestContext,
+				_ bool) error {
+				// The nursery will publish the timeout tx.
+				ctx.notifier.SpendChan <- &chainntnfs.SpendDetail{
+					SpendingTx:    timeoutTx,
+					SpenderTxHash: &timeoutTxid,
+				}
+
+				// The resolver should deliver a failure
+				// resolution message (indicating we
+				// successfully timed out the HTLC).
+				select {
+				case resolutionMsg := <-ctx.resolutionChan:
+					if resolutionMsg.Failure == nil {
+						t.Fatalf("expected failure resolution msg")
+					}
+				case <-time.After(time.Second * 1):
+					t.Fatalf("resolution not sent")
+				}
+
+				// Deliver spend of timeout tx.
+				ctx.notifier.SpendChan <- &chainntnfs.SpendDetail{
+					SpendingTx:    sweepTx,
+					SpenderTxHash: &sweepHash,
+				}
+
+				return nil
+			},
+
+			// After the sweep has confirmed, we expect the
+			// checkpoint to be resolved, and with the above
+			// reports.
+			incubating: true,
+			resolved:   true,
+			reports: []*channeldb.ResolverReport{
+				firstStage, secondState,
+			},
+		},
+	}
+
+	testHtlcTimeout(
+		t, twoStageResolution, checkpoints,
+	)
+}
+
+// TestHtlcTimeoutSingleStageRemoteSpend tests that when a local commitment
+// confirms, and the remote spends the HTLC output directly, we detect this and
+// extract the preimage.
+func TestHtlcTimeoutSingleStageRemoteSpend(t *testing.T) {
+	commitOutpoint := wire.OutPoint{Index: 2}
+	htlcOutpoint := wire.OutPoint{Index: 3}
+
+	spendTx := &wire.MsgTx{
+		TxIn:  []*wire.TxIn{{}},
+		TxOut: []*wire.TxOut{{}},
+	}
+
+	fakePreimageBytes := bytes.Repeat([]byte{1}, lntypes.HashSize)
+	var fakePreimage lntypes.Preimage
+	copy(fakePreimage[:], fakePreimageBytes)
+
+	signer := &mock.DummySigner{}
+	witness, err := input.SenderHtlcSpendRedeem(
+		signer, &testSignDesc, spendTx,
+		fakePreimageBytes,
+	)
+	require.NoError(t, err)
+	spendTx.TxIn[0].Witness = witness
+
+	spendTxHash := spendTx.TxHash()
+
+	timeoutTx := &wire.MsgTx{
+		TxIn: []*wire.TxIn{
+			{
+				PreviousOutPoint: commitOutpoint,
+			},
+		},
+		TxOut: []*wire.TxOut{
+			{
+				Value:    123,
+				PkScript: []byte{0xff, 0xff},
+			},
+		},
+	}
+
+	timeoutWitness, err := input.SenderHtlcSpendTimeout(
+		&mock.DummySignature{}, txscript.SigHashAll,
+		signer, &testSignDesc, timeoutTx,
+	)
+	require.NoError(t, err)
+	timeoutTx.TxIn[0].Witness = timeoutWitness
+
+	// twoStageResolution is a resolution for a htlc on the local
+	// party's commitment.
+	twoStageResolution := lnwallet.OutgoingHtlcResolution{
+		ClaimOutpoint:   htlcOutpoint,
+		SignedTimeoutTx: timeoutTx,
+		SweepSignDesc:   testSignDesc,
+	}
+
+	claim := &channeldb.ResolverReport{
+		OutPoint:        htlcOutpoint,
+		Amount:          btcutil.Amount(testSignDesc.Output.Value),
+		ResolverType:    channeldb.ResolverTypeOutgoingHtlc,
+		ResolverOutcome: channeldb.ResolverOutcomeClaimed,
+		SpendTxID:       &spendTxHash,
+	}
+
+	checkpoints := []checkpoint{
+		{
+			// Output should be handed off to the nursery.
+			incubating: true,
+		},
+		{
+			// We send a spend notification for a remote spend with
+			// the preimage.
+			preCheckpoint: func(ctx *htlcResolverTestContext,
+				_ bool) error {
+
+				witnessBeacon := ctx.resolver.(*htlcTimeoutResolver).PreimageDB.(*mockWitnessBeacon)
+
+				// The remote spends the output direcly with
+				// the preimage.
+				ctx.notifier.SpendChan <- &chainntnfs.SpendDetail{
+					SpendingTx:    spendTx,
+					SpenderTxHash: &spendTxHash,
+				}
+
+				// We should extract the preimage.
+				select {
+				case newPreimage := <-witnessBeacon.newPreimages:
+					if newPreimage[0] != fakePreimage {
+						t.Fatalf("wrong pre-image: "+
+							"expected %v, got %v",
+							fakePreimage, newPreimage)
+					}
+
+				case <-time.After(time.Second * 5):
+					t.Fatalf("pre-image not added")
+				}
+
+				// Finally, we should get a resolution message
+				// with the pre-image set within the message.
+				select {
+				case resolutionMsg := <-ctx.resolutionChan:
+					if *resolutionMsg.PreImage != fakePreimage {
+						t.Fatalf("wrong pre-image: "+
+							"expected %v, got %v",
+							fakePreimage, resolutionMsg.PreImage)
+					}
+				case <-time.After(time.Second * 5):
+					t.Fatalf("resolution not sent")
+				}
+
+				return nil
+			},
+
+			// After the success tx has confirmed, we expect the
+			// checkpoint to be resolved, and with the above
+			// report.
+			incubating: true,
+			resolved:   true,
+			reports: []*channeldb.ResolverReport{
+				claim,
+			},
+		},
+	}
+
+	testHtlcTimeout(
+		t, twoStageResolution, checkpoints,
+	)
+}
+
+// TestHtlcTimeoutSecondStageRemoteSpend tests that when a remite commitment
+// confirms, and the remote spends the output using the success tx, we
+// properly detect this and extract the preimage.
+func TestHtlcTimeoutSecondStageRemoteSpend(t *testing.T) {
+	commitOutpoint := wire.OutPoint{Index: 2}
+
+	remoteSuccessTx := &wire.MsgTx{
+		TxIn: []*wire.TxIn{
+			{
+				PreviousOutPoint: commitOutpoint,
+			},
+		},
+		TxOut: []*wire.TxOut{},
+	}
+
+	fakePreimageBytes := bytes.Repeat([]byte{1}, lntypes.HashSize)
+	var fakePreimage lntypes.Preimage
+	copy(fakePreimage[:], fakePreimageBytes)
+
+	signer := &mock.DummySigner{}
+	witness, err := input.ReceiverHtlcSpendRedeem(
+		&mock.DummySignature{}, txscript.SigHashAll,
+		fakePreimageBytes, signer,
+		&testSignDesc, remoteSuccessTx,
+	)
+	require.NoError(t, err)
+
+	remoteSuccessTx.TxIn[0].Witness = witness
+	successTxid := remoteSuccessTx.TxHash()
+
+	// singleStageResolution allwoing the local node to sweep HTLC output
+	// directly from the remote commitment after timeout.
+	singleStageResolution := lnwallet.OutgoingHtlcResolution{
+		ClaimOutpoint: commitOutpoint,
+		SweepSignDesc: testSignDesc,
+	}
+
+	claim := &channeldb.ResolverReport{
+		OutPoint:        commitOutpoint,
+		Amount:          btcutil.Amount(testSignDesc.Output.Value),
+		ResolverType:    channeldb.ResolverTypeOutgoingHtlc,
+		ResolverOutcome: channeldb.ResolverOutcomeClaimed,
+		SpendTxID:       &successTxid,
+	}
+
+	checkpoints := []checkpoint{
+		{
+			// Output should be handed off to the nursery.
+			incubating: true,
+		},
+		{
+			// We send a confirmation for the remote's second layer
+			// success transcation.
+			preCheckpoint: func(ctx *htlcResolverTestContext,
+				_ bool) error {
+				ctx.notifier.SpendChan <- &chainntnfs.SpendDetail{
+					SpendingTx:    remoteSuccessTx,
+					SpenderTxHash: &successTxid,
+				}
+
+				witnessBeacon := ctx.resolver.(*htlcTimeoutResolver).PreimageDB.(*mockWitnessBeacon)
+
+				// We expect the preimage to be extracted,
+				select {
+				case newPreimage := <-witnessBeacon.newPreimages:
+					if newPreimage[0] != fakePreimage {
+						t.Fatalf("wrong pre-image: "+
+							"expected %v, got %v",
+							fakePreimage, newPreimage)
+					}
+
+				case <-time.After(time.Second * 5):
+					t.Fatalf("pre-image not added")
+				}
+
+				// Finally, we should get a resolution message with the
+				// pre-image set within the message.
+				select {
+				case resolutionMsg := <-ctx.resolutionChan:
+					if *resolutionMsg.PreImage != fakePreimage {
+						t.Fatalf("wrong pre-image: "+
+							"expected %v, got %v",
+							fakePreimage, resolutionMsg.PreImage)
+					}
+				case <-time.After(time.Second * 5):
+					t.Fatalf("resolution not sent")
+				}
+
+				return nil
+			},
+
+			// After the sweep has confirmed, we expect the
+			// checkpoint to be resolved, and with the above
+			// report.
+			incubating: true,
+			resolved:   true,
+			reports: []*channeldb.ResolverReport{
+				claim,
+			},
+		},
+	}
+
+	testHtlcTimeout(
+		t, singleStageResolution, checkpoints,
+	)
+}
+
+// TestHtlcTimeoutSecondStageSweeper tests that for anchor channels, when a
+// local commitment confirms, the timeout tx is handed to the sweeper to claim
+// the HTLC output.
+func TestHtlcTimeoutSecondStageSweeper(t *testing.T) {
+	commitOutpoint := wire.OutPoint{Index: 2}
+	htlcOutpoint := wire.OutPoint{Index: 3}
+
+	sweepTx := &wire.MsgTx{
+		TxIn:  []*wire.TxIn{{}},
+		TxOut: []*wire.TxOut{{}},
+	}
+	sweepHash := sweepTx.TxHash()
+
+	timeoutTx := &wire.MsgTx{
+		TxIn: []*wire.TxIn{
+			{
+				PreviousOutPoint: commitOutpoint,
+			},
+		},
+		TxOut: []*wire.TxOut{
+			{
+				Value:    123,
+				PkScript: []byte{0xff, 0xff},
+			},
+		},
+	}
+
+	// We set the timeout witness since the script is used when subscribing
+	// to spends.
+	signer := &mock.DummySigner{}
+	timeoutWitness, err := input.SenderHtlcSpendTimeout(
+		&mock.DummySignature{}, txscript.SigHashAll,
+		signer, &testSignDesc, timeoutTx,
+	)
+	require.NoError(t, err)
+	timeoutTx.TxIn[0].Witness = timeoutWitness
+
+	reSignedTimeoutTx := &wire.MsgTx{
+		TxIn: []*wire.TxIn{
+			{
+				PreviousOutPoint: wire.OutPoint{
+					Hash:  chainhash.Hash{0xaa, 0xbb},
+					Index: 0,
+				},
+			},
+			timeoutTx.TxIn[0],
+			{
+				PreviousOutPoint: wire.OutPoint{
+					Hash:  chainhash.Hash{0xaa, 0xbb},
+					Index: 2,
+				},
+			},
+		},
+
+		TxOut: []*wire.TxOut{
+			{
+				Value:    111,
+				PkScript: []byte{0xaa, 0xaa},
+			},
+			timeoutTx.TxOut[0],
+		},
+	}
+	reSignedHash := reSignedTimeoutTx.TxHash()
+	reSignedOutPoint := wire.OutPoint{
+		Hash:  reSignedHash,
+		Index: 1,
+	}
+
+	// twoStageResolution is a resolution for a htlc on the local
+	// party's commitment, where the timout tx can be re-signed.
+	twoStageResolution := lnwallet.OutgoingHtlcResolution{
+		ClaimOutpoint:   htlcOutpoint,
+		SignedTimeoutTx: timeoutTx,
+		SignDetails: &input.SignDetails{
+			SignDesc: testSignDesc,
+			PeerSig:  testSig,
+		},
+		SweepSignDesc: testSignDesc,
+	}
+
+	firstStage := &channeldb.ResolverReport{
+		OutPoint:        commitOutpoint,
+		Amount:          testHtlcAmt.ToSatoshis(),
+		ResolverType:    channeldb.ResolverTypeOutgoingHtlc,
+		ResolverOutcome: channeldb.ResolverOutcomeFirstStage,
+		SpendTxID:       &reSignedHash,
+	}
+
+	secondState := &channeldb.ResolverReport{
+		OutPoint:        reSignedOutPoint,
+		Amount:          btcutil.Amount(testSignDesc.Output.Value),
+		ResolverType:    channeldb.ResolverTypeOutgoingHtlc,
+		ResolverOutcome: channeldb.ResolverOutcomeTimeout,
+		SpendTxID:       &sweepHash,
+	}
+
+	checkpoints := []checkpoint{
+		{
+			// The output should be given to the sweeper.
+			preCheckpoint: func(ctx *htlcResolverTestContext,
+				_ bool) error {
+
+				resolver := ctx.resolver.(*htlcTimeoutResolver)
+				inp := <-resolver.Sweeper.(*mockSweeper).sweptInputs
+				op := inp.OutPoint()
+				if *op != commitOutpoint {
+					return fmt.Errorf("outpoint %v swept, "+
+						"expected %v", op,
+						commitOutpoint)
+				}
+
+				// Emulat the sweeper spending using the
+				// re-signed timeout tx.
+				ctx.notifier.SpendChan <- &chainntnfs.SpendDetail{
+					SpendingTx:        reSignedTimeoutTx,
+					SpenderInputIndex: 1,
+					SpenderTxHash:     &reSignedHash,
+					SpendingHeight:    10,
+				}
+
+				return nil
+			},
+			// incubating=true is used to signal that the
+			// second-level transaction was confirmed.
+			incubating: true,
+		},
+		{
+			// We send a confirmation for our sweep tx to indicate
+			// that our sweep succeeded.
+			preCheckpoint: func(ctx *htlcResolverTestContext,
+				resumed bool) error {
+
+				// If we are resuming from a checkpoing, we
+				// expect the resolver to re-subscribe to a
+				// spend, hence we must resend it.
+				if resumed {
+					ctx.notifier.SpendChan <- &chainntnfs.SpendDetail{
+						SpendingTx:        reSignedTimeoutTx,
+						SpenderInputIndex: 1,
+						SpenderTxHash:     &reSignedHash,
+						SpendingHeight:    10,
+					}
+				}
+
+				// The resolver should deliver a failure
+				// resolution message (indicating we
+				// successfully timed out the HTLC).
+				select {
+				case resolutionMsg := <-ctx.resolutionChan:
+					if resolutionMsg.Failure == nil {
+						t.Fatalf("expected failure resolution msg")
+					}
+				case <-time.After(time.Second * 1):
+					t.Fatalf("resolution not sent")
+				}
+
+				// Mimic CSV lock expiring.
+				ctx.notifier.EpochChan <- &chainntnfs.BlockEpoch{
+					Height: 13,
+				}
+
+				// The timout tx output should now be given to
+				// the sweeper.
+				resolver := ctx.resolver.(*htlcTimeoutResolver)
+				inp := <-resolver.Sweeper.(*mockSweeper).sweptInputs
+				op := inp.OutPoint()
+				exp := wire.OutPoint{
+					Hash:  reSignedHash,
+					Index: 1,
+				}
+				if *op != exp {
+					return fmt.Errorf("wrong outpoint swept")
+				}
+
+				// Notify about the spend, which should resolve
+				// the resolver.
+				ctx.notifier.SpendChan <- &chainntnfs.SpendDetail{
+					SpendingTx:     sweepTx,
+					SpenderTxHash:  &sweepHash,
+					SpendingHeight: 14,
+				}
+
+				return nil
+			},
+
+			// After the sweep has confirmed, we expect the
+			// checkpoint to be resolved, and with the above
+			// reports.
+			incubating: true,
+			resolved:   true,
+			reports: []*channeldb.ResolverReport{
+				firstStage,
+				secondState,
+			},
+		},
+	}
+
+	testHtlcTimeout(
+		t, twoStageResolution, checkpoints,
+	)
+}
+
+// TestHtlcTimeoutSecondStageSweeperRemoteSpend tests that if a local timeout
+// tx is offered to the sweeper, but the output is swept by the remote node, we
+// properly detect this and extract the preimage.
+func TestHtlcTimeoutSecondStageSweeperRemoteSpend(t *testing.T) {
+	commitOutpoint := wire.OutPoint{Index: 2}
+	htlcOutpoint := wire.OutPoint{Index: 3}
+
+	timeoutTx := &wire.MsgTx{
+		TxIn: []*wire.TxIn{
+			{
+				PreviousOutPoint: commitOutpoint,
+			},
+		},
+		TxOut: []*wire.TxOut{
+			{
+				Value:    123,
+				PkScript: []byte{0xff, 0xff},
+			},
+		},
+	}
+
+	// We set the timeout witness since the script is used when subscribing
+	// to spends.
+	signer := &mock.DummySigner{}
+	timeoutWitness, err := input.SenderHtlcSpendTimeout(
+		&mock.DummySignature{}, txscript.SigHashAll,
+		signer, &testSignDesc, timeoutTx,
+	)
+	require.NoError(t, err)
+	timeoutTx.TxIn[0].Witness = timeoutWitness
+
+	spendTx := &wire.MsgTx{
+		TxIn:  []*wire.TxIn{{}},
+		TxOut: []*wire.TxOut{{}},
+	}
+
+	fakePreimageBytes := bytes.Repeat([]byte{1}, lntypes.HashSize)
+	var fakePreimage lntypes.Preimage
+	copy(fakePreimage[:], fakePreimageBytes)
+
+	witness, err := input.SenderHtlcSpendRedeem(
+		signer, &testSignDesc, spendTx,
+		fakePreimageBytes,
+	)
+	require.NoError(t, err)
+	spendTx.TxIn[0].Witness = witness
+
+	spendTxHash := spendTx.TxHash()
+
+	// twoStageResolution is a resolution for a htlc on the local
+	// party's commitment, where the timout tx can be re-signed.
+	twoStageResolution := lnwallet.OutgoingHtlcResolution{
+		ClaimOutpoint:   htlcOutpoint,
+		SignedTimeoutTx: timeoutTx,
+		SignDetails: &input.SignDetails{
+			SignDesc: testSignDesc,
+			PeerSig:  testSig,
+		},
+		SweepSignDesc: testSignDesc,
+	}
+
+	claim := &channeldb.ResolverReport{
+		OutPoint:        htlcOutpoint,
+		Amount:          btcutil.Amount(testSignDesc.Output.Value),
+		ResolverType:    channeldb.ResolverTypeOutgoingHtlc,
+		ResolverOutcome: channeldb.ResolverOutcomeClaimed,
+		SpendTxID:       &spendTxHash,
+	}
+
+	checkpoints := []checkpoint{
+		{
+			// The output should be given to the sweeper.
+			preCheckpoint: func(ctx *htlcResolverTestContext,
+				_ bool) error {
+
+				resolver := ctx.resolver.(*htlcTimeoutResolver)
+				inp := <-resolver.Sweeper.(*mockSweeper).sweptInputs
+				op := inp.OutPoint()
+				if *op != commitOutpoint {
+					return fmt.Errorf("outpoint %v swept, "+
+						"expected %v", op,
+						commitOutpoint)
+				}
+
+				// Emulate the remote sweeping the output with the preimage.
+				// re-signed timeout tx.
+				ctx.notifier.SpendChan <- &chainntnfs.SpendDetail{
+					SpendingTx:    spendTx,
+					SpenderTxHash: &spendTxHash,
+				}
+
+				return nil
+			},
+			// incubating=true is used to signal that the
+			// second-level transaction was confirmed.
+			incubating: true,
+		},
+		{
+			// We send a confirmation for our sweep tx to indicate
+			// that our sweep succeeded.
+			preCheckpoint: func(ctx *htlcResolverTestContext,
+				resumed bool) error {
+
+				// If we are resuming from a checkpoing, we
+				// expect the resolver to re-subscribe to a
+				// spend, hence we must resend it.
+				if resumed {
+					fmt.Println("resumed")
+					ctx.notifier.SpendChan <- &chainntnfs.SpendDetail{
+						SpendingTx:    spendTx,
+						SpenderTxHash: &spendTxHash,
+					}
+				}
+
+				witnessBeacon := ctx.resolver.(*htlcTimeoutResolver).PreimageDB.(*mockWitnessBeacon)
+
+				// We should extract the preimage.
+				select {
+				case newPreimage := <-witnessBeacon.newPreimages:
+					if newPreimage[0] != fakePreimage {
+						t.Fatalf("wrong pre-image: "+
+							"expected %v, got %v",
+							fakePreimage, newPreimage)
+					}
+
+				case <-time.After(time.Second * 5):
+					t.Fatalf("pre-image not added")
+				}
+
+				// Finally, we should get a resolution message
+				// with the pre-image set within the message.
+				select {
+				case resolutionMsg := <-ctx.resolutionChan:
+					if *resolutionMsg.PreImage != fakePreimage {
+						t.Fatalf("wrong pre-image: "+
+							"expected %v, got %v",
+							fakePreimage, resolutionMsg.PreImage)
+					}
+				case <-time.After(time.Second * 5):
+					t.Fatalf("resolution not sent")
+				}
+
+				return nil
+			},
+
+			// After the sweep has confirmed, we expect the
+			// checkpoint to be resolved, and with the above
+			// reports.
+			incubating: true,
+			resolved:   true,
+			reports: []*channeldb.ResolverReport{
+				claim,
+			},
+		},
+	}
+
+	testHtlcTimeout(
+		t, twoStageResolution, checkpoints,
+	)
+}
+
+func testHtlcTimeout(t *testing.T, resolution lnwallet.OutgoingHtlcResolution,
+	checkpoints []checkpoint) {
+
+	defer timeout(t)()
+
+	// We first run the resolver from start to finish, ensuring it gets
+	// checkpointed at every expected stage. We store the checkpointed data
+	// for the next portion of the test.
+	ctx := newHtlcResolverTestContext(t,
+		func(htlc channeldb.HTLC, cfg ResolverConfig) ContractResolver {
+			return &htlcTimeoutResolver{
+				contractResolverKit: *newContractResolverKit(cfg),
+				htlc:                htlc,
+				htlcResolution:      resolution,
+			}
+		},
+	)
+
+	checkpointedState := runFromCheckpoint(t, ctx, checkpoints)
+
+	// Now, from every checkpoint created, we re-create the resolver, and
+	// run the test from that checkpoint.
+	for i := range checkpointedState {
+		cp := bytes.NewReader(checkpointedState[i])
+		ctx := newHtlcResolverTestContext(t,
+			func(htlc channeldb.HTLC, cfg ResolverConfig) ContractResolver {
+				resolver, err := newTimeoutResolverFromReader(cp, cfg)
+				if err != nil {
+					t.Fatal(err)
+				}
+
+				resolver.Supplement(htlc)
+				resolver.htlcResolution = resolution
+				return resolver
+			},
+		)
+
+		// Run from the given checkpoint, ensuring we'll hit the rest.
+		_ = runFromCheckpoint(t, ctx, checkpoints[i+1:])
+	}
+}