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  }