From 0e22ae09bd490fd1d22e6d030dcba0d93fce7441 Mon Sep 17 00:00:00 2001
From: anthdm <anthony@academiclabs.co>
Date: Thu, 1 Feb 2018 08:19:29 +0100
Subject: [PATCH] added peerCount.

---
 pkg/network/server.go      | 72 +++++++++++++++++++++++++++-----------
 pkg/network/server_test.go | 13 +++++++
 2 files changed, 65 insertions(+), 20 deletions(-)

diff --git a/pkg/network/server.go b/pkg/network/server.go
index 7e73c57d7..c18c11695 100644
--- a/pkg/network/server.go
+++ b/pkg/network/server.go
@@ -5,7 +5,6 @@ import (
 	"log"
 	"net"
 	"os"
-	"strconv"
 	"time"
 
 	"github.com/anthdm/neo-go/pkg/network/payload"
@@ -53,6 +52,8 @@ type Server struct {
 	relay bool
 	// TCP listener of the server
 	listener net.Listener
+	// channel for safely responding the number of current connected peers.
+	peerCountCh chan peerCount
 }
 
 // NewServer returns a pointer to a new server.
@@ -64,28 +65,26 @@ func NewServer(net NetMode) *Server {
 	}
 
 	s := &Server{
-		id:         util.RandUint32(1111111, 9999999),
-		userAgent:  fmt.Sprintf("/NEO:%s/", version),
-		logger:     logger,
-		peers:      make(map[Peer]bool),
-		register:   make(chan Peer),
-		unregister: make(chan Peer),
-		message:    make(chan messageTuple),
-		relay:      true, // currently relay is not handled.
-		net:        net,
-		quit:       make(chan struct{}),
+		id:          util.RandUint32(1111111, 9999999),
+		userAgent:   fmt.Sprintf("/NEO:%s/", version),
+		logger:      logger,
+		peers:       make(map[Peer]bool),
+		register:    make(chan Peer),
+		unregister:  make(chan Peer),
+		message:     make(chan messageTuple),
+		relay:       true, // currently relay is not handled.
+		net:         net,
+		quit:        make(chan struct{}),
+		peerCountCh: make(chan peerCount),
 	}
 
 	return s
 }
 
 // Start run's the server.
-func (s *Server) Start(port string, seeds []string) {
-	p, err := strconv.Atoi(port[1:len(port)])
-	if err != nil {
-		s.logger.Fatalf("could not convert port to integer: %s", err)
-	}
-	s.port = uint16(p)
+// TODO: server should be initialized with a config.
+func (s *Server) Start(opts StartOpts) {
+	s.port = uint16(opts.TCP)
 
 	fmt.Println(logo())
 	fmt.Println(string(s.userAgent))
@@ -93,10 +92,14 @@ func (s *Server) Start(port string, seeds []string) {
 	s.logger.Printf("NET: %s - TCP: %d - RELAY: %v - ID: %d",
 		s.net, int(s.port), s.relay, s.id)
 
-	go listenTCP(s, port)
+	go listenTCP(s, opts.TCP)
 
-	if len(seeds) > 0 {
-		connectToSeeds(s, seeds)
+	if opts.RPC > 0 {
+		go listenHTTP(s, opts.RPC)
+	}
+
+	if len(opts.Seeds) > 0 {
+		connectToSeeds(s, opts.Seeds)
 	}
 
 	s.loop()
@@ -137,6 +140,9 @@ func (s *Server) loop() {
 				s.logger.Printf("peer %s disconnected", peer.addr())
 			}
 
+		case t := <-s.peerCountCh:
+			t.count <- len(s.peers)
+
 		case <-s.quit:
 			s.shutdown()
 		}
@@ -208,6 +214,7 @@ func (s *Server) peerAlreadyConnected(addr net.Addr) bool {
 
 func (s *Server) sendLoop(peer Peer) {
 	// TODO: check if this peer is still connected.
+	// dont keep asking (maxPeers and no new nodes)
 	for {
 		getaddrMsg := newMessage(s.net, cmdGetAddr, nil)
 		peer.callGetaddr(getaddrMsg)
@@ -216,6 +223,31 @@ func (s *Server) sendLoop(peer Peer) {
 	}
 }
 
+type peerCount struct {
+	count chan int
+}
+
+// peerCount returns the number of connected peers to this server.
+func (s *Server) peerCount() int {
+	ch := peerCount{
+		count: make(chan int),
+	}
+
+	s.peerCountCh <- ch
+
+	return <-ch.count
+}
+
+// StartOpts holds the server configuration.
+type StartOpts struct {
+	// tcp port
+	TCP int
+	// slice of peer addresses the server will connect to
+	Seeds []string
+	// JSON-RPC port. If 0 no RPC handler will be attached.
+	RPC int
+}
+
 func logo() string {
 	return `
     _   ____________        __________ 
diff --git a/pkg/network/server_test.go b/pkg/network/server_test.go
index ad07d93e1..3c1771296 100644
--- a/pkg/network/server_test.go
+++ b/pkg/network/server_test.go
@@ -23,6 +23,19 @@ func TestHandleVersion(t *testing.T) {
 		t.Fatal("verack payload should be nil")
 	}
 }
+func TestPeerCount(t *testing.T) {
+	s := NewServer(ModeDevNet)
+	go s.loop()
+
+	lenPeers := 10
+	for i := 0; i < lenPeers; i++ {
+		s.register <- NewLocalPeer(s)
+	}
+
+	if have, want := s.peerCount(), lenPeers; want != have {
+		t.Fatalf("expected %d connected peers got %d", want, have)
+	}
+}
 
 func TestHandleAddrCmd(t *testing.T) {
 	// todo