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