From 502e47362cf0703dba005363d29483ec158d0b1b Mon Sep 17 00:00:00 2001
From: carla <kirkcohenc@gmail.com>
Date: Wed, 24 Mar 2021 15:13:31 +0200
Subject: [PATCH] liquidity: add peers to swap interface and include in set
 reason

This commit adds a peer listing function to our generic swap
interface, which we will use to set reasons for swaps that are
specified by peer pubkey rather than channel.
---
 liquidity/interface.go      | 33 +++++++++++++++++++++++++++++++++
 liquidity/liquidity.go      | 14 ++++++++++++++
 liquidity/liquidity_test.go |  7 +++----
 3 files changed, 50 insertions(+), 4 deletions(-)

diff --git a/liquidity/interface.go b/liquidity/interface.go
index 4c695ad..c05091c 100644
--- a/liquidity/interface.go
+++ b/liquidity/interface.go
@@ -5,6 +5,7 @@ import (
 	"github.com/lightninglabs/loop"
 	"github.com/lightningnetwork/lnd/lnwallet/chainfee"
 	"github.com/lightningnetwork/lnd/lnwire"
+	"github.com/lightningnetwork/lnd/routing/route"
 )
 
 // FeeLimit is an interface implemented by different strategies for limiting
@@ -43,8 +44,17 @@ type swapSuggestion interface {
 
 	// channels returns the set of channels involved in the swap.
 	channels() []lnwire.ShortChannelID
+
+	// peers returns the set of peers involved in the swap, taking a map
+	// of known channel IDs to peers as an argument so that channel peers
+	// can be looked up.
+	peers(knownChans map[uint64]route.Vertex) []route.Vertex
 }
 
+// Compile-time assertion that loopOutSwapSuggestion satisfies the
+// swapSuggestion interface.
+var _ swapSuggestion = (*loopOutSwapSuggestion)(nil)
+
 type loopOutSwapSuggestion struct {
 	loop.OutRequest
 }
@@ -69,3 +79,26 @@ func (l *loopOutSwapSuggestion) channels() []lnwire.ShortChannelID {
 
 	return channels
 }
+
+// peers returns the set of peers that the loop out swap is restricted to.
+func (l *loopOutSwapSuggestion) peers(
+	knownChans map[uint64]route.Vertex) []route.Vertex {
+
+	peers := make(map[route.Vertex]struct{}, len(knownChans))
+
+	for _, channel := range l.OutgoingChanSet {
+		peer, ok := knownChans[channel]
+		if !ok {
+			log.Warnf("peer for channel: %v unknown", channel)
+		}
+
+		peers[peer] = struct{}{}
+	}
+
+	peerList := make([]route.Vertex, 0, len(peers))
+	for peer := range peers {
+		peerList = append(peerList, peer)
+	}
+
+	return peerList
+}
diff --git a/liquidity/liquidity.go b/liquidity/liquidity.go
index 91a6fb7..6189ed6 100644
--- a/liquidity/liquidity.go
+++ b/liquidity/liquidity.go
@@ -685,8 +685,13 @@ func (m *Manager) SuggestSwaps(ctx context.Context, autoloop bool) (
 		return nil, err
 	}
 
+	// Collect a map of channel IDs to peer pubkeys, and a set of per-peer
+	// balances which we will use for peer-level liquidity rules.
+	channelPeers := make(map[uint64]route.Vertex)
 	peerChannels := make(map[route.Vertex]*balances)
 	for _, channel := range channels {
+		channelPeers[channel.ChannelID] = channel.PubKeyBytes
+
 		bal, ok := peerChannels[channel.PubKeyBytes]
 		if !ok {
 			bal = &balances{}
@@ -777,6 +782,15 @@ func (m *Manager) SuggestSwaps(ctx context.Context, autoloop bool) (
 	// setReason is a helper that adds a swap's channels to our disqualified
 	// list with the reason provided.
 	setReason := func(reason Reason, swap swapSuggestion) {
+		for _, peer := range swap.peers(channelPeers) {
+			_, ok := m.params.PeerRules[peer]
+			if !ok {
+				continue
+			}
+
+			resp.DisqualifiedPeers[peer] = reason
+		}
+
 		for _, channel := range swap.channels() {
 			_, ok := m.params.ChannelRules[channel]
 			if !ok {
diff --git a/liquidity/liquidity_test.go b/liquidity/liquidity_test.go
index 34ea77b..ccded3a 100644
--- a/liquidity/liquidity_test.go
+++ b/liquidity/liquidity_test.go
@@ -1198,10 +1198,9 @@ func TestInFlightLimit(t *testing.T) {
 					chan1Rec,
 				},
 				DisqualifiedChans: noneDisqualified,
-				// This is a bug, we should list our second
-				// peer as disqualified because we've hit our
-				// in-flight limit for now.
-				DisqualifiedPeers: noPeersDisqualified,
+				DisqualifiedPeers: map[route.Vertex]Reason{
+					peer2: ReasonInFlight,
+				},
 			},
 		},
 	}