From 09b67af48d1ac7fc8fc32a93211c37115064530f Mon Sep 17 00:00:00 2001 From: yyforyongyu Date: Fri, 26 May 2023 06:50:19 +0800 Subject: [PATCH] channeldb: add method `Registrable` to decide adding HTLCs This commit adds a new method, `Registrable`, to help decide whether adding new HTLCs to a given payment is allowed. A new unit test, `TestRegistrable` is also added to test it. --- channeldb/mp_payment.go | 38 +++++++++++++ channeldb/mp_payment_test.go | 94 +++++++++++++++++++++++++++++++ channeldb/payment_control.go | 48 ++++------------ channeldb/payment_control_test.go | 42 +++++++++++--- 4 files changed, 175 insertions(+), 47 deletions(-) diff --git a/channeldb/mp_payment.go b/channeldb/mp_payment.go index 3137be640..6b1a6484e 100644 --- a/channeldb/mp_payment.go +++ b/channeldb/mp_payment.go @@ -235,6 +235,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 { diff --git a/channeldb/mp_payment_test.go b/channeldb/mp_payment_test.go index eb5f0f48c..c16dad8d4 100644 --- a/channeldb/mp_payment_test.go +++ b/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) + }) + } +} diff --git a/channeldb/payment_control.go b/channeldb/payment_control.go index dd7fe218d..e0d7fbcec 100644 --- a/channeldb/payment_control.go +++ b/channeldb/payment_control.go @@ -85,6 +85,14 @@ var ( ErrMPPTotalAmountMismatch = errors.New("mp payment total amount " + "mismatch") + // ErrPaymentPendingSettled is returned when we try to add a new + // attempt to a payment that has at least one of its HTLCs settled. + ErrPaymentPendingSettled = errors.New("payment has settled htlcs") + + // ErrPaymentAlreadyFailed is returned when we try to add a new attempt + // to a payment that already has a failure reason. + ErrPaymentPendingFailed = errors.New("payment has failure reason") + // errNoAttemptInfo is returned when no attempt info is stored yet. errNoAttemptInfo = errors.New("unable to find attempt info for " + "inflight payment") @@ -310,16 +318,8 @@ func (p *PaymentControl) RegisterAttempt(paymentHash lntypes.Hash, return err } - // We cannot register a new attempt if the payment already has - // reached a terminal condition. We check this before - // ensureInFlight because it is a more general check. - settle, fail := payment.TerminalInfo() - if settle != nil || fail != nil { - return ErrPaymentTerminal - } - - // Ensure the payment is in-flight. - if err := ensureInFlight(payment); err != nil { + // Check if registering a new attempt is allowed. + if err := payment.Registrable(); err != nil { return err } @@ -707,34 +707,6 @@ func fetchPaymentStatus(bucket kvdb.RBucket) (PaymentStatus, error) { return payment.Status, nil } -// ensureInFlight checks whether the payment found in the given bucket has -// status InFlight, and returns an error otherwise. This should be used to -// ensure we only mark in-flight payments as succeeded or failed. -func ensureInFlight(payment *MPPayment) error { - paymentStatus := payment.Status - - switch { - // Newly created payment is also inflight. - case paymentStatus == StatusInitiated: - return nil - - // The payment was indeed InFlight. - case paymentStatus == StatusInFlight: - return nil - - // The payment succeeded previously. - case paymentStatus == StatusSucceeded: - return ErrPaymentAlreadySucceeded - - // The payment was already failed. - case paymentStatus == StatusFailed: - return ErrPaymentAlreadyFailed - - default: - return ErrUnknownPaymentStatus - } -} - // FetchInFlightPayments returns all payments with status InFlight. func (p *PaymentControl) FetchInFlightPayments() ([]*MPPayment, error) { var inFlights []*MPPayment diff --git a/channeldb/payment_control_test.go b/channeldb/payment_control_test.go index 61c8e1da1..ffce597f9 100644 --- a/channeldb/payment_control_test.go +++ b/channeldb/payment_control_test.go @@ -818,8 +818,10 @@ func TestPaymentControlMultiShard(t *testing.T) { b = *attempt b.AttemptID = 3 _, err = pControl.RegisterAttempt(info.PaymentIdentifier, &b) - if err != ErrPaymentTerminal { - t.Fatalf("expected ErrPaymentTerminal, got: %v", err) + if test.settleFirst { + require.ErrorIs(t, err, ErrPaymentPendingSettled) + } else { + require.ErrorIs(t, err, ErrPaymentPendingFailed) } assertPaymentStatus(t, pControl, info.PaymentIdentifier, StatusInFlight) @@ -878,19 +880,41 @@ func TestPaymentControlMultiShard(t *testing.T) { } } - // If any of the two attempts settled, the payment should end - // up in the Succeeded state. If both failed the payment should - // also be Failed at this poinnt. - finalStatus := StatusFailed - if test.settleFirst || test.settleLast { + var ( + finalStatus PaymentStatus + registerErr error + ) + + switch { + // If one of the attempts settled but the other failed with + // terminal error, we would still consider the payment is + // settled. + case test.settleFirst && !test.settleLast: finalStatus = StatusSucceeded + registerErr = ErrPaymentAlreadySucceeded + + case !test.settleFirst && test.settleLast: + finalStatus = StatusSucceeded + registerErr = ErrPaymentAlreadySucceeded + + // If both failed, we end up in a failed status. + case !test.settleFirst && !test.settleLast: + finalStatus = StatusFailed + registerErr = ErrPaymentAlreadyFailed + + // Otherwise, the payment has a succeed status. + case test.settleFirst && test.settleLast: + finalStatus = StatusSucceeded + registerErr = ErrPaymentAlreadySucceeded } - assertPaymentStatus(t, pControl, info.PaymentIdentifier, finalStatus) + assertPaymentStatus( + t, pControl, info.PaymentIdentifier, finalStatus, + ) // Finally assert we cannot register more attempts. _, err = pControl.RegisterAttempt(info.PaymentIdentifier, &b) - require.Equal(t, ErrPaymentTerminal, err) + require.Equal(t, registerErr, err) } for _, test := range tests {