github.com/JFJun/bsc@v1.0.0/les/peer_test.go (about) 1 // Copyright 2019 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 les 18 19 import ( 20 "crypto/rand" 21 "math/big" 22 "reflect" 23 "sort" 24 "testing" 25 "time" 26 27 "github.com/JFJun/bsc/common" 28 "github.com/JFJun/bsc/p2p" 29 "github.com/JFJun/bsc/p2p/enode" 30 ) 31 32 type testServerPeerSub struct { 33 regCh chan *serverPeer 34 unregCh chan *serverPeer 35 } 36 37 func newTestServerPeerSub() *testServerPeerSub { 38 return &testServerPeerSub{ 39 regCh: make(chan *serverPeer, 1), 40 unregCh: make(chan *serverPeer, 1), 41 } 42 } 43 44 func (t *testServerPeerSub) registerPeer(p *serverPeer) { t.regCh <- p } 45 func (t *testServerPeerSub) unregisterPeer(p *serverPeer) { t.unregCh <- p } 46 47 func TestPeerSubscription(t *testing.T) { 48 peers := newServerPeerSet() 49 defer peers.close() 50 51 checkIds := func(expect []string) { 52 given := peers.ids() 53 if len(given) == 0 && len(expect) == 0 { 54 return 55 } 56 sort.Strings(given) 57 sort.Strings(expect) 58 if !reflect.DeepEqual(given, expect) { 59 t.Fatalf("all peer ids mismatch, want %v, given %v", expect, given) 60 } 61 } 62 checkPeers := func(peerCh chan *serverPeer) { 63 select { 64 case <-peerCh: 65 case <-time.NewTimer(100 * time.Millisecond).C: 66 t.Fatalf("timeout, no event received") 67 } 68 select { 69 case <-peerCh: 70 t.Fatalf("unexpected event received") 71 case <-time.NewTimer(10 * time.Millisecond).C: 72 } 73 } 74 checkIds([]string{}) 75 76 sub := newTestServerPeerSub() 77 peers.subscribe(sub) 78 79 // Generate a random id and create the peer 80 var id enode.ID 81 rand.Read(id[:]) 82 peer := newServerPeer(2, NetworkId, false, p2p.NewPeer(id, "name", nil), nil) 83 peers.register(peer) 84 85 checkIds([]string{peer.id}) 86 checkPeers(sub.regCh) 87 88 peers.unregister(peer.id) 89 checkIds([]string{}) 90 checkPeers(sub.unregCh) 91 } 92 93 func TestHandshake(t *testing.T) { 94 // Create a message pipe to communicate through 95 app, net := p2p.MsgPipe() 96 97 // Generate a random id and create the peer 98 var id enode.ID 99 rand.Read(id[:]) 100 101 peer1 := newClientPeer(2, NetworkId, p2p.NewPeer(id, "name", nil), net) 102 peer2 := newServerPeer(2, NetworkId, true, p2p.NewPeer(id, "name", nil), app) 103 104 var ( 105 errCh1 = make(chan error, 1) 106 errCh2 = make(chan error, 1) 107 108 td = big.NewInt(100) 109 head = common.HexToHash("deadbeef") 110 headNum = uint64(10) 111 genesis = common.HexToHash("cafebabe") 112 ) 113 go func() { 114 errCh1 <- peer1.handshake(td, head, headNum, genesis, func(list *keyValueList) { 115 var announceType uint64 = announceTypeSigned 116 *list = (*list).add("announceType", announceType) 117 }, nil) 118 }() 119 go func() { 120 errCh2 <- peer2.handshake(td, head, headNum, genesis, nil, func(recv keyValueMap) error { 121 var reqType uint64 122 err := recv.get("announceType", &reqType) 123 if err != nil { 124 t.Fatal(err) 125 } 126 if reqType != announceTypeSigned { 127 t.Fatal("Expected announceTypeSigned") 128 } 129 return nil 130 }) 131 }() 132 133 for i := 0; i < 2; i++ { 134 select { 135 case err := <-errCh1: 136 if err != nil { 137 t.Fatalf("handshake failed, %v", err) 138 } 139 case err := <-errCh2: 140 if err != nil { 141 t.Fatalf("handshake failed, %v", err) 142 } 143 case <-time.After(time.Second): 144 t.Fatalf("timeout") 145 } 146 } 147 }