From 105c275b91f30909d8bc165e87d0c982dbc29b36 Mon Sep 17 00:00:00 2001
From: yyforyongyu <yy2452@columbia.edu>
Date: Thu, 24 Nov 2022 17:53:58 +0800
Subject: [PATCH] channeldb: add method `updatable` to decide updating payments

This commit adds a new method, `updatable`, to help decide whether
updating a given payment is allowed.
---
 channeldb/payment_control.go |  2 +-
 channeldb/payment_status.go  | 24 ++++++++++++++++++++++++
 2 files changed, 25 insertions(+), 1 deletion(-)

diff --git a/channeldb/payment_control.go b/channeldb/payment_control.go
index 2d13e2585..fb88a3c94 100644
--- a/channeldb/payment_control.go
+++ b/channeldb/payment_control.go
@@ -451,7 +451,7 @@ func (p *PaymentControl) updateHtlcKey(paymentHash lntypes.Hash,
 		// We can only update keys of in-flight payments. We allow
 		// updating keys even if the payment has reached a terminal
 		// condition, since the HTLC outcomes must still be updated.
-		if err := ensureInFlight(p); err != nil {
+		if err := p.Status.updatable(); err != nil {
 			return err
 		}
 
diff --git a/channeldb/payment_status.go b/channeldb/payment_status.go
index 17ffd3e88..a6c7c1038 100644
--- a/channeldb/payment_status.go
+++ b/channeldb/payment_status.go
@@ -103,3 +103,27 @@ func (ps PaymentStatus) removable() error {
 		return fmt.Errorf("%w: %v", ErrUnknownPaymentStatus, ps)
 	}
 }
+
+// updatable returns an error to specify whether the payment's HTLCs can be
+// updated. A payment can update its HTLCs when it has inflight HTLCs.
+func (ps PaymentStatus) updatable() error {
+	switch ps {
+	// Newly created payments can be updated.
+	case StatusInitiated:
+		return nil
+
+	// Inflight payments can be updated.
+	case StatusInFlight:
+		return nil
+
+	// If the payment has a terminal condition, we won't allow any updates.
+	case StatusSucceeded:
+		return ErrPaymentAlreadySucceeded
+
+	case StatusFailed:
+		return ErrPaymentAlreadyFailed
+
+	default:
+		return fmt.Errorf("%w: %v", ErrUnknownPaymentStatus, ps)
+	}
+}