added whitelisted ips config option

distinct.go

@@ -7,23 +7,26 @@ import (
 	qubic "github.com/qubic/go-node-connector"
 	"log"
 	"math/rand"
+	"net"
 	"sync"
 	"time"
 )
 
 type distinctPeers struct {
 	bp                        *blacklistedPeers
+	whitelistedPeers          map[string]struct{}
 	peers                     map[string]struct{}
 	maxPeers                  int
 	db                        *pebble.DB
 	mux                       sync.RWMutex
 	exchangeConnectionTimeout time.Duration
 }
 
-func newDistinctPeers(startingPeers []string, maxPeers int, exchangeConnectionTimeout time.Duration, bp *blacklistedPeers, db *pebble.DB) *distinctPeers {
+func newDistinctPeers(startingPeers []string, whitelistedPeers []string, maxPeers int, exchangeConnectionTimeout time.Duration, bp *blacklistedPeers, db *pebble.DB) *distinctPeers {
 	dp := distinctPeers{
 		bp:                        bp,
 		peers:                     make(map[string]struct{}, maxPeers),
+		whitelistedPeers:          createWhitelistedPeersMap(whitelistedPeers),
 		maxPeers:                  maxPeers,
 		db:                        db,
 		exchangeConnectionTimeout: exchangeConnectionTimeout,
@@ -33,6 +36,34 @@ func newDistinctPeers(startingPeers []string, maxPeers int, exchangeConnectionTi
 	return &dp
 }
 
+func createWhitelistedPeersMap(peers []string) map[string]struct{} {
+	peersMap := make(map[string]struct{})
+
+	if peers == nil || len(peers) == 0 {
+		return peersMap
+	}
+
+	for _, peer := range peers {
+		ip := net.ParseIP(peer)
+		if ip == nil {
+			continue
+		}
+		peersMap[peer] = struct{}{}
+	}
+
+	return peersMap
+}
+
+// isWhitelisted checks if a peer is whitelisted, if the whitelistedPeers list is empty, it returns true
+func (p *distinctPeers) isWhitelisted(peer string) bool {
+	if len(p.whitelistedPeers) == 0 {
+		return true
+	}
+
+	_, ok := p.whitelistedPeers[peer]
+	return ok
+}
+
 func (p *distinctPeers) build() ([]string, error) {
 	peer, err := p.getRandomPeer()
 	if err != nil {
@@ -114,6 +145,10 @@ func (p *distinctPeers) setPeers(peers []string) {
 	defer p.mux.Unlock()
 
 	for _, peer := range peers {
+		if !p.isWhitelisted(peer) {
+			continue
+		}
+
 		p.peers[peer] = struct{}{}
 	}
 }

main.go

@@ -18,7 +18,6 @@ func main() {
 	if err != nil {
 		log.Fatal(err.Error())
 	}
-
 }
 
 func run() error {
@@ -29,7 +28,8 @@ func run() error {
 			ShutdownTimeout time.Duration `conf:"default:5s"`
 		}
 		Qubic struct {
-			StartingPeerIP  string        `conf:"default:92.50.108.113"`
+			StartingPeerIP  string `conf:"default:95.156.231.18"`
+			WhitelistPeers  []string
 			MaxPeers        int           `conf:"default:50"`
 			ExchangeTimeout time.Duration `conf:"default:2s"`
 			StorageFolder   string        `conf:"default:store"`
@@ -66,7 +66,7 @@ func run() error {
 		log.Fatalf("err opening pebble: %s", err.Error())
 	}
 
-	rp, err := NewPeers(cfg.Qubic.StartingPeerIP, cfg.Qubic.MaxPeers, cfg.Qubic.ExchangeTimeout, db)
+	rp, err := NewPeers(cfg.Qubic.StartingPeerIP, cfg.Qubic.WhitelistPeers, cfg.Qubic.MaxPeers, cfg.Qubic.ExchangeTimeout, db)
 	if err != nil {
 		return errors.Wrap(err, "creating new peers")
 	}

peers.go

@@ -16,7 +16,7 @@ type PeersResponse struct {
 	UpdatedAt int64
 }
 
-func NewPeers(startingPeer string, maxPeers int, exchangeConnectionTimeout time.Duration, db *pebble.DB) (*Peers, error) {
+func NewPeers(startingPeer string, whitelistedPeers []string, maxPeers int, exchangeConnectionTimeout time.Duration, db *pebble.DB) (*Peers, error) {
 	storedPeers, err := retrievePeers(db)
 	if err != nil {
 		return nil, errors.Wrap(err, "retrieving peers from store")
@@ -25,7 +25,7 @@ func NewPeers(startingPeer string, maxPeers int, exchangeConnectionTimeout time.
 	initialPeers := append(storedPeers, startingPeer)
 	bp := newBlacklistedPeers()
 	p := Peers{
-		dp: newDistinctPeers(initialPeers, maxPeers, exchangeConnectionTimeout, bp, db),
+		dp: newDistinctPeers(initialPeers, whitelistedPeers, maxPeers, exchangeConnectionTimeout, bp, db),
 		rp: newReliablePeers(bp),
 	}