github.com/Blockdaemon/celo-blockchain@v0.0.0-20200129231733-e667f6b08419/p2p/protocols/protocol_test.go (about) 1 // Copyright 2017 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 protocols 18 19 import ( 20 "bytes" 21 "context" 22 "errors" 23 "fmt" 24 "testing" 25 "time" 26 27 "github.com/ethereum/go-ethereum/p2p" 28 "github.com/ethereum/go-ethereum/p2p/enode" 29 "github.com/ethereum/go-ethereum/p2p/simulations/adapters" 30 p2ptest "github.com/ethereum/go-ethereum/p2p/testing" 31 "github.com/ethereum/go-ethereum/rlp" 32 ) 33 34 // handshake message type 35 type hs0 struct { 36 C uint 37 } 38 39 // message to kill/drop the peer with nodeID 40 type kill struct { 41 C enode.ID 42 } 43 44 // message to drop connection 45 type drop struct { 46 } 47 48 /// protoHandshake represents module-independent aspects of the protocol and is 49 // the first message peers send and receive as part the initial exchange 50 type protoHandshake struct { 51 Version uint // local and remote peer should have identical version 52 NetworkID string // local and remote peer should have identical network id 53 } 54 55 // checkProtoHandshake verifies local and remote protoHandshakes match 56 func checkProtoHandshake(testVersion uint, testNetworkID string) func(interface{}) error { 57 return func(rhs interface{}) error { 58 remote := rhs.(*protoHandshake) 59 if remote.NetworkID != testNetworkID { 60 return fmt.Errorf("%s (!= %s)", remote.NetworkID, testNetworkID) 61 } 62 63 if remote.Version != testVersion { 64 return fmt.Errorf("%d (!= %d)", remote.Version, testVersion) 65 } 66 return nil 67 } 68 } 69 70 // newProtocol sets up a protocol 71 // the run function here demonstrates a typical protocol using peerPool, handshake 72 // and messages registered to handlers 73 func newProtocol(pp *p2ptest.TestPeerPool) func(*p2p.Peer, p2p.MsgReadWriter) error { 74 spec := &Spec{ 75 Name: "test", 76 Version: 42, 77 MaxMsgSize: 10 * 1024, 78 Messages: []interface{}{ 79 protoHandshake{}, 80 hs0{}, 81 kill{}, 82 drop{}, 83 }, 84 } 85 return func(p *p2p.Peer, rw p2p.MsgReadWriter) error { 86 peer := NewPeer(p, rw, spec) 87 88 // initiate one-off protohandshake and check validity 89 ctx, cancel := context.WithTimeout(context.Background(), time.Second) 90 defer cancel() 91 phs := &protoHandshake{42, "420"} 92 hsCheck := checkProtoHandshake(phs.Version, phs.NetworkID) 93 _, err := peer.Handshake(ctx, phs, hsCheck) 94 if err != nil { 95 return err 96 } 97 98 lhs := &hs0{42} 99 // module handshake demonstrating a simple repeatable exchange of same-type message 100 hs, err := peer.Handshake(ctx, lhs, nil) 101 if err != nil { 102 return err 103 } 104 105 if rmhs := hs.(*hs0); rmhs.C > lhs.C { 106 return fmt.Errorf("handshake mismatch remote %v > local %v", rmhs.C, lhs.C) 107 } 108 109 handle := func(ctx context.Context, msg interface{}) error { 110 switch msg := msg.(type) { 111 112 case *protoHandshake: 113 return errors.New("duplicate handshake") 114 115 case *hs0: 116 rhs := msg 117 if rhs.C > lhs.C { 118 return fmt.Errorf("handshake mismatch remote %v > local %v", rhs.C, lhs.C) 119 } 120 lhs.C += rhs.C 121 return peer.Send(ctx, lhs) 122 123 case *kill: 124 // demonstrates use of peerPool, killing another peer connection as a response to a message 125 id := msg.C 126 pp.Get(id).Drop(errors.New("killed")) 127 return nil 128 129 case *drop: 130 // for testing we can trigger self induced disconnect upon receiving drop message 131 return errors.New("dropped") 132 133 default: 134 return fmt.Errorf("unknown message type: %T", msg) 135 } 136 } 137 138 pp.Add(peer) 139 defer pp.Remove(peer) 140 return peer.Run(handle) 141 } 142 } 143 144 func protocolTester(pp *p2ptest.TestPeerPool) *p2ptest.ProtocolTester { 145 conf := adapters.RandomNodeConfig() 146 return p2ptest.NewProtocolTester(conf.ID, 2, newProtocol(pp)) 147 } 148 149 func protoHandshakeExchange(id enode.ID, proto *protoHandshake) []p2ptest.Exchange { 150 151 return []p2ptest.Exchange{ 152 { 153 Expects: []p2ptest.Expect{ 154 { 155 Code: 0, 156 Msg: &protoHandshake{42, "420"}, 157 Peer: id, 158 }, 159 }, 160 }, 161 { 162 Triggers: []p2ptest.Trigger{ 163 { 164 Code: 0, 165 Msg: proto, 166 Peer: id, 167 }, 168 }, 169 }, 170 } 171 } 172 173 func runProtoHandshake(t *testing.T, proto *protoHandshake, errs ...error) { 174 pp := p2ptest.NewTestPeerPool() 175 s := protocolTester(pp) 176 // TODO: make this more than one handshake 177 node := s.Nodes[0] 178 if err := s.TestExchanges(protoHandshakeExchange(node.ID(), proto)...); err != nil { 179 t.Fatal(err) 180 } 181 var disconnects []*p2ptest.Disconnect 182 for i, err := range errs { 183 disconnects = append(disconnects, &p2ptest.Disconnect{Peer: s.Nodes[i].ID(), Error: err}) 184 } 185 if err := s.TestDisconnected(disconnects...); err != nil { 186 t.Fatal(err) 187 } 188 } 189 190 type dummyHook struct { 191 peer *Peer 192 size uint32 193 msg interface{} 194 send bool 195 err error 196 waitC chan struct{} 197 } 198 199 type dummyMsg struct { 200 Content string 201 } 202 203 func (d *dummyHook) Send(peer *Peer, size uint32, msg interface{}) error { 204 d.peer = peer 205 d.size = size 206 d.msg = msg 207 d.send = true 208 return d.err 209 } 210 211 func (d *dummyHook) Receive(peer *Peer, size uint32, msg interface{}) error { 212 d.peer = peer 213 d.size = size 214 d.msg = msg 215 d.send = false 216 d.waitC <- struct{}{} 217 return d.err 218 } 219 220 func TestProtocolHook(t *testing.T) { 221 testHook := &dummyHook{ 222 waitC: make(chan struct{}, 1), 223 } 224 spec := &Spec{ 225 Name: "test", 226 Version: 42, 227 MaxMsgSize: 10 * 1024, 228 Messages: []interface{}{ 229 dummyMsg{}, 230 }, 231 Hook: testHook, 232 } 233 234 runFunc := func(p *p2p.Peer, rw p2p.MsgReadWriter) error { 235 peer := NewPeer(p, rw, spec) 236 ctx := context.TODO() 237 err := peer.Send(ctx, &dummyMsg{ 238 Content: "handshake"}) 239 240 if err != nil { 241 t.Fatal(err) 242 } 243 244 handle := func(ctx context.Context, msg interface{}) error { 245 return nil 246 } 247 248 return peer.Run(handle) 249 } 250 251 conf := adapters.RandomNodeConfig() 252 tester := p2ptest.NewProtocolTester(conf.ID, 2, runFunc) 253 err := tester.TestExchanges(p2ptest.Exchange{ 254 Expects: []p2ptest.Expect{ 255 { 256 Code: 0, 257 Msg: &dummyMsg{Content: "handshake"}, 258 Peer: tester.Nodes[0].ID(), 259 }, 260 }, 261 }) 262 if err != nil { 263 t.Fatal(err) 264 } 265 if testHook.msg == nil || testHook.msg.(*dummyMsg).Content != "handshake" { 266 t.Fatal("Expected msg to be set, but it is not") 267 } 268 if !testHook.send { 269 t.Fatal("Expected a send message, but it is not") 270 } 271 if testHook.peer == nil { 272 t.Fatal("Expected peer ID to be set correctly, but it is not (it is nil)") 273 } else if testHook.peer.ID() != tester.Nodes[0].ID() { 274 t.Fatalf("Expected peer ID to be set correctly, but it is not (%v != %v)", testHook.peer.ID(), tester.Nodes[0].ID()) 275 } 276 if testHook.size != 11 { //11 is the length of the encoded message 277 t.Fatalf("Expected size to be %d, but it is %d ", 1, testHook.size) 278 } 279 280 err = tester.TestExchanges(p2ptest.Exchange{ 281 Triggers: []p2ptest.Trigger{ 282 { 283 Code: 0, 284 Msg: &dummyMsg{Content: "response"}, 285 Peer: tester.Nodes[1].ID(), 286 }, 287 }, 288 }) 289 290 <-testHook.waitC 291 292 if err != nil { 293 t.Fatal(err) 294 } 295 if testHook.msg == nil || testHook.msg.(*dummyMsg).Content != "response" { 296 t.Fatal("Expected msg to be set, but it is not") 297 } 298 if testHook.send { 299 t.Fatal("Expected a send message, but it is not") 300 } 301 if testHook.peer == nil || testHook.peer.ID() != tester.Nodes[1].ID() { 302 t.Fatal("Expected peer ID to be set correctly, but it is not") 303 } 304 if testHook.size != 10 { //11 is the length of the encoded message 305 t.Fatalf("Expected size to be %d, but it is %d ", 1, testHook.size) 306 } 307 308 testHook.err = fmt.Errorf("dummy error") 309 err = tester.TestExchanges(p2ptest.Exchange{ 310 Triggers: []p2ptest.Trigger{ 311 { 312 Code: 0, 313 Msg: &dummyMsg{Content: "response"}, 314 Peer: tester.Nodes[1].ID(), 315 }, 316 }, 317 }) 318 319 <-testHook.waitC 320 321 time.Sleep(100 * time.Millisecond) 322 err = tester.TestDisconnected(&p2ptest.Disconnect{Peer: tester.Nodes[1].ID(), Error: testHook.err}) 323 if err != nil { 324 t.Fatalf("Expected a specific disconnect error, but got different one: %v", err) 325 } 326 327 } 328 329 //We need to test that if the hook is not defined, then message infrastructure 330 //(send,receive) still works 331 func TestNoHook(t *testing.T) { 332 //create a test spec 333 spec := createTestSpec() 334 //a random node 335 id := adapters.RandomNodeConfig().ID 336 //a peer 337 p := p2p.NewPeer(id, "testPeer", nil) 338 rw := &dummyRW{} 339 peer := NewPeer(p, rw, spec) 340 ctx := context.TODO() 341 msg := &perBytesMsgSenderPays{Content: "testBalance"} 342 //send a message 343 err := peer.Send(ctx, msg) 344 if err != nil { 345 t.Fatal(err) 346 } 347 //simulate receiving a message 348 rw.msg = msg 349 peer.handleIncoming(func(ctx context.Context, msg interface{}) error { 350 return nil 351 }) 352 //all should just work and not result in any error 353 } 354 355 func TestProtoHandshakeVersionMismatch(t *testing.T) { 356 runProtoHandshake(t, &protoHandshake{41, "420"}, errorf(ErrHandshake, errorf(ErrHandler, "(msg code 0): 41 (!= 42)").Error())) 357 } 358 359 func TestProtoHandshakeNetworkIDMismatch(t *testing.T) { 360 runProtoHandshake(t, &protoHandshake{42, "421"}, errorf(ErrHandshake, errorf(ErrHandler, "(msg code 0): 421 (!= 420)").Error())) 361 } 362 363 func TestProtoHandshakeSuccess(t *testing.T) { 364 runProtoHandshake(t, &protoHandshake{42, "420"}) 365 } 366 367 func moduleHandshakeExchange(id enode.ID, resp uint) []p2ptest.Exchange { 368 369 return []p2ptest.Exchange{ 370 { 371 Expects: []p2ptest.Expect{ 372 { 373 Code: 1, 374 Msg: &hs0{42}, 375 Peer: id, 376 }, 377 }, 378 }, 379 { 380 Triggers: []p2ptest.Trigger{ 381 { 382 Code: 1, 383 Msg: &hs0{resp}, 384 Peer: id, 385 }, 386 }, 387 }, 388 } 389 } 390 391 func runModuleHandshake(t *testing.T, resp uint, errs ...error) { 392 pp := p2ptest.NewTestPeerPool() 393 s := protocolTester(pp) 394 node := s.Nodes[0] 395 if err := s.TestExchanges(protoHandshakeExchange(node.ID(), &protoHandshake{42, "420"})...); err != nil { 396 t.Fatal(err) 397 } 398 if err := s.TestExchanges(moduleHandshakeExchange(node.ID(), resp)...); err != nil { 399 t.Fatal(err) 400 } 401 var disconnects []*p2ptest.Disconnect 402 for i, err := range errs { 403 disconnects = append(disconnects, &p2ptest.Disconnect{Peer: s.Nodes[i].ID(), Error: err}) 404 } 405 if err := s.TestDisconnected(disconnects...); err != nil { 406 t.Fatal(err) 407 } 408 } 409 410 func TestModuleHandshakeError(t *testing.T) { 411 runModuleHandshake(t, 43, fmt.Errorf("handshake mismatch remote 43 > local 42")) 412 } 413 414 func TestModuleHandshakeSuccess(t *testing.T) { 415 runModuleHandshake(t, 42) 416 } 417 418 // testing complex interactions over multiple peers, relaying, dropping 419 func testMultiPeerSetup(a, b enode.ID) []p2ptest.Exchange { 420 421 return []p2ptest.Exchange{ 422 { 423 Label: "primary handshake", 424 Expects: []p2ptest.Expect{ 425 { 426 Code: 0, 427 Msg: &protoHandshake{42, "420"}, 428 Peer: a, 429 }, 430 { 431 Code: 0, 432 Msg: &protoHandshake{42, "420"}, 433 Peer: b, 434 }, 435 }, 436 }, 437 { 438 Label: "module handshake", 439 Triggers: []p2ptest.Trigger{ 440 { 441 Code: 0, 442 Msg: &protoHandshake{42, "420"}, 443 Peer: a, 444 }, 445 { 446 Code: 0, 447 Msg: &protoHandshake{42, "420"}, 448 Peer: b, 449 }, 450 }, 451 Expects: []p2ptest.Expect{ 452 { 453 Code: 1, 454 Msg: &hs0{42}, 455 Peer: a, 456 }, 457 { 458 Code: 1, 459 Msg: &hs0{42}, 460 Peer: b, 461 }, 462 }, 463 }, 464 465 {Label: "alternative module handshake", Triggers: []p2ptest.Trigger{{Code: 1, Msg: &hs0{41}, Peer: a}, 466 {Code: 1, Msg: &hs0{41}, Peer: b}}}, 467 {Label: "repeated module handshake", Triggers: []p2ptest.Trigger{{Code: 1, Msg: &hs0{1}, Peer: a}}}, 468 {Label: "receiving repeated module handshake", Expects: []p2ptest.Expect{{Code: 1, Msg: &hs0{43}, Peer: a}}}} 469 } 470 471 func runMultiplePeers(t *testing.T, peer int, errs ...error) { 472 pp := p2ptest.NewTestPeerPool() 473 s := protocolTester(pp) 474 475 if err := s.TestExchanges(testMultiPeerSetup(s.Nodes[0].ID(), s.Nodes[1].ID())...); err != nil { 476 t.Fatal(err) 477 } 478 // after some exchanges of messages, we can test state changes 479 // here this is simply demonstrated by the peerPool 480 // after the handshake negotiations peers must be added to the pool 481 // time.Sleep(1) 482 tick := time.NewTicker(10 * time.Millisecond) 483 timeout := time.NewTimer(1 * time.Second) 484 WAIT: 485 for { 486 select { 487 case <-tick.C: 488 if pp.Has(s.Nodes[0].ID()) { 489 break WAIT 490 } 491 case <-timeout.C: 492 t.Fatal("timeout") 493 } 494 } 495 if !pp.Has(s.Nodes[1].ID()) { 496 t.Fatalf("missing peer test-1: %v (%v)", pp, s.Nodes) 497 } 498 499 // peer 0 sends kill request for peer with index <peer> 500 err := s.TestExchanges(p2ptest.Exchange{ 501 Triggers: []p2ptest.Trigger{ 502 { 503 Code: 2, 504 Msg: &kill{s.Nodes[peer].ID()}, 505 Peer: s.Nodes[0].ID(), 506 }, 507 }, 508 }) 509 510 if err != nil { 511 t.Fatal(err) 512 } 513 514 // the peer not killed sends a drop request 515 err = s.TestExchanges(p2ptest.Exchange{ 516 Triggers: []p2ptest.Trigger{ 517 { 518 Code: 3, 519 Msg: &drop{}, 520 Peer: s.Nodes[(peer+1)%2].ID(), 521 }, 522 }, 523 }) 524 525 if err != nil { 526 t.Fatal(err) 527 } 528 529 // check the actual discconnect errors on the individual peers 530 var disconnects []*p2ptest.Disconnect 531 for i, err := range errs { 532 disconnects = append(disconnects, &p2ptest.Disconnect{Peer: s.Nodes[i].ID(), Error: err}) 533 } 534 if err := s.TestDisconnected(disconnects...); err != nil { 535 t.Fatal(err) 536 } 537 // test if disconnected peers have been removed from peerPool 538 if pp.Has(s.Nodes[peer].ID()) { 539 t.Fatalf("peer test-%v not dropped: %v (%v)", peer, pp, s.Nodes) 540 } 541 542 } 543 func XTestMultiplePeersDropSelf(t *testing.T) { 544 runMultiplePeers(t, 0, 545 fmt.Errorf("subprotocol error"), 546 fmt.Errorf("Message handler error: (msg code 3): dropped"), 547 ) 548 } 549 550 func XTestMultiplePeersDropOther(t *testing.T) { 551 runMultiplePeers(t, 1, 552 fmt.Errorf("Message handler error: (msg code 3): dropped"), 553 fmt.Errorf("subprotocol error"), 554 ) 555 } 556 557 //dummy implementation of a MsgReadWriter 558 //this allows for quick and easy unit tests without 559 //having to build up the complete protocol 560 type dummyRW struct { 561 msg interface{} 562 size uint32 563 code uint64 564 } 565 566 func (d *dummyRW) WriteMsg(msg p2p.Msg) error { 567 return nil 568 } 569 570 func (d *dummyRW) ReadMsg() (p2p.Msg, error) { 571 enc := bytes.NewReader(d.getDummyMsg()) 572 return p2p.Msg{ 573 Code: d.code, 574 Size: d.size, 575 Payload: enc, 576 ReceivedAt: time.Now(), 577 }, nil 578 } 579 580 func (d *dummyRW) getDummyMsg() []byte { 581 r, _ := rlp.EncodeToBytes(d.msg) 582 var b bytes.Buffer 583 wmsg := WrappedMsg{ 584 Context: b.Bytes(), 585 Size: uint32(len(r)), 586 Payload: r, 587 } 588 rr, _ := rlp.EncodeToBytes(wmsg) 589 d.size = uint32(len(rr)) 590 return rr 591 }