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