@@ -148,7 +148,8 @@ func TestServerDial(t *testing.T) {
148148
149149 // tell the server to connect
150150 tcpAddr := listener .Addr ().(* net.TCPAddr )
151- srv .AddPeer (& discover.Node {ID : remid , IP : tcpAddr .IP , TCP : uint16 (tcpAddr .Port )})
151+ node := & discover.Node {ID : remid , IP : tcpAddr .IP , TCP : uint16 (tcpAddr .Port )}
152+ srv .AddPeer (node )
152153
153154 select {
154155 case conn := <- accepted :
@@ -169,6 +170,29 @@ func TestServerDial(t *testing.T) {
169170 if ! reflect .DeepEqual (peers , []* Peer {peer }) {
170171 t .Errorf ("Peers mismatch: got %v, want %v" , peers , []* Peer {peer })
171172 }
173+
174+ // Test AddTrustedPeer/RemoveTrustedPeer and changing Trusted flags
175+ // Particularly for race conditions on changing the flag state.
176+ if peer := srv .Peers ()[0 ]; peer .Info ().Network .Trusted {
177+ t .Errorf ("peer is trusted prematurely: %v" , peer )
178+ }
179+ done := make (chan bool )
180+ go func () {
181+ srv .AddTrustedPeer (node )
182+ if peer := srv .Peers ()[0 ]; ! peer .Info ().Network .Trusted {
183+ t .Errorf ("peer is not trusted after AddTrustedPeer: %v" , peer )
184+ }
185+ srv .RemoveTrustedPeer (node )
186+ if peer := srv .Peers ()[0 ]; peer .Info ().Network .Trusted {
187+ t .Errorf ("peer is trusted after RemoveTrustedPeer: %v" , peer )
188+ }
189+ done <- true
190+ }()
191+ // Trigger potential race conditions
192+ peer = srv .Peers ()[0 ]
193+ _ = peer .Inbound ()
194+ _ = peer .Info ()
195+ <- done
172196 case <- time .After (1 * time .Second ):
173197 t .Error ("server did not launch peer within one second" )
174198 }
@@ -365,7 +389,8 @@ func TestServerAtCap(t *testing.T) {
365389 }
366390 }
367391 // Try inserting a non-trusted connection.
368- c := newconn (randomID ())
392+ anotherID := randomID ()
393+ c := newconn (anotherID )
369394 if err := srv .checkpoint (c , srv .posthandshake ); err != DiscTooManyPeers {
370395 t .Error ("wrong error for insert:" , err )
371396 }
@@ -378,6 +403,87 @@ func TestServerAtCap(t *testing.T) {
378403 t .Error ("Server did not set trusted flag" )
379404 }
380405
406+ // Remove from trusted set and try again
407+ srv .RemoveTrustedPeer (& discover.Node {ID : trustedID })
408+ c = newconn (trustedID )
409+ if err := srv .checkpoint (c , srv .posthandshake ); err != DiscTooManyPeers {
410+ t .Error ("wrong error for insert:" , err )
411+ }
412+
413+ // Add anotherID to trusted set and try again
414+ srv .AddTrustedPeer (& discover.Node {ID : anotherID })
415+ c = newconn (anotherID )
416+ if err := srv .checkpoint (c , srv .posthandshake ); err != nil {
417+ t .Error ("unexpected error for trusted conn @posthandshake:" , err )
418+ }
419+ if ! c .is (trustedConn ) {
420+ t .Error ("Server did not set trusted flag" )
421+ }
422+ }
423+
424+ func TestServerPeerLimits (t * testing.T ) {
425+ srvkey := newkey ()
426+
427+ clientid := randomID ()
428+ clientnode := & discover.Node {ID : clientid }
429+
430+ var tp * setupTransport = & setupTransport {
431+ id : clientid ,
432+ phs : & protoHandshake {
433+ ID : clientid ,
434+ // Force "DiscUselessPeer" due to unmatching caps
435+ // Caps: []Cap{discard.cap()},
436+ },
437+ }
438+ var flags connFlag = dynDialedConn
439+ var dialDest * discover.Node = & discover.Node {ID : clientid }
440+
441+ srv := & Server {
442+ Config : Config {
443+ PrivateKey : srvkey ,
444+ MaxPeers : 0 ,
445+ NoDial : true ,
446+ Protocols : []Protocol {discard },
447+ },
448+ newTransport : func (fd net.Conn ) transport { return tp },
449+ log : log .New (),
450+ }
451+ if err := srv .Start (); err != nil {
452+ t .Fatalf ("couldn't start server: %v" , err )
453+ }
454+ defer srv .Stop ()
455+
456+ // Check that server is full (MaxPeers=0)
457+ conn , _ := net .Pipe ()
458+ srv .SetupConn (conn , flags , dialDest )
459+ if tp .closeErr != DiscTooManyPeers {
460+ t .Errorf ("unexpected close error: %q" , tp .closeErr )
461+ }
462+ conn .Close ()
463+
464+ srv .AddTrustedPeer (clientnode )
465+
466+ // Check that server allows a trusted peer despite being full.
467+ conn , _ = net .Pipe ()
468+ srv .SetupConn (conn , flags , dialDest )
469+ if tp .closeErr == DiscTooManyPeers {
470+ t .Errorf ("failed to bypass MaxPeers with trusted node: %q" , tp .closeErr )
471+ }
472+
473+ if tp .closeErr != DiscUselessPeer {
474+ t .Errorf ("unexpected close error: %q" , tp .closeErr )
475+ }
476+ conn .Close ()
477+
478+ srv .RemoveTrustedPeer (clientnode )
479+
480+ // Check that server is full again.
481+ conn , _ = net .Pipe ()
482+ srv .SetupConn (conn , flags , dialDest )
483+ if tp .closeErr != DiscTooManyPeers {
484+ t .Errorf ("unexpected close error: %q" , tp .closeErr )
485+ }
486+ conn .Close ()
381487}
382488
383489func TestServerSetupConn (t * testing.T ) {
0 commit comments