From 699794d88d86c30fa8ac74c0bbe3e0ac9cde88a2 Mon Sep 17 00:00:00 2001
From: Andrey Petrov <andrey.petrov@shazow.net>
Date: Tue, 5 Jun 2018 15:45:43 -0400
Subject: [PATCH] p2p: More tests for AddTrustedPeer/RemoveTrustedPeer

---
 p2p/server_test.go | 60 +++++++++++++++++++++++++++++++++++++++-------
 1 file changed, 52 insertions(+), 8 deletions(-)

diff --git a/p2p/server_test.go b/p2p/server_test.go
index efab9bb5e..5fad1d0a7 100644
--- a/p2p/server_test.go
+++ b/p2p/server_test.go
@@ -148,7 +148,8 @@ func TestServerDial(t *testing.T) {
 
 	// tell the server to connect
 	tcpAddr := listener.Addr().(*net.TCPAddr)
-	srv.AddPeer(&discover.Node{ID: remid, IP: tcpAddr.IP, TCP: uint16(tcpAddr.Port)})
+	node := &discover.Node{ID: remid, IP: tcpAddr.IP, TCP: uint16(tcpAddr.Port)}
+	srv.AddPeer(node)
 
 	select {
 	case conn := <-accepted:
@@ -170,6 +171,21 @@ func TestServerDial(t *testing.T) {
 			if !reflect.DeepEqual(peers, []*Peer{peer}) {
 				t.Errorf("Peers mismatch: got %v, want %v", peers, []*Peer{peer})
 			}
+
+			// Test AddTrustedPeer/RemoveTrustedPeer and changing Trusted flags
+			// Particularly for race conditions on changing the flag state.
+			if peer := srv.Peers()[0]; peer.Info().Network.Trusted {
+				t.Errorf("peer is trusted prematurely: %v", peer)
+			}
+			srv.AddTrustedPeer(node)
+			if peer := srv.Peers()[0]; !peer.Info().Network.Trusted {
+				t.Errorf("peer is not trusted after AddTrustedPeer: %v", peer)
+			}
+			srv.RemoveTrustedPeer(node)
+			if peer := srv.Peers()[0]; peer.Info().Network.Trusted {
+				t.Errorf("peer is trusted after RemoveTrustedPeer: %v", peer)
+			}
+
 		case <-time.After(1 * time.Second):
 			t.Error("server did not launch peer within one second")
 		}
@@ -351,7 +367,8 @@ func TestServerAtCap(t *testing.T) {
 		}
 	}
 	// Try inserting a non-trusted connection.
-	c := newconn(randomID())
+	anotherID := randomID()
+	c := newconn(anotherID)
 	if err := srv.checkpoint(c, srv.posthandshake); err != DiscTooManyPeers {
 		t.Error("wrong error for insert:", err)
 	}
@@ -363,15 +380,41 @@ func TestServerAtCap(t *testing.T) {
 	if !c.is(trustedConn) {
 		t.Error("Server did not set trusted flag")
 	}
+
+	// Remove from trusted set and try again
+	srv.RemoveTrustedPeer(&discover.Node{ID: trustedID})
+	c = newconn(trustedID)
+	if err := srv.checkpoint(c, srv.posthandshake); err != DiscTooManyPeers {
+		t.Error("wrong error for insert:", err)
+	}
+
+	// Add anotherID to trusted set and try again
+	srv.AddTrustedPeer(&discover.Node{ID: anotherID})
+	c = newconn(anotherID)
+	if err := srv.checkpoint(c, srv.posthandshake); err != nil {
+		t.Error("unexpected error for trusted conn @posthandshake:", err)
+	}
+	if !c.is(trustedConn) {
+		t.Error("Server did not set trusted flag")
+	}
 }
 
 func TestServerPeerLimits(t *testing.T) {
 	srvkey := newkey()
-	srvid := discover.PubkeyID(&srvkey.PublicKey)
 
-	var tp *setupTransport = &setupTransport{id: srvid, phs: &protoHandshake{ID: srvid}}
+	clientid := randomID()
+	clientnode := &discover.Node{ID: clientid}
+
+	var tp *setupTransport = &setupTransport{
+		id: clientid,
+		phs: &protoHandshake{
+			ID: clientid,
+			// Force "DiscUselessPeer" due to unmatching caps
+			// Caps: []Cap{discard.cap()},
+		},
+	}
 	var flags connFlag = dynDialedConn
-	var dialDest *discover.Node = &discover.Node{ID: srvid}
+	var dialDest *discover.Node = &discover.Node{ID: clientid}
 
 	srv := &Server{
 		Config: Config{
@@ -386,6 +429,7 @@ func TestServerPeerLimits(t *testing.T) {
 	if err := srv.Start(); err != nil {
 		t.Fatalf("couldn't start server: %v", err)
 	}
+	defer srv.Stop()
 
 	// Check that server is full (MaxPeers=0)
 	conn, _ := net.Pipe()
@@ -395,7 +439,7 @@ func TestServerPeerLimits(t *testing.T) {
 	}
 	conn.Close()
 
-	srv.AddTrustedPeer(dialDest)
+	srv.AddTrustedPeer(clientnode)
 
 	// Check that server allows a trusted peer despite being full.
 	conn, _ = net.Pipe()
@@ -404,12 +448,12 @@ func TestServerPeerLimits(t *testing.T) {
 		t.Errorf("failed to bypass MaxPeers with trusted node: %q", tp.closeErr)
 	}
 
-	if tp.closeErr != DiscSelf {
+	if tp.closeErr != DiscUselessPeer {
 		t.Errorf("unexpected close error: %q", tp.closeErr)
 	}
 	conn.Close()
 
-	srv.RemoveTrustedPeer(dialDest)
+	srv.RemoveTrustedPeer(clientnode)
 
 	// Check that server is full again.
 	conn, _ = net.Pipe()