From 658ba445eac419de3a87af8702a6d42b6603ba88 Mon Sep 17 00:00:00 2001
From: yyforyongyu <yy2452@columbia.edu>
Date: Wed, 20 Mar 2024 16:48:58 +0800
Subject: [PATCH] sweep: delay sweeping inputs with future locktimes

This commit fixes an edge case that the sweeper's best known block
height is behind arbitrator's, which may cause an issue when creating
sweeping tx, as we may end up using an old block height from
arbitrator's view.
---
 sweep/sweeper.go      | 11 +++++++++++
 sweep/sweeper_test.go | 42 +++++++++++++++++++++++++++++++++++-------
 2 files changed, 46 insertions(+), 7 deletions(-)

diff --git a/sweep/sweeper.go b/sweep/sweeper.go
index 879e15c13..581ab8460 100644
--- a/sweep/sweeper.go
+++ b/sweep/sweeper.go
@@ -1520,6 +1520,17 @@ func (s *UtxoSweeper) updateSweeperInputs() InputsMap {
 			continue
 		}
 
+		// If the input has a locktime that's not yet reached, we will
+		// skip this input and wait for the locktime to be reached.
+		locktime, _ := input.RequiredLockTime()
+		if uint32(s.currentHeight) < locktime {
+			log.Warnf("Skipping input %v due to locktime=%v not "+
+				"reached, current height is %v", op, locktime,
+				s.currentHeight)
+
+			continue
+		}
+
 		// If this input is new or has been failed to be published,
 		// we'd retry it. The assumption here is that when an error is
 		// returned from `PublishTransaction`, it means the tx has
diff --git a/sweep/sweeper_test.go b/sweep/sweeper_test.go
index 7aa7a260f..800e0900c 100644
--- a/sweep/sweeper_test.go
+++ b/sweep/sweeper_test.go
@@ -2478,14 +2478,40 @@ func TestUpdateSweeperInputs(t *testing.T) {
 	// Create a test sweeper.
 	s := New(nil)
 
+	// Create mock inputs.
+	inp1 := &input.MockInput{}
+	defer inp1.AssertExpectations(t)
+	inp2 := &input.MockInput{}
+	defer inp2.AssertExpectations(t)
+
 	// Create a list of inputs using all the states.
-	input0 := &SweeperInput{state: Init}
-	input1 := &SweeperInput{state: PendingPublish}
-	input2 := &SweeperInput{state: Published}
-	input3 := &SweeperInput{state: PublishFailed}
-	input4 := &SweeperInput{state: Swept}
-	input5 := &SweeperInput{state: Excluded}
-	input6 := &SweeperInput{state: Failed}
+	//
+	// Mock the input to have a locktime that's matured so it will be
+	// returned.
+	inp1.On("RequiredLockTime").Return(
+		uint32(s.currentHeight), false).Once()
+	input0 := &SweeperInput{state: Init, Input: inp1}
+
+	// These inputs won't hit RequiredLockTime so we won't mock.
+	input1 := &SweeperInput{state: PendingPublish, Input: inp1}
+	input2 := &SweeperInput{state: Published, Input: inp1}
+
+	// Mock the input to have a locktime that's matured so it will be
+	// returned.
+	inp1.On("RequiredLockTime").Return(
+		uint32(s.currentHeight), false).Once()
+	input3 := &SweeperInput{state: PublishFailed, Input: inp1}
+
+	// These inputs won't hit RequiredLockTime so we won't mock.
+	input4 := &SweeperInput{state: Swept, Input: inp1}
+	input5 := &SweeperInput{state: Excluded, Input: inp1}
+	input6 := &SweeperInput{state: Failed, Input: inp1}
+
+	// Mock the input to have a locktime in the future so it will NOT be
+	// returned.
+	inp2.On("RequiredLockTime").Return(
+		uint32(s.currentHeight+1), true).Once()
+	input7 := &SweeperInput{state: Init, Input: inp2}
 
 	// Add the inputs to the sweeper. After the update, we should see the
 	// terminated inputs being removed.
@@ -2497,6 +2523,7 @@ func TestUpdateSweeperInputs(t *testing.T) {
 		{Index: 4}: input4,
 		{Index: 5}: input5,
 		{Index: 6}: input6,
+		{Index: 7}: input7,
 	}
 
 	// We expect the inputs with `Swept`, `Excluded`, and `Failed` to be
@@ -2506,6 +2533,7 @@ func TestUpdateSweeperInputs(t *testing.T) {
 		{Index: 1}: input1,
 		{Index: 2}: input2,
 		{Index: 3}: input3,
+		{Index: 7}: input7,
 	}
 
 	// We expect only the inputs with `Init` and `PublishFailed` to be