p2p: implement AddNotaryPeer and RemoveNotaryPeer

AddNotaryPeer adds node to static node set so that server will maintain
the connection with the notary node.
AddNotaryPeer also sets the notaryConn flag to allow the node to always
connect, even if the slot are full.

RemoveNotaryPeer removes node from static, then disconnect and unsets
the notaryConn flag.

1 files changed, 118 insertions, 3 deletions

diff --git a/p2p/server_test.go b/p2p/server_test.go
index f665c1424..c3ff825a3 100644
--- a/p2p/server_test.go
+++ b/p2p/server_test.go
@@ -27,6 +27,7 @@ import (
 
 	"github.com/ethereum/go-ethereum/crypto"
 	"github.com/ethereum/go-ethereum/log"
+	"github.com/ethereum/go-ethereum/p2p/discover"
 	"github.com/ethereum/go-ethereum/p2p/enode"
 	"github.com/ethereum/go-ethereum/p2p/enr"
 	"golang.org/x/crypto/sha3"
@@ -172,10 +173,14 @@ func TestServerDial(t *testing.T) {
 			}
 
 			// Test AddTrustedPeer/RemoveTrustedPeer and changing Trusted flags
+			// Test AddNotaryPeer/RemoveTrustedPeer and changing Notary 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)
 			}
+			if peer := srv.Peers()[0]; peer.Info().Network.Notary {
+				t.Errorf("peer is notary prematurely: %v", peer)
+			}
 			done := make(chan bool)
 			go func() {
 				srv.AddTrustedPeer(node)
@@ -186,6 +191,15 @@ func TestServerDial(t *testing.T) {
 				if peer := srv.Peers()[0]; peer.Info().Network.Trusted {
 					t.Errorf("peer is trusted after RemoveTrustedPeer: %v", peer)
 				}
+
+				srv.AddNotaryPeer(node)
+				if peer := srv.Peers()[0]; !peer.Info().Network.Notary {
+					t.Errorf("peer is not notary after AddNotaryPeer: %v", peer)
+				}
+				srv.RemoveNotaryPeer(node)
+				if peer := srv.Peers()[0]; peer.Info().Network.Notary {
+					t.Errorf("peer is notary after RemoveNotaryPeer: %v", peer)
+				}
 				done <- true
 			}()
 			// Trigger potential race conditions
@@ -202,6 +216,73 @@ func TestServerDial(t *testing.T) {
 	}
 }
 
