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  }