From d92e193e63929717c1c90c868821b88b59e202b2 Mon Sep 17 00:00:00 2001
From: Roman Khimov <roman@nspcc.ru>
Date: Fri, 10 Jan 2020 15:13:29 +0300
Subject: [PATCH] rpc/network: refactor getpeers logic

Deduplicate and simplify code.
---
 pkg/network/server.go   | 13 +++++++++
 pkg/rpc/result/peers.go | 64 +++++++++++++++++++++--------------------
 pkg/rpc/server.go       | 17 +++--------
 3 files changed, 50 insertions(+), 44 deletions(-)

diff --git a/pkg/network/server.go b/pkg/network/server.go
index dacac6461..3772a6018 100644
--- a/pkg/network/server.go
+++ b/pkg/network/server.go
@@ -198,6 +198,19 @@ func (s *Server) BadPeers() []string {
 	return []string{}
 }
 
+// ConnectedPeers returns a list of currently connected peers.
+func (s *Server) ConnectedPeers() []string {
+	s.lock.RLock()
+	defer s.lock.RUnlock()
+
+	peers := make([]string, 0, len(s.peers))
+	for k := range s.peers {
+		peers = append(peers, k.PeerAddr().String())
+	}
+
+	return peers
+}
+
 // run is a goroutine that starts another goroutine to manage protocol specifics
 // while itself dealing with peers management (handling connects/disconnects).
 func (s *Server) run() {
diff --git a/pkg/rpc/result/peers.go b/pkg/rpc/result/peers.go
index 98442acaa..b18acc0e8 100644
--- a/pkg/rpc/result/peers.go
+++ b/pkg/rpc/result/peers.go
@@ -5,13 +5,16 @@ import (
 )
 
 type (
-	// Peers payload for outputting peers in `getpeers` RPC call.
-	Peers struct {
-		Unconnected []Peer `json:"unconnected"`
-		Connected   []Peer `json:"connected"`
-		Bad         []Peer `json:"bad"`
+	// GetPeers payload for outputting peers in `getpeers` RPC call.
+	GetPeers struct {
+		Unconnected Peers `json:"unconnected"`
+		Connected   Peers `json:"connected"`
+		Bad         Peers `json:"bad"`
 	}
 
+	// Peers represent a slice of peers.
+	Peers []Peer
+
 	// Peer represents the peer.
 	Peer struct {
 		Address string `json:"address"`
@@ -19,40 +22,39 @@ type (
 	}
 )
 
-// NewPeers creates a new Peers struct.
-func NewPeers() Peers {
-	return Peers{
+// NewGetPeers creates a new GetPeers structure.
+func NewGetPeers() GetPeers {
+	return GetPeers{
 		Unconnected: []Peer{},
 		Connected:   []Peer{},
 		Bad:         []Peer{},
 	}
 }
 
-// AddPeer adds a peer to the given peer type slice.
-func (p *Peers) AddPeer(peerType string, addr string) {
-	addressParts := strings.Split(addr, ":")
-	peer := Peer{
-		Address: addressParts[0],
-		Port:    addressParts[1],
-	}
+// AddUnconnected adds a set of peers to the unconnected peers slice.
+func (g *GetPeers) AddUnconnected(addrs []string) {
+	g.Unconnected.addPeers(addrs)
+}
 
-	switch peerType {
-	case "unconnected":
-		p.Unconnected = append(
-			p.Unconnected,
-			peer,
-		)
+// AddConnected adds a set of peers to the connected peers slice.
+func (g *GetPeers) AddConnected(addrs []string) {
+	g.Connected.addPeers(addrs)
+}
 
-	case "connected":
-		p.Connected = append(
-			p.Connected,
-			peer,
-		)
+// AddBad adds a set of peers to the bad peers slice.
+func (g *GetPeers) AddBad(addrs []string) {
+	g.Bad.addPeers(addrs)
+}
 
-	case "bad":
-		p.Bad = append(
-			p.Bad,
-			peer,
-		)
+// addPeers adds a set of peers to the given peer slice.
+func (p *Peers) addPeers(addrs []string) {
+	for i := range addrs {
+		addressParts := strings.Split(addrs[i], ":")
+		peer := Peer{
+			Address: addressParts[0],
+			Port:    addressParts[1],
+		}
+
+		*p = append(*p, peer)
 	}
 }
diff --git a/pkg/rpc/server.go b/pkg/rpc/server.go
index 183019e67..1ebb345b6 100644
--- a/pkg/rpc/server.go
+++ b/pkg/rpc/server.go
@@ -192,19 +192,10 @@ Methods:
 
 	case "getpeers":
 		getpeersCalled.Inc()
-		peers := result.NewPeers()
-		for _, addr := range s.coreServer.UnconnectedPeers() {
-			peers.AddPeer("unconnected", addr)
-		}
-
-		for _, addr := range s.coreServer.BadPeers() {
-			peers.AddPeer("bad", addr)
-		}
-
-		for addr := range s.coreServer.Peers() {
-			peers.AddPeer("connected", addr.PeerAddr().String())
-		}
-
+		peers := result.NewGetPeers()
+		peers.AddUnconnected(s.coreServer.UnconnectedPeers())
+		peers.AddConnected(s.coreServer.ConnectedPeers())
+		peers.AddBad(s.coreServer.BadPeers())
 		results = peers
 
 	case "getstorage":