From d76b86febdfa4782a48a606eace7a0117a83a46c Mon Sep 17 00:00:00 2001
From: anthdm <anthony@academiclabs.co>
Date: Sat, 27 Jan 2018 08:37:07 +0100
Subject: [PATCH] wip

---
 pkg/network/message.go |  4 +++-
 pkg/network/server.go  | 42 ++++++++++++++++++++++++++++++------------
 2 files changed, 33 insertions(+), 13 deletions(-)

diff --git a/pkg/network/message.go b/pkg/network/message.go
index 940811c02..ac89a7f31 100644
--- a/pkg/network/message.go
+++ b/pkg/network/message.go
@@ -98,6 +98,8 @@ func (m *Message) commandType() commandType {
 		return cmdVersion
 	case "verack":
 		return cmdVerack
+	case "getaddr":
+		return cmdGetAddr
 	case "addr":
 		return cmdAddr
 	case "getheaders":
@@ -211,7 +213,7 @@ func newVersionPayload(p uint16, ua string, h uint32, r bool) *Version {
 		Services:    1,
 		Timestamp:   12345,
 		Port:        p,
-		Nonce:       1911099534,
+		Nonce:       19110,
 		UserAgent:   []byte(ua),
 		StartHeight: 0,
 		Relay:       r,
diff --git a/pkg/network/server.go b/pkg/network/server.go
index 8d1496bac..ea3e3fc8e 100644
--- a/pkg/network/server.go
+++ b/pkg/network/server.go
@@ -68,7 +68,9 @@ func NewServer(net NetMode) *Server {
 	}
 
 	s := &Server{
-		userAgent:  fmt.Sprintf("/NEO:%s/", version),
+		// It is important to have this user agent correct. Otherwise we will get
+		// disconnected.
+		userAgent:  fmt.Sprintf("\v/NEO:%s/", version),
 		logger:     logger,
 		peers:      make(map[*Peer]bool),
 		register:   make(chan *Peer),
@@ -118,6 +120,12 @@ func (s *Server) shutdown() {
 	}
 }
 
+func (s *Server) disconnect(p *Peer) {
+	p.conn.Close()
+	close(p.send)
+	s.unregister <- p
+}
+
 func (s *Server) loop() {
 	for {
 		select {
@@ -133,7 +141,10 @@ func (s *Server) loop() {
 				peer.send <- resp
 			}
 		case peer := <-s.unregister:
-			s.logger.Printf("peer %s disconnected", peer.conn.RemoteAddr())
+			if _, ok := s.peers[peer]; ok {
+				delete(s.peers, peer)
+				s.logger.Printf("peer %s disconnected", peer.conn.RemoteAddr())
+			}
 		case tuple := <-s.message:
 			if err := s.processMessage(tuple.msg, tuple.peer); err != nil {
 				s.logger.Fatalf("failed to process message: %s", err)
@@ -147,16 +158,15 @@ func (s *Server) loop() {
 // TODO: unregister peers on error.
 // processMessage processes the received message from a remote node.
 func (s *Server) processMessage(msg *Message, peer *Peer) error {
-	rpcLogger.Printf("IN :: %+v", msg)
+	rpcLogger.Printf("IN :: %+v", string(msg.Command))
 
 	switch msg.commandType() {
 	case cmdVersion:
-		v, _ := msg.decodePayload()
-		resp, err := s.handleVersionCmd(v.(*Version))
+		v, err := msg.decodePayload()
 		if err != nil {
 			return err
 		}
-		peer.send <- resp
+		s.handleVersionCmd(v.(*Version), peer)
 	case cmdVerack:
 	case cmdGetAddr:
 	case cmdAddr:
@@ -187,14 +197,22 @@ func (s *Server) handlePeerConnected() (*Message, error) {
 	return msg, nil
 }
 
-// Version declares the server's version when a new connection is been made.
-// We respond with a instant "verack" message.
-func (s *Server) handleVersionCmd(v *Version) (*Message, error) {
+// Version declares the server's version.
+func (s *Server) handleVersionCmd(v *Version, peer *Peer) {
 	// TODO: check version and verify to trust that node.
 
-	// Empty payload for the verack message.
-	msg := newMessage(s.net, cmdVerack, nil)
-	return msg, nil
+	payload := newVersionPayload(s.port, s.userAgent, 0, s.relay)
+	b, err := payload.encode()
+	if err != nil {
+		s.disconnect(peer)
+		return
+	}
+	versionMsg := newMessage(s.net, cmdVersion, b)
+	peer.send <- versionMsg
+
+	peer.verack = true
+	verackMsg := newMessage(s.net, cmdVerack, nil)
+	peer.send <- verackMsg
 }
 
 func logo() string {