From bb6aca11300abda156bdf16cbb881bb2f427bf69 Mon Sep 17 00:00:00 2001
From: Olaoluwa Osuntokun <laolu32@gmail.com>
Date: Fri, 9 Apr 2021 15:58:15 -0700
Subject: [PATCH] discovery: add unit tests for chan anns in the reject cache

---
 discovery/gossiper_test.go | 84 ++++++++++++++++++++++++++++++++++----
 1 file changed, 75 insertions(+), 9 deletions(-)

diff --git a/discovery/gossiper_test.go b/discovery/gossiper_test.go
index f32860b9..a8b3154a 100644
--- a/discovery/gossiper_test.go
+++ b/discovery/gossiper_test.go
@@ -100,19 +100,21 @@ func makeTestDB() (*channeldb.DB, func(), error) {
 type mockGraphSource struct {
 	bestHeight uint32
 
-	mu      sync.Mutex
-	nodes   []channeldb.LightningNode
-	infos   map[uint64]channeldb.ChannelEdgeInfo
-	edges   map[uint64][]channeldb.ChannelEdgePolicy
-	zombies map[uint64][][33]byte
+	mu            sync.Mutex
+	nodes         []channeldb.LightningNode
+	infos         map[uint64]channeldb.ChannelEdgeInfo
+	edges         map[uint64][]channeldb.ChannelEdgePolicy
+	zombies       map[uint64][][33]byte
+	chansToReject map[uint64]struct{}
 }
 
 func newMockRouter(height uint32) *mockGraphSource {
 	return &mockGraphSource{
-		bestHeight: height,
-		infos:      make(map[uint64]channeldb.ChannelEdgeInfo),
-		edges:      make(map[uint64][]channeldb.ChannelEdgePolicy),
-		zombies:    make(map[uint64][][33]byte),
+		bestHeight:    height,
+		infos:         make(map[uint64]channeldb.ChannelEdgeInfo),
+		edges:         make(map[uint64][]channeldb.ChannelEdgePolicy),
+		zombies:       make(map[uint64][][33]byte),
+		chansToReject: make(map[uint64]struct{}),
 	}
 }
 
@@ -138,10 +140,21 @@ func (r *mockGraphSource) AddEdge(info *channeldb.ChannelEdgeInfo,
 		return errors.New("info already exist")
 	}
 
+	if _, ok := r.chansToReject[info.ChannelID]; ok {
+		return errors.New("validation failed")
+	}
+
 	r.infos[info.ChannelID] = *info
 	return nil
 }
 
+func (r *mockGraphSource) queueValidationFail(chanID uint64) {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+
+	r.chansToReject[chanID] = struct{}{}
+}
+
 func (r *mockGraphSource) UpdateEdge(edge *channeldb.ChannelEdgePolicy,
 	_ ...batch.SchedulerOption) error {
 
@@ -4177,3 +4190,56 @@ func TestIgnoreOwnAnnouncement(t *testing.T) {
 		t.Fatalf("expected gossiper to ignore announcement, got: %v", err)
 	}
 }
+
+// TestRejectCacheChannelAnn checks that if we reject a channel announcement,
+// then if we attempt to validate it again, we'll reject it with the proper
+// error.
+func TestRejectCacheChannelAnn(t *testing.T) {
+	t.Parallel()
+
+	ctx, cleanup, err := createTestCtx(proofMatureDelta)
+	if err != nil {
+		t.Fatalf("can't create context: %v", err)
+	}
+	defer cleanup()
+
+	// First, we create a channel announcement to send over to our test
+	// peer.
+	batch, err := createRemoteAnnouncements(0)
+	if err != nil {
+		t.Fatalf("can't generate announcements: %v", err)
+	}
+
+	remoteKey, err := btcec.ParsePubKey(batch.nodeAnn2.NodeID[:], btcec.S256())
+	if err != nil {
+		t.Fatalf("unable to parse pubkey: %v", err)
+	}
+	remotePeer := &mockPeer{remoteKey, nil, nil}
+
+	// Before sending over the announcement, we'll modify it such that we
+	// know it will always fail.
+	chanID := batch.chanAnn.ShortChannelID.ToUint64()
+	ctx.router.queueValidationFail(chanID)
+
+	// If we process the batch the first time we should get an error.
+	select {
+	case err = <-ctx.gossiper.ProcessRemoteAnnouncement(
+		batch.chanAnn, remotePeer,
+	):
+		require.NotNil(t, err)
+	case <-time.After(2 * time.Second):
+		t.Fatal("did not process remote announcement")
+	}
+
+	// If we process it a *second* time, then we should get an error saying
+	// we rejected it already.
+	select {
+	case err = <-ctx.gossiper.ProcessRemoteAnnouncement(
+		batch.chanAnn, remotePeer,
+	):
+		errStr := err.Error()
+		require.Contains(t, errStr, "recently rejected")
+	case <-time.After(2 * time.Second):
+		t.Fatal("did not process remote announcement")
+	}
+}