Merge pull request lightningnetwork#7175 from yyforyongyu/refactor-pa…

…yment-status

channeldb: properly decide and handle payment status

channeldb/duplicate_payments.go

@@ -59,22 +59,22 @@ type duplicateHTLCAttemptInfo struct {
 	route route.Route
 }
 
-// fetchDuplicatePaymentStatus fetches the payment status of the payment. If the
-// payment isn't found, it will default to "StatusUnknown".
-func fetchDuplicatePaymentStatus(bucket kvdb.RBucket) PaymentStatus {
+// fetchDuplicatePaymentStatus fetches the payment status of the payment. If
+// the payment isn't found, it will return error `ErrPaymentNotInitiated`.
+func fetchDuplicatePaymentStatus(bucket kvdb.RBucket) (PaymentStatus, error) {
 	if bucket.Get(duplicatePaymentSettleInfoKey) != nil {
-		return StatusSucceeded
+		return StatusSucceeded, nil
 	}
 
 	if bucket.Get(duplicatePaymentFailInfoKey) != nil {
-		return StatusFailed
+		return StatusFailed, nil
 	}
 
 	if bucket.Get(duplicatePaymentCreationInfoKey) != nil {
-		return StatusInFlight
+		return StatusInFlight, nil
 	}
 
-	return StatusUnknown
+	return 0, ErrPaymentNotInitiated
 }
 
 func deserializeDuplicateHTLCAttemptInfo(r io.Reader) (
@@ -138,7 +138,10 @@ func fetchDuplicatePayment(bucket kvdb.RBucket) (*MPPayment, error) {
 	sequenceNum := binary.BigEndian.Uint64(seqBytes)
 
 	// Get the payment status.
-	paymentStatus := fetchDuplicatePaymentStatus(bucket)
+	paymentStatus, err := fetchDuplicatePaymentStatus(bucket)
+	if err != nil {
+		return nil, err
+	}
 
 	// Get the PaymentCreationInfo.
 	b := bucket.Get(duplicatePaymentCreationInfoKey)

channeldb/mp_payment.go

@@ -177,6 +177,13 @@ type MPPayment struct {
 	Status PaymentStatus
 }
 
+// Terminated returns a bool to specify whether the payment is in a terminal
+// state.
+func (m *MPPayment) Terminated() bool {
+	// If the payment is in terminal state, it cannot be updated.
+	return m.Status.updatable() != nil
+}
+
 // TerminalInfo returns any HTLC settle info recorded. If no settle info is
 // recorded, any payment level failure will be returned. If neither a settle
 // nor a failure is recorded, both return values will be nil.
@@ -235,6 +242,44 @@ func (m *MPPayment) GetAttempt(id uint64) (*HTLCAttempt, error) {
 	return nil, errors.New("htlc attempt not found on payment")
 }
 
+// Registrable returns an error to specify whether adding more HTLCs to the
+// payment with its current status is allowed. A payment can accept new HTLC
+// registrations when it's newly created, or none of its HTLCs is in a terminal
+// state.
+func (m *MPPayment) Registrable() error {
+	// Get the terminal info.
+	settle, reason := m.TerminalInfo()
+	settled := settle != nil
+	failed := reason != nil
+
+	// If updating the payment is not allowed, we can't register new HTLCs.
+	// Otherwise, the status must be either `StatusInitiated` or
+	// `StatusInFlight`.
+	if err := m.Status.updatable(); err != nil {
+		return err
+	}
+
+	// Exit early if this is not inflight.
+	if m.Status != StatusInFlight {
+		return nil
+	}
+
+	// There are still inflight HTLCs and we need to check whether there
+	// are settled HTLCs or the payment is failed. If we already have
+	// settled HTLCs, we won't allow adding more HTLCs.
+	if settled {
+		return ErrPaymentPendingSettled
+	}
+
+	// If the payment is already failed, we won't allow adding more HTLCs.
+	if failed {
+		return ErrPaymentPendingFailed
+	}
+
+	// Otherwise we can add more HTLCs.
+	return nil
+}
+
 // serializeHTLCSettleInfo serializes the details of a settled htlc.
 func serializeHTLCSettleInfo(w io.Writer, s *HTLCSettleInfo) error {
 	if _, err := w.Write(s.Preimage[:]); err != nil {

channeldb/mp_payment_test.go

@@ -2,6 +2,7 @@ package channeldb
 
 import (
 	"bytes"
+	"fmt"
 	"testing"
 
 	"github.com/stretchr/testify/require"
@@ -25,3 +26,96 @@ func TestLazySessionKeyDeserialize(t *testing.T) {
 	sessionKey := attempt.SessionKey()
 	require.Equal(t, priv, sessionKey)
 }
+
+// TestRegistrable checks the method `Registrable` behaves as expected for ALL
+// possible payment statuses.
+func TestRegistrable(t *testing.T) {
+	t.Parallel()
+
+	testCases := []struct {
+		status         PaymentStatus
+		registryErr    error
+		hasSettledHTLC bool
+		paymentFailed  bool
+	}{
+		{
+			status:      StatusInitiated,
+			registryErr: nil,
+		},
+		{
+			// Test inflight status with no settled HTLC and no
+			// failed payment.
+			status:      StatusInFlight,
+			registryErr: nil,
+		},
+		{
+			// Test inflight status with settled HTLC but no failed
+			// payment.
+			status:         StatusInFlight,
+			registryErr:    ErrPaymentPendingSettled,
+			hasSettledHTLC: true,
+		},
+		{
+			// Test inflight status with no settled HTLC but failed
+			// payment.
+			status:        StatusInFlight,
+			registryErr:   ErrPaymentPendingFailed,
+			paymentFailed: true,
+		},
+		{
+			// Test error state with settled HTLC and failed
+			// payment.
+			status:         0,
+			registryErr:    ErrUnknownPaymentStatus,
+			hasSettledHTLC: true,
+			paymentFailed:  true,
+		},
+		{
+			status:      StatusSucceeded,
+			registryErr: ErrPaymentAlreadySucceeded,
+		},
+		{
+			status:      StatusFailed,
+			registryErr: ErrPaymentAlreadyFailed,
+		},
+		{
+			status:      0,
+			registryErr: ErrUnknownPaymentStatus,
+		},
+	}
+
+	// Create test objects.
+	reason := FailureReasonError
+	htlcSettled := HTLCAttempt{
+		Settle: &HTLCSettleInfo{},
+	}
+
+	for i, tc := range testCases {
+		i, tc := i, tc
+
+		p := &MPPayment{
+			Status: tc.status,
+		}
+
+		// Add the settled htlc to the payment if needed.
+		htlcs := make([]HTLCAttempt, 0)
+		if tc.hasSettledHTLC {
+			htlcs = append(htlcs, htlcSettled)
+		}
+		p.HTLCs = htlcs
+
+		// Add the failure reason if needed.
+		if tc.paymentFailed {
+			p.FailureReason = &reason
+		}
+
+		name := fmt.Sprintf("test_%d_%s", i, p.Status.String())
+		t.Run(name, func(t *testing.T) {
+			t.Parallel()
+
+			err := p.Registrable()
+			require.ErrorIs(t, err, tc.registryErr,
+				"registrable under state %v", tc.status)
+		})
+	}
+}