From 7f5f57e9322371a859c1f5fbf6f2cf9fc89447c9 Mon Sep 17 00:00:00 2001
From: George Tsagkarelis <george.tsagkarelis@gmail.com>
Date: Mon, 11 Nov 2024 14:31:15 +0100
Subject: [PATCH] rfq: policies track accepted htlcs

---
 rfq/order.go | 45 +++++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 43 insertions(+), 2 deletions(-)

diff --git a/rfq/order.go b/rfq/order.go
index 7d29e1709..8cbc0b8ad 100644
--- a/rfq/order.go
+++ b/rfq/order.go
@@ -70,6 +70,11 @@ type Policy interface {
 	// which the policy applies.
 	Scid() uint64
 
+	// TrackAcceptedHtlc makes the policy aware of this new accepted HTLC.
+	// This is important in cases where the set of existing HTLCs may affect
+	// whether the next compliance check passes.
+	TrackAcceptedHtlc(htlc lndclient.InterceptedHtlc)
+
 	// GenerateInterceptorResponse generates an interceptor response for the
 	// HTLC interceptor from the policy.
 	GenerateInterceptorResponse(
@@ -94,6 +99,10 @@ type AssetSalePolicy struct {
 	// the policy.
 	MaxOutboundAssetAmount uint64
 
+	// CurrentAssetAmountMsat is the total amount that is held currently in
+	// accepted htlcs.
+	CurrentAmountMsat lnwire.MilliSatoshi
+
 	// AskAssetRate is the quote's asking asset unit to BTC conversion rate.
 	AskAssetRate rfqmath.BigIntFixedPoint
 
@@ -151,7 +160,7 @@ func (c *AssetSalePolicy) CheckHtlcCompliance(
 		maxAssetAmount, c.AskAssetRate,
 	)
 
-	if htlc.AmountOutMsat > policyMaxOutMsat {
+	if (c.CurrentAmountMsat + htlc.AmountOutMsat) > policyMaxOutMsat {
 		return fmt.Errorf("htlc out amount is greater than the policy "+
 			"maximum (htlc_out_msat=%d, policy_max_out_msat=%d)",
 			htlc.AmountOutMsat, policyMaxOutMsat)
@@ -166,6 +175,12 @@ func (c *AssetSalePolicy) CheckHtlcCompliance(
 	return nil
 }
 
+// TrackAcceptedHtlc accounts for the newly accepted htlc. This may affect the
+// acceptance of future htlcs.
+func (c *AssetSalePolicy) TrackAcceptedHtlc(htlc lndclient.InterceptedHtlc) {
+	c.CurrentAmountMsat += htlc.AmountOutMsat
+}
+
 // Expiry returns the policy's expiry time as a unix timestamp.
 func (c *AssetSalePolicy) Expiry() uint64 {
 	return c.expiry
@@ -245,6 +260,10 @@ type AssetPurchasePolicy struct {
 	// AcceptedQuoteId is the ID of the accepted quote.
 	AcceptedQuoteId rfqmsg.ID
 
+	// CurrentAssetAmountMsat is the total amount that is held currently in
+	// accepted htlcs.
+	CurrentAmountMsat lnwire.MilliSatoshi
+
 	// BidAssetRate is the quote's asset to BTC conversion rate.
 	BidAssetRate rfqmath.BigIntFixedPoint
 
@@ -314,7 +333,7 @@ func (c *AssetPurchasePolicy) CheckHtlcCompliance(
 
 	// Ensure that the outbound HTLC amount is less than the maximum agreed
 	// BTC payment.
-	if htlc.AmountOutMsat > c.PaymentMaxAmt {
+	if (c.CurrentAmountMsat + htlc.AmountOutMsat) > c.PaymentMaxAmt {
 		return fmt.Errorf("htlc out amount is more than the maximum "+
 			"agreed BTC payment (htlc_out_msat=%d, "+
 			"payment_max_amt=%d)", htlc.AmountOutMsat,
@@ -330,6 +349,14 @@ func (c *AssetPurchasePolicy) CheckHtlcCompliance(
 	return nil
 }
 
+// TrackAcceptedHtlc accounts for the newly accepted htlc. This may affect the
+// acceptance of future htlcs.
+func (c *AssetPurchasePolicy) TrackAcceptedHtlc(
+	htlc lndclient.InterceptedHtlc) {
+
+	c.CurrentAmountMsat += htlc.AmountOutMsat
+}
+
 // Expiry returns the policy's expiry time as a unix timestamp in seconds.
 func (c *AssetPurchasePolicy) Expiry() uint64 {
 	return c.expiry
@@ -428,6 +455,16 @@ func (a *AssetForwardPolicy) CheckHtlcCompliance(
 	return nil
 }
 
+// TrackAcceptedHtlc accounts for the newly accepted htlc. This may affect the
+// acceptance of future htlcs.
+func (a *AssetForwardPolicy) TrackAcceptedHtlc(htlc lndclient.InterceptedHtlc) {
+	// Track accepted htlc in the incoming policy.
+	a.incomingPolicy.TrackAcceptedHtlc(htlc)
+
+	// Track accepted htlc in the outgoing policy.
+	a.outgoingPolicy.TrackAcceptedHtlc(htlc)
+}
+
 // Expiry returns the policy's expiry time as a unix timestamp in seconds. The
 // returned expiry time is the earliest expiry time of the incoming and outgoing
 // policies.
@@ -585,6 +622,10 @@ func (h *OrderHandler) handleIncomingHtlc(_ context.Context,
 		}, nil
 	}
 
+	// The htlc passed the compliance checks, so now we keep track of the
+	// accepted htlc.
+	policy.TrackAcceptedHtlc(htlc)
+
 	log.Debug("HTLC complies with policy. Broadcasting accept event.")
 	h.cfg.AcceptHtlcEvents <- NewAcceptHtlcEvent(htlc, policy)