update code and add some test cases

hadv · Oct 26, 2022 · 08e00ba · 08e00ba
1 parent 208c28f
commit 08e00ba
Show file tree

Hide file tree

Showing 3 changed files with 53 additions and 7 deletions.
diff --git a/node/node.go b/node/node.go
@@ -129,7 +129,7 @@ func New(conf *Config) (*Node, error) {
 		node.server.Config.StaticNodes = node.config.StaticNodes()
 	}
 	if node.server.Config.TrustedNodes == nil {
-		node.server.Config.TrustedNodes = append(node.config.TrustedNodes(), node.config.PrivateNodes()...)
+		node.server.Config.TrustedNodes = node.config.TrustedNodes()
 	}
 	if node.server.Config.NodeDatabase == "" {
 		node.server.Config.NodeDatabase = node.config.NodeDB()

diff --git a/p2p/server.go b/p2p/server.go
@@ -737,6 +737,9 @@ func (srv *Server) run() {
 	for _, n := range srv.TrustedNodes {
 		trusted[n.ID()] = true
 	}
+	for _, n := range srv.PrivateNodes {
+		trusted[n.ID()] = true
+	}
 
 running:
 	for {

diff --git a/p2p/server_test.go b/p2p/server_test.go
@@ -237,14 +237,17 @@ func TestServerRemovePeerDisconnect(t *testing.T) {
 // when the server is at capacity. Trusted connections should still be accepted.
 func TestServerAtCap(t *testing.T) {
 	trustedNode := newkey()
+	privateNode := newkey()
 	trustedID := enode.PubkeyToIDV4(&trustedNode.PublicKey)
+	privateID := enode.PubkeyToIDV4(&privateNode.PublicKey)
 	srv := &Server{
 		Config: Config{
 			PrivateKey:   newkey(),
 			MaxPeers:     10,
 			NoDial:       true,
 			NoDiscovery:  true,
 			TrustedNodes: []*enode.Node{newNode(trustedID, "")},
+			PrivateNodes: []*enode.Node{newNode(privateID, "")},
 			Logger:       testlog.Logger(t, log.LvlTrace),
 		},
 	}
@@ -260,6 +263,9 @@ func TestServerAtCap(t *testing.T) {
 		return &conn{fd: fd, transport: tx, flags: inboundConn, node: node, cont: make(chan error)}
 	}
 
+	// remove private node (it was added a static node at the beginning)
+	srv.RemovePeer(newNode(privateID, ""))
+
 	// Inject a few connections to fill up the peer set.
 	for i := 0; i < 10; i++ {
 		c := newconn(randomID())
@@ -282,6 +288,21 @@ func TestServerAtCap(t *testing.T) {
 		t.Error("Server did not set trusted flag")
 	}
 
+	// Try inserting a non-trusted connection.
+	anotherID = randomID()
+	c = newconn(anotherID)
+	if err := srv.checkpoint(c, srv.checkpointPostHandshake); err != DiscTooManyPeers {
+		t.Error("wrong error for insert:", err)
+	}
+	// Try inserting the private node as trusted connection as well
+	c = newconn(privateID)
+	if err := srv.checkpoint(c, srv.checkpointPostHandshake); err != nil {
+		t.Error("unexpected error for trusted conn @posthandshake:", err)
+	}
+	if !c.is(trustedConn) {
+		t.Error("Server did not set trusted flag")
+	}
+
 	// Remove from trusted set and try again
 	srv.RemoveTrustedPeer(newNode(trustedID, ""))
 	c = newconn(trustedID)
@@ -304,6 +325,8 @@ func TestServerPeerLimits(t *testing.T) {
 	srvkey := newkey()
 	clientkey := newkey()
 	clientnode := enode.NewV4(&clientkey.PublicKey, nil, 0, 0)
+	privatekey := newkey()
+	privatenode := enode.NewV4(&privatekey.PublicKey, nil, 0, 0)
 
 	var tp = &setupTransport{
 		pubkey: &clientkey.PublicKey,
@@ -316,12 +339,13 @@ func TestServerPeerLimits(t *testing.T) {
 
 	srv := &Server{
 		Config: Config{
-			PrivateKey:  srvkey,
-			MaxPeers:    0,
-			NoDial:      true,
-			NoDiscovery: true,
-			Protocols:   []Protocol{discard},
-			Logger:      testlog.Logger(t, log.LvlTrace),
+			PrivateKey:   srvkey,
+			MaxPeers:     0,
+			NoDial:       true,
+			NoDiscovery:  true,
+			Protocols:    []Protocol{discard},
+			PrivateNodes: []*enode.Node{privatenode},
+			Logger:       testlog.Logger(t, log.LvlTrace),
 		},
 		newTransport: func(fd net.Conn, dialDest *ecdsa.PublicKey) transport { return tp },
 	}
@@ -363,6 +387,25 @@ func TestServerPeerLimits(t *testing.T) {
 		t.Errorf("unexpected close error: %q", tp.closeErr)
 	}
 	conn.Close()
+
+	srv.RemovePeer(privatenode)
+
+	// Check that server is full again.
+	conn, _ = net.Pipe()
+	srv.SetupConn(conn, flags, dialDest)
+	if tp.closeErr != DiscTooManyPeers {
+		t.Errorf("unexpected close error: %q", tp.closeErr)
+	}
+	conn.Close()
+
+	// Check that server allows a private peer despite being full.
+	dialDest = privatenode
+	conn, _ = net.Pipe()
+	srv.SetupConn(conn, flags, dialDest)
+	if tp.closeErr == DiscTooManyPeers {
+		t.Errorf("failed to bypass MaxPeers with private node: %q", tp.closeErr)
+	}
+	conn.Close()
 }
 
 func TestServerSetupConn(t *testing.T) {