@@ -22,8 +22,11 @@ func startTestServer(t *testing.T, pf newPeerHook) *Server {
22
22
ListenAddr : "127.0.0.1:0" ,
23
23
PrivateKey : newkey (),
24
24
newPeerHook : pf ,
25
- setupFunc : func (fd net.Conn , prv * ecdsa.PrivateKey , our * protoHandshake , dial * discover.Node , atcap bool , trusted map [ discover.NodeID ] bool ) (* conn , error ) {
25
+ setupFunc : func (fd net.Conn , prv * ecdsa.PrivateKey , our * protoHandshake , dial * discover.Node , keepconn func ( discover.NodeID ) bool ) (* conn , error ) {
26
26
id := randomID ()
27
+ if ! keepconn (id ) {
28
+ return nil , DiscAlreadyConnected
29
+ }
27
30
rw := newRlpxFrameRW (fd , secrets {
28
31
MAC : zero16 ,
29
32
AES : zero16 ,
@@ -200,7 +203,7 @@ func TestServerDisconnectAtCap(t *testing.T) {
200
203
// Run the handshakes just like a real peer would.
201
204
key := newkey ()
202
205
hs := & protoHandshake {Version : baseProtocolVersion , ID : discover .PubkeyID (& key .PublicKey )}
203
- _ , err = setupConn (conn , key , hs , srv .Self (), false , srv . trustedNodes )
206
+ _ , err = setupConn (conn , key , hs , srv .Self (), keepalways )
204
207
if i == nconns - 1 {
205
208
// When handling the last connection, the server should
206
209
// disconnect immediately instead of running the protocol
@@ -250,7 +253,7 @@ func TestServerStaticPeers(t *testing.T) {
250
253
// Run the handshakes just like a real peer would, and wait for completion
251
254
key := newkey ()
252
255
shake := & protoHandshake {Version : baseProtocolVersion , ID : discover .PubkeyID (& key .PublicKey )}
253
- if _ , err = setupConn (conn , key , shake , server .Self (), false , server . trustedNodes ); err != nil {
256
+ if _ , err = setupConn (conn , key , shake , server .Self (), keepalways ); err != nil {
254
257
t .Fatalf ("conn %d: unexpected error: %v" , i , err )
255
258
}
256
259
<- started
@@ -344,7 +347,7 @@ func TestServerTrustedPeers(t *testing.T) {
344
347
// Run the handshakes just like a real peer would, and wait for completion
345
348
key := newkey ()
346
349
shake := & protoHandshake {Version : baseProtocolVersion , ID : discover .PubkeyID (& key .PublicKey )}
347
- if _ , err = setupConn (conn , key , shake , server .Self (), false , server . trustedNodes ); err != nil {
350
+ if _ , err = setupConn (conn , key , shake , server .Self (), keepalways ); err != nil {
348
351
t .Fatalf ("conn %d: unexpected error: %v" , i , err )
349
352
}
350
353
<- started
@@ -357,7 +360,7 @@ func TestServerTrustedPeers(t *testing.T) {
357
360
defer conn .Close ()
358
361
359
362
shake := & protoHandshake {Version : baseProtocolVersion , ID : trusted .ID }
360
- if _ , err = setupConn (conn , key , shake , server .Self (), false , server . trustedNodes ); err != nil {
363
+ if _ , err = setupConn (conn , key , shake , server .Self (), keepalways ); err != nil {
361
364
t .Fatalf ("trusted node: unexpected error: %v" , err )
362
365
}
363
366
select {
@@ -472,7 +475,7 @@ func TestServerMaxPendingAccepts(t *testing.T) {
472
475
go func () {
473
476
key := newkey ()
474
477
shake := & protoHandshake {Version : baseProtocolVersion , ID : discover .PubkeyID (& key .PublicKey )}
475
- if _ , err := setupConn (conns [1 ], key , shake , server .Self (), false , server . trustedNodes ); err != nil {
478
+ if _ , err := setupConn (conns [1 ], key , shake , server .Self (), keepalways ); err != nil {
476
479
t .Fatalf ("failed to run handshake: %v" , err )
477
480
}
478
481
}()
@@ -486,7 +489,7 @@ func TestServerMaxPendingAccepts(t *testing.T) {
486
489
go func () {
487
490
key := newkey ()
488
491
shake := & protoHandshake {Version : baseProtocolVersion , ID : discover .PubkeyID (& key .PublicKey )}
489
- if _ , err := setupConn (conns [0 ], key , shake , server .Self (), false , server . trustedNodes ); err != nil {
492
+ if _ , err := setupConn (conns [0 ], key , shake , server .Self (), keepalways ); err != nil {
490
493
t .Fatalf ("failed to run handshake: %v" , err )
491
494
}
492
495
}()
@@ -513,3 +516,7 @@ func randomID() (id discover.NodeID) {
513
516
}
514
517
return id
515
518
}
519
+
520
+ func keepalways (id discover.NodeID ) bool {
521
+ return true
522
+ }
0 commit comments