htlcswitch/control_tower_test: test strict and non-strict ctrltwrs

htlcswitch/control_tower_test.go

@@ -24,10 +24,62 @@ func genHtlc() (*lnwire.UpdateAddHTLC, error) {
 	return htlc, nil
 }
 
-// TestPaymentControlSwitchFail checks that payment status returns to Grounded
+type paymentControlTestCase func(*testing.T, bool)
+
+var paymentControlTests = []struct {
+	name     string
+	strict   bool
+	testcase paymentControlTestCase
+}{
+	{
+		name:     "fail-strict",
+		strict:   true,
+		testcase: testPaymentControlSwitchFail,
+	},
+	{
+		name:     "double-send-strict",
+		strict:   true,
+		testcase: testPaymentControlSwitchDoubleSend,
+	},
+	{
+		name:     "double-pay-strict",
+		strict:   true,
+		testcase: testPaymentControlSwitchDoublePay,
+	},
+	{
+		name:     "fail-not-strict",
+		strict:   false,
+		testcase: testPaymentControlSwitchFail,
+	},
+	{
+		name:     "double-send-not-strict",
+		strict:   false,
+		testcase: testPaymentControlSwitchDoubleSend,
+	},
+	{
+		name:     "double-pay-not-strict",
+		strict:   false,
+		testcase: testPaymentControlSwitchDoublePay,
+	},
+}
+
+// TestPaymentControls runs a set of common tests against both the strict and
+// non-strict payment control instances. This ensures that the two both behave
+// identically when making the expected state-transitions of the stricter
+// implementation. Behavioral differences in the strict and non-strict
+// implementations are tested separately.
+func TestPaymentControls(t *testing.T) {
+	for _, test := range paymentControlTests {
+		t.Run(test.name, func(t *testing.T) {
+			test.testcase(t, test.strict)
+		})
+	}
+}
+
+// testPaymentControlSwitchFail checks that payment status returns to Grounded
 // status after failing, and that ClearForTakeoff allows another HTLC for the
 // same payment hash.