+// TestNotaryPeer checks that the node is added to and remove from static when
+// AddNotaryPeer and RemoveNotaryPeer is called.
+func TestNotaryPeer(t *testing.T) {
+	var (
+		returned    = make(chan struct{})
+		add, remove = make(chan *discover.Node), make(chan *discover.Node)
+		tg          = taskgen{
+			newFunc: func(running int, peers map[discover.NodeID]*Peer) []task {
+				return []task{}
+			},
+			doneFunc: func(t task) {},
+			addFunc: func(n *discover.Node) {
+				add <- n
+			},
+			removeFunc: func(n *discover.Node) {
+				remove <- n
+			},
+		}
+	)
+
+	srv := &Server{
+		Config:       Config{MaxPeers: 10},
+		quit:         make(chan struct{}),
+		ntab:         fakeTable{},
+		addnotary:    make(chan *discover.Node),
+		removenotary: make(chan *discover.Node),
+		running:      true,
+		log:          log.New(),
+	}
+	srv.loopWG.Add(1)
+	go func() {
+		srv.run(tg)
+		close(returned)
+	}()
+
+	notaryID := randomID()
+	go srv.AddNotaryPeer(&discover.Node{ID: notaryID})
+
+	select {
+	case n := <-add:
+		if n.ID != notaryID {
+			t.Errorf("node ID mismatched: got %s, want %s",
+				n.ID.String(), notaryID.String())
+		}
+	case <-time.After(1 * time.Second):
+		t.Error("add static is not called within one second")
+	}
+
+	go srv.RemoveNotaryPeer(&discover.Node{ID: notaryID})
+	select {
+	case n := <-remove:
+		if n.ID != notaryID {
+			t.Errorf("node ID mismatched: got %s, want %s",
+				n.ID.String(), notaryID.String())
+		}
+	case <-time.After(1 * time.Second):
+		t.Error("remove static is not called within one second")
+	}
+
+	srv.Stop()
+	select {
+	case <-returned:
+	case <-time.After(500 * time.Millisecond):
+		t.Error("Server.run did not return within 500ms")
+	}
+}
+
 // This test checks that tasks generated by dialstate are
 // actually executed and taskdone is called for them.
 func TestServerTaskScheduling(t *testing.T) {
@@ -326,6 +407,9 @@ func TestServerManyTasks(t *testing.T) {
 type taskgen struct {
 	newFunc  func(running int, peers map[enode.ID]*Peer) []task
 	doneFunc func(task)
+
+	addFunc    func(*discover.Node)
+	removeFunc func(*discover.Node)
 }
 
 func (tg taskgen) newTasks(running int, peers map[enode.ID]*Peer, now time.Time) []task {
@@ -334,9 +418,11 @@ func (tg taskgen) newTasks(running int, peers map[enode.ID]*Peer, now time.Time)
 func (tg taskgen) taskDone(t task, now time.Time) {
 	tg.doneFunc(t)
 }
-func (tg taskgen) addStatic(*enode.Node) {
+func (tg taskgen) addStatic(n *enode.Node) {
+	tg.addFunc(n)
 }
-func (tg taskgen) removeStatic(*enode.Node) {
+func (tg taskgen) removeStatic(n *enode.Node) {
+	tg.removeFunc(n)
 }
 
 type testTask struct {
@@ -350,10 +436,11 @@ func (t *testTask) Do(srv *Server) {
 
 // This test checks that connections are disconnected
 // just after the encryption handshake when the server is
-// at capacity. Trusted connections should still be accepted.
+// at capacity. Trusted and Notary connections should still be accepted.
 func TestServerAtCap(t *testing.T) {
 	trustedNode := newkey()
 	trustedID := enode.PubkeyToIDV4(&trustedNode.PublicKey)
+	notaryID := randomID()
 	srv := &Server{
 		Config: Config{
 			PrivateKey:   newkey(),
@@ -366,6 +453,7 @@ func TestServerAtCap(t *testing.T) {
 		t.Fatalf("could not start: %v", err)
 	}
 	defer srv.Stop()
+	srv.AddNotaryPeer(&discover.Node{ID: notaryID})
 
 	newconn := func(id enode.ID) *conn {
 		fd, _ := net.Pipe()
@@ -396,6 +484,15 @@ func TestServerAtCap(t *testing.T) {
 		t.Error("Server did not set trusted flag")
 	}
 
+	// Try inserting a notary connection.
+	c = newconn(notaryID)
+	if err := srv.checkpoint(c, srv.posthandshake); err != nil {
+		t.Error("unexpected error for notary conn @posthandshake:", err)
+	}
+	if !c.is(notaryConn) {
+		t.Error("Server did not set notary flag")
+	}
+
 	// Remove from trusted set and try again
 	srv.RemoveTrustedPeer(newNode(trustedID, nil))
 	c = newconn(trustedID)
@@ -412,6 +509,24 @@ func TestServerAtCap(t *testing.T) {
 	if !c.is(trustedConn) {
 		t.Error("Server did not set trusted flag")
 	}
+
+	// Remove from notary set and try again
+	srv.RemoveNotaryPeer(&discover.Node{ID: notaryID})
+	c = newconn(notaryID)
+	if err := srv.checkpoint(c, srv.posthandshake); err != DiscTooManyPeers {
+		t.Error("wrong error for insert:", err)
+	}
+
+	// Add anotherID to notary set and try again
+	anotherNotaryID := randomID()
+	srv.AddNotaryPeer(&discover.Node{ID: anotherNotaryID})
+	c = newconn(anotherNotaryID)
+	if err := srv.checkpoint(c, srv.posthandshake); err != nil {
+		t.Error("unexpected error for notary conn @posthandshake:", err)
+	}
+	if !c.is(notaryConn) {
+		t.Error("Server did not set notary flag")
+	}
 }
 
 func TestServerPeerLimits(t *testing.T) {