From 601afaf1b2d62e751bc9c076c506aa497aaedb20 Mon Sep 17 00:00:00 2001
From: Keagan McClelland <keagan.mcclelland@gmail.com>
Date: Mon, 27 Nov 2023 17:32:03 -0800
Subject: [PATCH] htlcswitch: implement flush api for channelLink

---
 htlcswitch/link.go | 47 +++++++++++++++++++++++++++++++++++++++-------
 1 file changed, 40 insertions(+), 7 deletions(-)

diff --git a/htlcswitch/link.go b/htlcswitch/link.go
index ebcd2031a..97525d8be 100644
--- a/htlcswitch/link.go
+++ b/htlcswitch/link.go
@@ -367,6 +367,14 @@ type channelLink struct {
 	// log is a link-specific logging instance.
 	log btclog.Logger
 
+	// isOutgoingAddBlocked tracks whether the channelLink can send an
+	// UpdateAddHTLC.
+	isOutgoingAddBlocked atomic.Bool
+
+	// isIncomingAddBlocked tracks whether the channelLink can receive an
+	// UpdateAddHTLC.
+	isIncomingAddBlocked atomic.Bool
+
 	wg   sync.WaitGroup
 	quit chan struct{}
 }
@@ -559,24 +567,49 @@ func (l *channelLink) EligibleToUpdate() bool {
 // EnableAdds sets the ChannelUpdateHandler state to allow UpdateAddHtlc's in
 // the specified direction. It returns an error if the state already allowed
 // those adds.
-func (l *channelLink) EnableAdds(LinkDirection) error {
-	// TODO(proofofkeags): Implement
+func (l *channelLink) EnableAdds(linkDirection LinkDirection) error {
+	if linkDirection == Outgoing {
+		if !l.isOutgoingAddBlocked.Swap(false) {
+			return errors.New("outgoing adds already enabled")
+		}
+	}
+
+	if linkDirection == Incoming {
+		if !l.isIncomingAddBlocked.Swap(false) {
+			return errors.New("incoming adds already enabled")
+		}
+	}
+
 	return nil
 }
 
 // DiableAdds sets the ChannelUpdateHandler state to allow UpdateAddHtlc's in
 // the specified direction. It returns an error if the state already disallowed
 // those adds.
-func (l *channelLink) DisableAdds(LinkDirection) error {
-	// TODO(proofofkeags): Implement
+func (l *channelLink) DisableAdds(linkDirection LinkDirection) error {
+	if linkDirection == Outgoing {
+		if l.isOutgoingAddBlocked.Swap(true) {
+			return errors.New("outgoing adds already disabled")
+		}
+	}
+
+	if linkDirection == Incoming {
+		if l.isIncomingAddBlocked.Swap(true) {
+			return errors.New("incoming adds already disabled")
+		}
+	}
+
 	return nil
 }
 
 // IsFlushing returns true when UpdateAddHtlc's are disabled in the direction of
 // the argument.
-func (l *channelLink) IsFlushing(LinkDirection) bool {
-	// TODO(proofofkeags): Implement
-	return false
+func (l *channelLink) IsFlushing(linkDirection LinkDirection) bool {
+	if linkDirection == Outgoing {
+		return l.isOutgoingAddBlocked.Load()
+	}
+
+	return l.isIncomingAddBlocked.Load()
 }
 
 // OnFlushedOnce adds a hook that will be called the next time the