From 58d95be4dd7e2b39e79f675bda4c14bd225fc477 Mon Sep 17 00:00:00 2001
From: carla <kirkcohenc@gmail.com>
Date: Fri, 23 Apr 2021 08:39:45 +0200
Subject: [PATCH] multi: change RegisterAttempt error checking order

Move our more generic terminal check forward so that we only
need to handle a single class of expected errors. This change
is mirrored in our mock, and our reproducing tests are updated
to assert that this move catches both classes of errors we get.
---
 channeldb/payment_control.go      | 13 +++++++------
 channeldb/payment_control_test.go |  6 +-----
 routing/mock_test.go              |  8 ++++----
 routing/payment_lifecycle_test.go |  4 ++--
 4 files changed, 14 insertions(+), 17 deletions(-)

diff --git a/channeldb/payment_control.go b/channeldb/payment_control.go
index 9ed5e6392..c820778cb 100644
--- a/channeldb/payment_control.go
+++ b/channeldb/payment_control.go
@@ -290,18 +290,19 @@ func (p *PaymentControl) RegisterAttempt(paymentHash lntypes.Hash,
 			return err
 		}
 
-		// Ensure the payment is in-flight.
-		if err := ensureInFlight(p); err != nil {
-			return err
-		}
-
 		// We cannot register a new attempt if the payment already has
-		// reached a terminal condition:
+		// reached a terminal condition. We check this before
+		// ensureInFlight because it is a more general check.
 		settle, fail := p.TerminalInfo()
 		if settle != nil || fail != nil {
 			return ErrPaymentTerminal
 		}
 
+		// Ensure the payment is in-flight.
+		if err := ensureInFlight(p); err != nil {
+			return err
+		}
+
 		// Make sure any existing shards match the new one with regards
 		// to MPP options.
 		mpp := attempt.Route.FinalHop().MPP
diff --git a/channeldb/payment_control_test.go b/channeldb/payment_control_test.go
index e13f9d720..5bf6a05d1 100644
--- a/channeldb/payment_control_test.go
+++ b/channeldb/payment_control_test.go
@@ -1013,19 +1013,15 @@ func TestPaymentControlMultiShard(t *testing.T) {
 		// up in the Succeeded state. If both failed the payment should
 		// also be Failed at this poinnt.
 		finalStatus := StatusFailed
-		expRegErr := ErrPaymentAlreadyFailed
 		if test.settleFirst || test.settleLast {
 			finalStatus = StatusSucceeded
-			expRegErr = ErrPaymentAlreadySucceeded
 		}
 
 		assertPaymentStatus(t, pControl, info.PaymentHash, finalStatus)
 
 		// Finally assert we cannot register more attempts.
 		_, err = pControl.RegisterAttempt(info.PaymentHash, &b)
-		if err != expRegErr {
-			t.Fatalf("expected error %v, got: %v", expRegErr, err)
-		}
+		require.Equal(t, ErrPaymentTerminal, err)
 	}
 
 	for _, test := range tests {
diff --git a/routing/mock_test.go b/routing/mock_test.go
index 0d0b7a5b8..a477fb10e 100644
--- a/routing/mock_test.go
+++ b/routing/mock_test.go
@@ -325,6 +325,10 @@ func (m *mockControlTower) RegisterAttempt(phash lntypes.Hash,
 	_, settled := m.successful[phash]
 	_, failed := m.failed[phash]
 
+	if settled || failed {
+		return channeldb.ErrPaymentTerminal
+	}
+
 	if settled && !inFlight {
 		return channeldb.ErrPaymentAlreadySucceeded
 	}
@@ -333,10 +337,6 @@ func (m *mockControlTower) RegisterAttempt(phash lntypes.Hash,
 		return channeldb.ErrPaymentAlreadyFailed
 	}
 
-	if settled || failed {
-		return channeldb.ErrPaymentTerminal
-	}
-
 	// Add attempt to payment.
 	p.attempts = append(p.attempts, channeldb.HTLCAttempt{
 		HTLCAttemptInfo: *a,
diff --git a/routing/payment_lifecycle_test.go b/routing/payment_lifecycle_test.go
index 01383d215..f9b6bbb65 100644
--- a/routing/payment_lifecycle_test.go
+++ b/routing/payment_lifecycle_test.go
@@ -619,7 +619,7 @@ func TestRouterPaymentStateMachine(t *testing.T) {
 			// A MP payment scenario when our path finding returns
 			// after we've just received a terminal failure, and
 			// demonstrates a bug where the payment will return with
-			// and "unexpected" ErrPaymentAlreadyFailed rather than
+			// and "unexpected" ErrPaymentTerminal rather than
 			// failing with a permanent error. This results in the
 			// payment getting stuck.
 			name: "MP path found after failure",
@@ -648,7 +648,7 @@ func TestRouterPaymentStateMachine(t *testing.T) {
 			routes: []*route.Route{
 				shard, shard,
 			},
-			paymentErr: channeldb.ErrPaymentAlreadyFailed,
+			paymentErr: channeldb.ErrPaymentTerminal,
 		},
 		{
 			// A MP payment scenario when our path finding returns