github.com/annchain/OG@v0.0.9/p2p/server_test.go (about) 1 // Copyright 2014 The go-ethereum Authors 2 // This file is part of the go-ethereum library. 3 // 4 // The go-ethereum library is free software: you can redistribute it and/or modify 5 // it under the terms of the GNU Lesser General Public License as published by 6 // the Free Software Foundation, either version 3 of the License, or 7 // (at your option) any later version. 8 // 9 // The go-ethereum library is distributed in the hope that it will be useful, 10 // but WITHOUT ANY WARRANTY; without even the implied warranty of 11 // MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 12 // GNU Lesser General Public License for more details. 13 // 14 // You should have received a copy of the GNU Lesser General Public License 15 // along with the go-ethereum library. If not, see <http://www.gnu.org/licenses/>. 16 17 package p2p 18 19 import ( 20 "crypto/ecdsa" 21 "errors" 22 ogcrypto2 "github.com/annchain/OG/deprecated/ogcrypto" 23 "math/rand" 24 "net" 25 "reflect" 26 "testing" 27 "time" 28 29 "github.com/annchain/OG/p2p/enr" 30 "github.com/annchain/OG/p2p/onode" 31 "golang.org/x/crypto/sha3" 32 ) 33 34 func init() { 35 // log.Root().SetHandler(log.LvlFilterHandler(log.LvlError, log.StreamHandler(os.Stderr, log.TerminalFormat(false)))) 36 } 37 38 type testTransport struct { 39 rpub *ecdsa.PublicKey 40 *rlpx 41 42 closeErr error 43 } 44 45 func newTestTransport(rpub *ecdsa.PublicKey, fd net.Conn) transport { 46 wrapped := newRLPX(fd).(*rlpx) 47 wrapped.rw = newRLPXFrameRW(fd, secrets{ 48 MAC: zero16, 49 AES: zero16, 50 IngressMAC: sha3.NewLegacyKeccak256(), 51 EgressMAC: sha3.NewLegacyKeccak256(), 52 }) 53 return &testTransport{rpub: rpub, rlpx: wrapped} 54 } 55 56 func (c *testTransport) doEncHandshake(prv *ecdsa.PrivateKey, dialDest *ecdsa.PublicKey) (*ecdsa.PublicKey, error) { 57 return c.rpub, nil 58 } 59 60 func (c *testTransport) doProtoHandshake(our *ProtoHandshake) (*ProtoHandshake, error) { 61 pubkey := ogcrypto2.FromECDSAPub(c.rpub)[1:] 62 return &ProtoHandshake{ID: pubkey, Name: "test"}, nil 63 } 64 65 func (c *testTransport) close(err error) { 66 c.rlpx.fd.Close() 67 c.closeErr = err 68 } 69 func startTestServer(t *testing.T, remoteKey *ecdsa.PublicKey, pf func(*Peer)) *Server { 70 config := Config{ 71 NodeName: "test", 72 MaxPeers: 10, 73 ListenAddr: "127.0.0.1:0", 74 PrivateKey: newkey(), 75 } 76 server := &Server{ 77 Config: config, 78 newPeerHook: pf, 79 newTransport: func(fd net.Conn) transport { return newTestTransport(remoteKey, fd) }, 80 } 81 if err := server.start(); err != nil { 82 t.Fatalf("Could not start server: %v", err) 83 } 84 return server 85 } 86 87 func TestServerListen(t *testing.T) { 88 // start the test server 89 connected := make(chan *Peer) 90 remid := &newkey().PublicKey 91 srv := startTestServer(t, remid, func(p *Peer) { 92 if p.ID() != onode.PubkeyToIDV4(remid) { 93 t.Error("peer func called with wrong node id") 94 } 95 connected <- p 96 }) 97 defer close(connected) 98 defer srv.Stop() 99 100 // dial the test server 101 conn, err := net.DialTimeout("tcp", srv.ListenAddr, 5*time.Second) 102 if err != nil { 103 t.Fatalf("could not dial: %v", err) 104 } 105 defer conn.Close() 106 107 select { 108 case peer := <-connected: 109 if peer.LocalAddr().String() != conn.RemoteAddr().String() { 110 t.Errorf("peer started with wrong conn: got %v, want %v", 111 peer.LocalAddr(), conn.RemoteAddr()) 112 } 113 peers := srv.Peers() 114 if !reflect.DeepEqual(peers, []*Peer{peer}) { 115 t.Errorf("PeerPipeIns mismatch: got %v, want %v", peers, []*Peer{peer}) 116 } 117 case <-time.After(1 * time.Second): 118 t.Error("server did not accept within one second") 119 } 120 } 121 func TestServerDial(t *testing.T) { 122 // run a one-shot TCP server to handle the connection. 123 listener, err := net.Listen("tcp", "127.0.0.1:0") 124 if err != nil { 125 t.Fatalf("could not setup listener: %v", err) 126 } 127 defer listener.Close() 128 accepted := make(chan net.Conn) 129 go func() { 130 conn, err := listener.Accept() 131 if err != nil { 132 t.Error("accept error:", err) 133 return 134 } 135 accepted <- conn 136 }() 137 138 // start the server 139 connected := make(chan *Peer) 140 remid := &newkey().PublicKey 141 srv := startTestServer(t, remid, func(p *Peer) { connected <- p }) 142 defer close(connected) 143 defer srv.Stop() 144 145 // tell the server to connect 146 tcpAddr := listener.Addr().(*net.TCPAddr) 147 node := onode.NewV4(remid, tcpAddr.IP, tcpAddr.Port, 0) 148 srv.AddPeer(node) 149 150 select { 151 case conn := <-accepted: 152 defer conn.Close() 153 154 select { 155 case peer := <-connected: 156 if peer.ID() != onode.PubkeyToIDV4(remid) { 157 t.Errorf("peer has wrong id") 158 } 159 if peer.Name() != "test" { 160 t.Errorf("peer has wrong name") 161 } 162 if peer.RemoteAddr().String() != conn.LocalAddr().String() { 163 t.Errorf("peer started with wrong conn: got %v, want %v", 164 peer.RemoteAddr(), conn.LocalAddr()) 165 } 166 peers := srv.Peers() 167 if !reflect.DeepEqual(peers, []*Peer{peer}) { 168 t.Errorf("PeerPipeIns mismatch: got %v, want %v", peers, []*Peer{peer}) 169 } 170 171 // Test AddTrustedPeer/RemoveTrustedPeer and changing Trusted flags 172 // Particularly for race conditions on changing the flag state. 173 if peer := srv.Peers()[0]; peer.Info().Network.Trusted { 174 t.Errorf("peer is trusted prematurely: %v", peer) 175 } 176 done := make(chan bool) 177 go func() { 178 srv.AddTrustedPeer(node) 179 if peer := srv.Peers()[0]; !peer.Info().Network.Trusted { 180 t.Errorf("peer is not trusted after AddTrustedPeer: %v", peer) 181 } 182 srv.RemoveTrustedPeer(node) 183 if peer := srv.Peers()[0]; peer.Info().Network.Trusted { 184 t.Errorf("peer is trusted after RemoveTrustedPeer: %v", peer) 185 } 186 done <- true 187 }() 188 // Trigger potential race conditions 189 peer = srv.Peers()[0] 190 _ = peer.Inbound() 191 _ = peer.Info() 192 <-done 193 case <-time.After(1 * time.Second): 194 t.Error("server did not launch peer within one second") 195 } 196 197 case <-time.After(1 * time.Second): 198 t.Error("server did not connect within one second") 199 } 200 } 201 202 // This test checks that tasks generated by dialstate are 203 // actually executed and taskdone is called for them. 204 func TestServerTaskScheduling(t *testing.T) { 205 var ( 206 done = make(chan *testTask) 207 quit, returned = make(chan struct{}), make(chan struct{}) 208 tc = 0 209 tg = taskgen{ 210 newFunc: func(running int, peers map[onode.ID]*Peer) []task { 211 tc++ 212 return []task{&testTask{index: tc - 1}} 213 }, 214 doneFunc: func(t task) { 215 select { 216 case done <- t.(*testTask): 217 case <-quit: 218 } 219 }, 220 } 221 ) 222 223 // The Server in this test isn't actually running 224 // because we're only interested in what run does. 225 // The Server in this test isn't actually running 226 // because we're only interested in what run does. 227 db, _ := onode.OpenDB("") 228 srv := &Server{ 229 Config: Config{MaxPeers: 10}, 230 localnode: onode.NewLocalNode(db, newkey()), 231 nodedb: db, 232 quit: make(chan struct{}), 233 ntab: fakeTable{}, 234 running: true, 235 } 236 srv.loopWG.Add(1) 237 go func() { 238 srv.run(tg) 239 close(returned) 240 }() 241 242 var gotdone []*testTask 243 for i := 0; i < 100; i++ { 244 gotdone = append(gotdone, <-done) 245 } 246 for i, task := range gotdone { 247 if task.index != i { 248 t.Errorf("task %d has wrong index, got %d", i, task.index) 249 break 250 } 251 if !task.called { 252 t.Errorf("task %d was not called", i) 253 break 254 } 255 } 256 257 close(quit) 258 srv.Stop() 259 select { 260 case <-returned: 261 case <-time.After(500 * time.Millisecond): 262 t.Error("Server.run did not return within 500ms") 263 } 264 } 265 266 // This test checks that Server doesn't drop tasks, 267 // even if newTasks returns more than the maximum number of tasks. 268 func TestServerManyTasks(t *testing.T) { 269 alltasks := make([]task, 300) 270 for i := range alltasks { 271 alltasks[i] = &testTask{index: i} 272 } 273 274 var ( 275 db, _ = onode.OpenDB("") 276 srv = &Server{ 277 quit: make(chan struct{}), 278 localnode: onode.NewLocalNode(db, newkey()), 279 nodedb: db, 280 ntab: fakeTable{}, 281 running: true, 282 } 283 done = make(chan *testTask) 284 start, end = 0, 0 285 ) 286 defer srv.Stop() 287 srv.loopWG.Add(1) 288 go srv.run(taskgen{ 289 newFunc: func(running int, peers map[onode.ID]*Peer) []task { 290 start, end = end, end+maxActiveDialTasks+10 291 if end > len(alltasks) { 292 end = len(alltasks) 293 } 294 return alltasks[start:end] 295 }, 296 doneFunc: func(tt task) { 297 done <- tt.(*testTask) 298 }, 299 }) 300 301 doneset := make(map[int]bool) 302 timeout := time.After(2 * time.Second) 303 for len(doneset) < len(alltasks) { 304 select { 305 case tt := <-done: 306 if doneset[tt.index] { 307 t.Errorf("task %d got done more than once", tt.index) 308 } else { 309 doneset[tt.index] = true 310 } 311 case <-timeout: 312 t.Errorf("%d of %d tasks got done within 2s", len(doneset), len(alltasks)) 313 for i := 0; i < len(alltasks); i++ { 314 if !doneset[i] { 315 t.Logf("task %d not done", i) 316 } 317 } 318 return 319 } 320 } 321 } 322 323 type taskgen struct { 324 newFunc func(running int, peers map[onode.ID]*Peer) []task 325 doneFunc func(task) 326 } 327 328 func (tg taskgen) newTasks(running int, peers map[onode.ID]*Peer, now time.Time) []task { 329 return tg.newFunc(running, peers) 330 } 331 func (tg taskgen) taskDone(t task, now time.Time) { 332 tg.doneFunc(t) 333 } 334 func (tg taskgen) addStatic(*onode.Node) { 335 } 336 func (tg taskgen) removeStatic(*onode.Node) { 337 } 338 339 type testTask struct { 340 index int 341 called bool 342 } 343 344 func (t *testTask) Do(srv *Server) { 345 t.called = true 346 } 347 348 // This test checks that connections are disconnected 349 // just after the encryption handshake when the server is 350 // at capacity. Trusted connections should still be accepted. 351 func TestServerAtCap(t *testing.T) { 352 trustedNode := newkey() 353 trustedID := onode.PubkeyToIDV4(&trustedNode.PublicKey) 354 srv := &Server{ 355 Config: Config{ 356 PrivateKey: newkey(), 357 MaxPeers: 10, 358 NoDial: true, 359 TrustedNodes: []*onode.Node{newNode(trustedID, nil)}, 360 }, 361 } 362 if err := srv.start(); err != nil { 363 t.Fatalf("could not start: %v", err) 364 } 365 defer srv.Stop() 366 367 newconn := func(id onode.ID) *conn { 368 fd, _ := net.Pipe() 369 tx := newTestTransport(&trustedNode.PublicKey, fd) 370 node := onode.SignNull(new(enr.Record), id) 371 return &conn{fd: fd, transport: tx, flags: inboundConn, node: node, cont: make(chan error)} 372 } 373 374 // Inject a few connections to fill up the peer set. 375 for i := 0; i < 10; i++ { 376 c := newconn(randomID()) 377 if err := srv.checkpoint(c, srv.addpeer); err != nil { 378 t.Fatalf("could not add conn %d: %v", i, err) 379 } 380 } 381 // Try inserting a non-trusted connection. 382 anotherID := randomID() 383 c := newconn(anotherID) 384 if err := srv.checkpoint(c, srv.posthandshake); err != DiscTooManyPeers { 385 t.Error("wrong error for insert:", err) 386 } 387 // Try inserting a trusted connection. 388 c = newconn(trustedID) 389 if err := srv.checkpoint(c, srv.posthandshake); err != nil { 390 t.Error("unexpected error for trusted conn @posthandshake:", err) 391 } 392 if !c.is(trustedConn) { 393 t.Error("Server did not set trusted flag") 394 } 395 396 // Remove from trusted set and try again 397 srv.RemoveTrustedPeer(newNode(trustedID, nil)) 398 c = newconn(trustedID) 399 if err := srv.checkpoint(c, srv.posthandshake); err != DiscTooManyPeers { 400 t.Error("wrong error for insert:", err) 401 } 402 403 // Add anotherID to trusted set and try again 404 srv.AddTrustedPeer(newNode(anotherID, nil)) 405 c = newconn(anotherID) 406 if err := srv.checkpoint(c, srv.posthandshake); err != nil { 407 t.Error("unexpected error for trusted conn @posthandshake:", err) 408 } 409 if !c.is(trustedConn) { 410 t.Error("Server did not set trusted flag") 411 } 412 } 413 414 func TestServerPeerLimits(t *testing.T) { 415 srvkey := newkey() 416 clientkey := newkey() 417 clientnode := onode.NewV4(&clientkey.PublicKey, nil, 0, 0) 418 419 var tp = &setupTransport{ 420 pubkey: &clientkey.PublicKey, 421 phs: ProtoHandshake{ 422 ID: ogcrypto2.FromECDSAPub(&clientkey.PublicKey)[1:], 423 // Force "DiscUselessPeer" due to unmatching caps 424 // Caps: []Cap{discard.cap()}, 425 }, 426 } 427 428 srv := &Server{ 429 Config: Config{ 430 PrivateKey: srvkey, 431 MaxPeers: 0, 432 NoDial: true, 433 Protocols: []Protocol{discard}, 434 }, 435 newTransport: func(fd net.Conn) transport { return tp }, 436 } 437 if err := srv.start(); err != nil { 438 t.Fatalf("couldn't start server: %v", err) 439 } 440 defer srv.Stop() 441 442 // Check that server is full (MaxPeers=0) 443 flags := dynDialedConn 444 dialDest := clientnode 445 conn, _ := net.Pipe() 446 srv.SetupConn(conn, flags, dialDest) 447 if tp.closeErr != DiscTooManyPeers { 448 t.Errorf("unexpected close error: %q", tp.closeErr) 449 } 450 conn.Close() 451 452 srv.AddTrustedPeer(clientnode) 453 454 // Check that server allows a trusted peer despite being full. 455 conn, _ = net.Pipe() 456 srv.SetupConn(conn, flags, dialDest) 457 if tp.closeErr == DiscTooManyPeers { 458 t.Errorf("failed to bypass MaxPeers with trusted node: %q", tp.closeErr) 459 } 460 461 if tp.closeErr != DiscUselessPeer { 462 t.Errorf("unexpected close error: %q", tp.closeErr) 463 } 464 conn.Close() 465 466 srv.RemoveTrustedPeer(clientnode) 467 468 // Check that server is full again. 469 conn, _ = net.Pipe() 470 srv.SetupConn(conn, flags, dialDest) 471 if tp.closeErr != DiscTooManyPeers { 472 t.Errorf("unexpected close error: %q", tp.closeErr) 473 } 474 conn.Close() 475 } 476 477 func TestServerSetupConn(t *testing.T) { 478 var ( 479 clientkey, srvkey = newkey(), newkey() 480 clientpub = &clientkey.PublicKey 481 srvpub = &srvkey.PublicKey 482 ) 483 tests := []struct { 484 dontstart bool 485 tt *setupTransport 486 flags connFlag 487 dialDest *onode.Node 488 489 wantCloseErr error 490 wantCalls string 491 }{ 492 { 493 dontstart: true, 494 tt: &setupTransport{pubkey: clientpub}, 495 wantCalls: "close,", 496 wantCloseErr: errServerStopped, 497 }, 498 { 499 tt: &setupTransport{pubkey: clientpub, encHandshakeErr: errors.New("read error")}, 500 flags: inboundConn, 501 wantCalls: "doEncHandshake,close,", 502 wantCloseErr: errors.New("read error"), 503 }, 504 { 505 tt: &setupTransport{pubkey: clientpub}, 506 dialDest: onode.NewV4(&newkey().PublicKey, nil, 0, 0), 507 flags: dynDialedConn, 508 wantCalls: "doEncHandshake,close,", 509 wantCloseErr: DiscUnexpectedIdentity, 510 }, 511 { 512 tt: &setupTransport{pubkey: clientpub, phs: ProtoHandshake{ID: randomID().Bytes()}}, 513 dialDest: onode.NewV4(clientpub, nil, 0, 0), 514 flags: dynDialedConn, 515 wantCalls: "doEncHandshake,doProtoHandshake,close,", 516 wantCloseErr: DiscUnexpectedIdentity, 517 }, 518 { 519 tt: &setupTransport{pubkey: clientpub, protoHandshakeErr: errors.New("foo")}, 520 dialDest: onode.NewV4(clientpub, nil, 0, 0), 521 flags: dynDialedConn, 522 wantCalls: "doEncHandshake,doProtoHandshake,close,", 523 wantCloseErr: errors.New("foo"), 524 }, 525 { 526 tt: &setupTransport{pubkey: srvpub, phs: ProtoHandshake{ID: ogcrypto2.FromECDSAPub(srvpub)[1:]}}, 527 flags: inboundConn, 528 wantCalls: "doEncHandshake,close,", 529 wantCloseErr: DiscSelf, 530 }, 531 { 532 tt: &setupTransport{pubkey: clientpub, phs: ProtoHandshake{ID: ogcrypto2.FromECDSAPub(clientpub)[1:]}}, 533 flags: inboundConn, 534 wantCalls: "doEncHandshake,doProtoHandshake,close,", 535 wantCloseErr: DiscUselessPeer, 536 }, 537 } 538 539 for i, test := range tests { 540 srv := &Server{ 541 Config: Config{ 542 PrivateKey: srvkey, 543 MaxPeers: 10, 544 NoDial: true, 545 Protocols: []Protocol{discard}, 546 }, 547 newTransport: func(fd net.Conn) transport { return test.tt }, 548 } 549 if !test.dontstart { 550 if err := srv.start(); err != nil { 551 t.Fatalf("couldn't start server: %v", err) 552 } 553 } 554 p1, _ := net.Pipe() 555 srv.SetupConn(p1, test.flags, test.dialDest) 556 if !reflect.DeepEqual(test.tt.closeErr, test.wantCloseErr) { 557 t.Errorf("test %d: close error mismatch: got %q, want %q", i, test.tt.closeErr, test.wantCloseErr) 558 } 559 if test.tt.calls != test.wantCalls { 560 t.Errorf("test %d: calls mismatch: got %q, want %q", i, test.tt.calls, test.wantCalls) 561 } 562 } 563 } 564 565 type setupTransport struct { 566 pubkey *ecdsa.PublicKey 567 encHandshakeErr error 568 phs ProtoHandshake 569 protoHandshakeErr error 570 571 calls string 572 closeErr error 573 } 574 575 func (c *setupTransport) doEncHandshake(prv *ecdsa.PrivateKey, dialDest *ecdsa.PublicKey) (*ecdsa.PublicKey, error) { 576 c.calls += "doEncHandshake," 577 return c.pubkey, c.encHandshakeErr 578 } 579 580 func (c *setupTransport) doProtoHandshake(our *ProtoHandshake) (*ProtoHandshake, error) { 581 c.calls += "doProtoHandshake," 582 if c.protoHandshakeErr != nil { 583 return nil, c.protoHandshakeErr 584 } 585 return &c.phs, nil 586 } 587 func (c *setupTransport) close(err error) { 588 c.calls += "close," 589 c.closeErr = err 590 } 591 592 // setupConn shouldn't write to/read from the connection. 593 func (c *setupTransport) WriteMsg(Msg) error { 594 panic("WriteMsg called on setupTransport") 595 } 596 func (c *setupTransport) ReadMsg() (Msg, error) { 597 panic("ReadMsg called on setupTransport") 598 } 599 600 func newkey() *ecdsa.PrivateKey { 601 key, err := ogcrypto2.GenerateKey() 602 if err != nil { 603 panic("couldn't generate key: " + err.Error()) 604 } 605 return key 606 } 607 608 func randomID() (id onode.ID) { 609 for i := range id { 610 id[i] = byte(rand.Intn(255)) 611 } 612 return id 613 }