-func TestPaymentControlSwitchFail(t *testing.T) {
+func testPaymentControlSwitchFail(t *testing.T, strict bool) {
 	t.Parallel()
 
 	db, err := initDB()
@@ -35,7 +87,7 @@ func TestPaymentControlSwitchFail(t *testing.T) {
 		t.Fatalf("unable to init db: %v", err)
 	}
 
-	pControl := NewPaymentControl(db)
+	pControl := NewPaymentControl(strict, db)
 
 	htlc, err := genHtlc()
 	if err != nil {
@@ -47,15 +99,7 @@ func TestPaymentControlSwitchFail(t *testing.T) {
 		t.Fatalf("unable to send htlc message: %v", err)
 	}
 
-	pStatus, err := db.FetchPaymentStatus(htlc.PaymentHash)
+	assertPaymentStatus(t, db, htlc.PaymentHash, channeldb.StatusInFlight)
-	if err != nil {
-		t.Fatalf("unable to fetch payment status: %v", err)
-	}
-
-	if pStatus != channeldb.StatusInFlight {
-		t.Fatalf("payment status mismatch: expected %v, got %v",
-			channeldb.StatusInFlight, pStatus)
-	}
 
 	// Fail the payment, which should moved it to Grounded.
 	if err := pControl.Fail(htlc.PaymentHash); err != nil {
@@ -63,15 +107,7 @@ func TestPaymentControlSwitchFail(t *testing.T) {
 	}
 
 	// Verify the status is indeed Grounded.
-	pStatus, err = db.FetchPaymentStatus(htlc.PaymentHash)
+	assertPaymentStatus(t, db, htlc.PaymentHash, channeldb.StatusGrounded)
-	if err != nil {
-		t.Fatalf("unable to fetch payment status: %v", err)
-	}
-
-	if pStatus != channeldb.StatusGrounded {
-		t.Fatalf("payment status mismatch: expected %v, got %v",
-			channeldb.StatusGrounded, pStatus)
-	}
 
 	// Sends the htlc again, which should succeed since the prior payment
 	// failed.
@@ -79,30 +115,14 @@ func TestPaymentControlSwitchFail(t *testing.T) {
 		t.Fatalf("unable to send htlc message: %v", err)
 	}
 
-	pStatus, err = db.FetchPaymentStatus(htlc.PaymentHash)
+	assertPaymentStatus(t, db, htlc.PaymentHash, channeldb.StatusInFlight)
-	if err != nil {
-		t.Fatalf("unable to fetch payment status: %v", err)
-	}
-
-	if pStatus != channeldb.StatusInFlight {
-		t.Fatalf("payment status mismatch: expected %v, got %v",
-			channeldb.StatusInFlight, pStatus)
-	}
 
 	// Verifies that status was changed to StatusCompleted.
 	if err := pControl.Success(htlc.PaymentHash); err != nil {
 		t.Fatalf("error shouldn't have been received, got: %v", err)
 	}
 
-	pStatus, err = db.FetchPaymentStatus(htlc.PaymentHash)
+	assertPaymentStatus(t, db, htlc.PaymentHash, channeldb.StatusCompleted)
-	if err != nil {
-		t.Fatalf("unable to fetch payment status: %v", err)
-	}
-
-	if pStatus != channeldb.StatusCompleted {
-		t.Fatalf("payment status mismatch: expected %v, got %v",
-			channeldb.StatusCompleted, pStatus)
-	}
 
 	// Attempt a final payment, which should now fail since the prior
 	// payment succeed.
@@ -111,9 +131,9 @@ func TestPaymentControlSwitchFail(t *testing.T) {
 	}
 }
 
-// TestPaymentControlSwitchDoubleSend checks the ability of payment control to
+// testPaymentControlSwitchDoubleSend checks the ability of payment control to
 // prevent double sending of htlc message, when message is in StatusInFlight.
-func TestPaymentControlSwitchDoubleSend(t *testing.T) {
+func testPaymentControlSwitchDoubleSend(t *testing.T, strict bool) {
 	t.Parallel()
 
 	db, err := initDB()
@@ -121,7 +141,7 @@ func TestPaymentControlSwitchDoubleSend(t *testing.T) {
 		t.Fatalf("unable to init db: %v", err)
 	}
 
-	pControl := NewPaymentControl(db)
+	pControl := NewPaymentControl(strict, db)
 
 	htlc, err := genHtlc()
 	if err != nil {
@@ -134,15 +154,7 @@ func TestPaymentControlSwitchDoubleSend(t *testing.T) {
 		t.Fatalf("unable to send htlc message: %v", err)
 	}
 
-	pStatus, err := db.FetchPaymentStatus(htlc.PaymentHash)
+	assertPaymentStatus(t, db, htlc.PaymentHash, channeldb.StatusInFlight)
-	if err != nil {
-		t.Fatalf("unable to fetch payment status: %v", err)
-	}
-
-	if pStatus != channeldb.StatusInFlight {
-		t.Fatalf("payment status mismatch: expected %v, got %v",
-			channeldb.StatusInFlight, pStatus)
-	}
 
 	// Try to initiate double sending of htlc message with the same
 	// payment hash, should result in error indicating that payment has
@@ -155,7 +167,7 @@ func TestPaymentControlSwitchDoubleSend(t *testing.T) {
 
 // TestPaymentControlSwitchDoublePay checks the ability of payment control to
 // prevent double payment.
-func TestPaymentControlSwitchDoublePay(t *testing.T) {
+func testPaymentControlSwitchDoublePay(t *testing.T, strict bool) {
 	t.Parallel()
 
 	db, err := initDB()
@@ -163,7 +175,7 @@ func TestPaymentControlSwitchDoublePay(t *testing.T) {
 		t.Fatalf("unable to init db: %v", err)
 	}
 
-	pControl := NewPaymentControl(db)
+	pControl := NewPaymentControl(strict, db)
 
 	htlc, err := genHtlc()
 	if err != nil {
@@ -176,15 +188,7 @@ func TestPaymentControlSwitchDoublePay(t *testing.T) {
 	}
 
 	// Verify that payment is InFlight.
-	pStatus, err := db.FetchPaymentStatus(htlc.PaymentHash)
+	assertPaymentStatus(t, db, htlc.PaymentHash, channeldb.StatusInFlight)
-	if err != nil {
-		t.Fatalf("unable to fetch payment status: %v", err)
-	}
-
-	if pStatus != channeldb.StatusInFlight {
-		t.Fatalf("payment status mismatch: expected %v, got %v",
-			channeldb.StatusInFlight, pStatus)
-	}
 
 	// Move payment to completed status, second payment should return error.
 	if err := pControl.Success(htlc.PaymentHash); err != nil {
@@ -192,18 +196,156 @@ func TestPaymentControlSwitchDoublePay(t *testing.T) {
 	}
 
 	// Verify that payment is Completed.
-	pStatus, err = db.FetchPaymentStatus(htlc.PaymentHash)
+	assertPaymentStatus(t, db, htlc.PaymentHash, channeldb.StatusCompleted)
-	if err != nil {
-		t.Fatalf("unable to fetch payment status: %v", err)
-	}
-
-	if pStatus != channeldb.StatusCompleted {
-		t.Fatalf("payment status mismatch: expected %v, got %v",
-			channeldb.StatusCompleted, pStatus)
-	}
 
 	if err := pControl.ClearForTakeoff(htlc); err != ErrAlreadyPaid {
 		t.Fatalf("payment control wrong behaviour:" +
 			" double payment must trigger ErrAlreadyPaid")
 	}
 }
+
+// TestPaymentControlNonStrictSuccessesWithoutInFlight checks that a non-strict
+// payment control will allow calls to Success when no payment is in flight. This
+// is necessary to gracefully handle the case in which the switch already sent
+// out a payment for a particular payment hash in a prior db version that didn't
+// have payment statuses.
+func TestPaymentControlNonStrictSuccessesWithoutInFlight(t *testing.T) {
+	t.Parallel()
+
+	db, err := initDB()
+	if err != nil {
+		t.Fatalf("unable to init db: %v", err)
+	}
+
+	pControl := NewPaymentControl(false, db)
+
+	htlc, err := genHtlc()
+	if err != nil {
+		t.Fatalf("unable to generate htlc message: %v", err)
+	}
+
+	if err := pControl.Success(htlc.PaymentHash); err != nil {
+		t.Fatalf("unable to mark payment hash success: %v", err)
+	}
+
+	assertPaymentStatus(t, db, htlc.PaymentHash, channeldb.StatusCompleted)
+
+	err = pControl.Success(htlc.PaymentHash)
+	if err != ErrPaymentAlreadyCompleted {
+		t.Fatalf("unable to remark payment hash failed: %v", err)
+	}
+}
+
+// TestPaymentControlNonStrictFailsWithoutInFlight checks that a non-strict
+// payment control will allow calls to Fail when no payment is in flight. This
+// is necessary to gracefully handle the case in which the switch already sent
+// out a payment for a particular payment hash in a prior db version that didn't
+// have payment statuses.
+func TestPaymentControlNonStrictFailsWithoutInFlight(t *testing.T) {
+	t.Parallel()
+
+	db, err := initDB()
+	if err != nil {
+		t.Fatalf("unable to init db: %v", err)
+	}
+
+	pControl := NewPaymentControl(false, db)
+
+	htlc, err := genHtlc()
+	if err != nil {
+		t.Fatalf("unable to generate htlc message: %v", err)
+	}
+
+	if err := pControl.Fail(htlc.PaymentHash); err != nil {
+		t.Fatalf("unable to mark payment hash failed: %v", err)
+	}
+
+	assertPaymentStatus(t, db, htlc.PaymentHash, channeldb.StatusGrounded)
+
+	err = pControl.Fail(htlc.PaymentHash)
+	if err != nil {
+		t.Fatalf("unable to remark payment hash failed: %v", err)
+	}
+
+	assertPaymentStatus(t, db, htlc.PaymentHash, channeldb.StatusGrounded)
+
+	err = pControl.Success(htlc.PaymentHash)
+	if err != nil {
+		t.Fatalf("unable to remark payment hash success: %v", err)
+	}
+
+	assertPaymentStatus(t, db, htlc.PaymentHash, channeldb.StatusCompleted)
+
+	err = pControl.Fail(htlc.PaymentHash)
+	if err != ErrPaymentAlreadyCompleted {
+		t.Fatalf("unable to remark payment hash failed: %v", err)
+	}
+
+	assertPaymentStatus(t, db, htlc.PaymentHash, channeldb.StatusCompleted)
+}
+
+// TestPaymentControlStrictSuccessesWithoutInFlight checks that a strict payment
+// control will disallow calls to Success when no payment is in flight.
+func TestPaymentControlStrictSuccessesWithoutInFlight(t *testing.T) {
+	t.Parallel()
+
+	db, err := initDB()
+	if err != nil {
+		t.Fatalf("unable to init db: %v", err)
+	}
+
+	pControl := NewPaymentControl(true, db)
+
+	htlc, err := genHtlc()
+	if err != nil {
+		t.Fatalf("unable to generate htlc message: %v", err)
+	}
+
+	err = pControl.Success(htlc.PaymentHash)
+	if err != ErrPaymentNotInitiated {
+		t.Fatalf("expected ErrPaymentNotInitiated, got %v", err)
+	}
+
+	assertPaymentStatus(t, db, htlc.PaymentHash, channeldb.StatusGrounded)
+}
+
+// TestPaymentControlStrictFailsWithoutInFlight checks that a strict payment
+// control will disallow calls to Fail when no payment is in flight.
+func TestPaymentControlStrictFailsWithoutInFlight(t *testing.T) {
+	t.Parallel()
+
+	db, err := initDB()
+	if err != nil {
+		t.Fatalf("unable to init db: %v", err)
+	}
+
+	pControl := NewPaymentControl(true, db)
+
+	htlc, err := genHtlc()
+	if err != nil {
+		t.Fatalf("unable to generate htlc message: %v", err)
+	}
+
+	err = pControl.Fail(htlc.PaymentHash)
+	if err != ErrPaymentNotInitiated {
+		t.Fatalf("expected ErrPaymentNotInitiated, got %v", err)
+	}
+
+	assertPaymentStatus(t, db, htlc.PaymentHash, channeldb.StatusGrounded)
+}
+
+func assertPaymentStatus(t *testing.T, db *channeldb.DB,
+	hash [32]byte, expStatus channeldb.PaymentStatus) {
+
+	t.Helper()
+
+	pStatus, err := db.FetchPaymentStatus(hash)
+	if err != nil {
+		t.Fatalf("unable to fetch payment status: %v", err)
+	}
+
+	if pStatus != expStatus {
+		t.Fatalf("payment status mismatch: expected %v, got %v",
+			expStatus, pStatus)
+	}
+}