From 4f4854880939928e9e1fc469ab2cd3fa3be78948 Mon Sep 17 00:00:00 2001
From: Viacheslav Gonkivskyi <viacheslavgonkivskyi@gmail.com>
Date: Tue, 2 Apr 2024 20:37:20 +0300
Subject: [PATCH] fix(p2p/session): return err if peer tracker is empty

---
 p2p/exchange.go     |  6 +++++-
 p2p/session.go      | 11 ++++++++---
 p2p/session_test.go | 28 ++++++++++++++++++++--------
 3 files changed, 33 insertions(+), 12 deletions(-)

diff --git a/p2p/exchange.go b/p2p/exchange.go
index 0681aeff..de503119 100644
--- a/p2p/exchange.go
+++ b/p2p/exchange.go
@@ -292,9 +292,13 @@ func (ex *Exchange[H]) GetRangeByHeight(
 			attribute.Int64("to", int64(to)),
 		))
 	defer span.End()
-	session := newSession[H](
+	session, err := newSession[H](
 		ex.ctx, ex.host, ex.peerTracker, ex.protocolID, ex.Params.RangeRequestTimeout, ex.metrics, withValidation(from),
 	)
+	// TODO(@vgonkivs): decide what to do with this error. Maybe we should fall into "discovery mode" and try to collect peers???
+	if err != nil {
+		return nil, err
+	}
 	defer session.close()
 	// we request the next header height that we don't have: `fromHead`+1
 	amount := to - (from.Height() + 1)
diff --git a/p2p/session.go b/p2p/session.go
index 6fc24f70..44559d77 100644
--- a/p2p/session.go
+++ b/p2p/session.go
@@ -62,23 +62,28 @@ func newSession[H header.Header[H]](
 	requestTimeout time.Duration,
 	metrics *exchangeMetrics,
 	options ...option[H],
-) *session[H] {
+) (*session[H], error) {
 	ctx, cancel := context.WithCancel(ctx)
 	ses := &session[H]{
 		ctx:            ctx,
 		cancel:         cancel,
 		protocolID:     protocolID,
 		host:           h,
-		queue:          newPeerQueue(ctx, peerTracker.peers()),
 		peerTracker:    peerTracker,
 		requestTimeout: requestTimeout,
 		metrics:        metrics,
 	}
 
+	peers := peerTracker.peers()
+	if len(peers) == 0 {
+		return nil, errors.New("empty peer tracker")
+	}
+	ses.queue = newPeerQueue(ctx, peers)
+
 	for _, opt := range options {
 		opt(ses)
 	}
-	return ses
+	return ses, nil
 }
 
 // getRangeByHeight requests headers from different peers.
diff --git a/p2p/session_test.go b/p2p/session_test.go
index 7c8599f5..8aa1776a 100644
--- a/p2p/session_test.go
+++ b/p2p/session_test.go
@@ -6,6 +6,8 @@ import (
 	"time"
 
 	"github.com/libp2p/go-libp2p/core/peer"
+	blankhost "github.com/libp2p/go-libp2p/p2p/host/blank"
+	swarm "github.com/libp2p/go-libp2p/p2p/net/swarm/testing"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
@@ -25,16 +27,21 @@ func Test_PrepareRequests(t *testing.T) {
 func Test_Validate(t *testing.T) {
 	suite := headertest.NewTestSuite(t)
 	head := suite.Head()
-	ses := newSession(
+	peerId := peer.ID("test")
+	pT := &peerTracker{trackedPeers: make(map[peer.ID]struct{})}
+	pT.trackedPeers[peerId] = struct{}{}
+	pT.host = blankhost.NewBlankHost(swarm.GenSwarm(t))
+	ses, err := newSession(
 		context.Background(),
 		nil,
-		&peerTracker{trackedPeers: make(map[peer.ID]struct{})},
+		pT,
 		"", time.Second, nil,
 		withValidation(head),
 	)
 
+	require.NoError(t, err)
 	headers := suite.GenDummyHeaders(5)
-	err := ses.verify(headers)
+	err = ses.verify(headers)
 	assert.NoError(t, err)
 }
 
@@ -42,17 +49,22 @@ func Test_Validate(t *testing.T) {
 func Test_ValidateFails(t *testing.T) {
 	suite := headertest.NewTestSuite(t)
 	head := suite.Head()
-	ses := newSession(
+
+	peerId := peer.ID("test")
+	pT := &peerTracker{trackedPeers: make(map[peer.ID]struct{})}
+	pT.trackedPeers[peerId] = struct{}{}
+	pT.host = blankhost.NewBlankHost(swarm.GenSwarm(t))
+	ses, err := newSession(
 		context.Background(),
-		nil,
-		&peerTracker{trackedPeers: make(map[peer.ID]struct{})},
+		blankhost.NewBlankHost(swarm.GenSwarm(t)),
+		pT,
 		"", time.Second, nil,
 		withValidation(head),
 	)
-
+	require.NoError(t, err)
 	headers := suite.GenDummyHeaders(5)
 	// break adjacency
 	headers[2] = headers[4]
-	err := ses.verify(headers)
+	err = ses.verify(headers)
 	assert.Error(t, err)
 }