github.com/daeglee/go-ethereum@v0.0.0-20190504220456-cad3e8d18e9b/p2p/testing/protocoltester.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  /*
    18  the p2p/testing package provides a unit test scheme to check simple
    19  protocol message exchanges with one pivot node and a number of dummy peers
    20  The pivot test node runs a node.Service, the dummy peers run a mock node
    21  that can be used to send and receive messages
    22  */
    23  
    24  package testing
    25  
    26  import (
    27  	"bytes"
    28  	"crypto/ecdsa"
    29  	"fmt"
    30  	"io"
    31  	"io/ioutil"
    32  	"strings"
    33  	"sync"
    34  
    35  	"github.com/ethereum/go-ethereum/log"
    36  	"github.com/ethereum/go-ethereum/node"
    37  	"github.com/ethereum/go-ethereum/p2p"
    38  	"github.com/ethereum/go-ethereum/p2p/enode"
    39  	"github.com/ethereum/go-ethereum/p2p/simulations"
    40  	"github.com/ethereum/go-ethereum/p2p/simulations/adapters"
    41  	"github.com/ethereum/go-ethereum/rlp"
    42  	"github.com/ethereum/go-ethereum/rpc"
    43  )
    44  
    45  // ProtocolTester is the tester environment used for unit testing protocol
    46  // message exchanges. It uses p2p/simulations framework
    47  type ProtocolTester struct {
    48  	*ProtocolSession
    49  	network *simulations.Network
    50  }
    51  
    52  // NewProtocolTester constructs a new ProtocolTester
    53  // it takes as argument the pivot node id, the number of dummy peers and the
    54  // protocol run function called on a peer connection by the p2p server
    55  func NewProtocolTester(prvkey *ecdsa.PrivateKey, nodeCount int, run func(*p2p.Peer, p2p.MsgReadWriter) error) *ProtocolTester {
    56  	services := adapters.Services{
    57  		"test": func(ctx *adapters.ServiceContext) (node.Service, error) {
    58  			return &testNode{run}, nil
    59  		},
    60  		"mock": func(ctx *adapters.ServiceContext) (node.Service, error) {
    61  			return newMockNode(), nil
    62  		},
    63  	}
    64  	adapter := adapters.NewSimAdapter(services)
    65  	net := simulations.NewNetwork(adapter, &simulations.NetworkConfig{})
    66  	nodeConfig := &adapters.NodeConfig{
    67  		PrivateKey:      prvkey,
    68  		EnableMsgEvents: true,
    69  		Services:        []string{"test"},
    70  	}
    71  	if _, err := net.NewNodeWithConfig(nodeConfig); err != nil {
    72  		panic(err.Error())
    73  	}
    74  	if err := net.Start(nodeConfig.ID); err != nil {
    75  		panic(err.Error())
    76  	}
    77  
    78  	node := net.GetNode(nodeConfig.ID).Node.(*adapters.SimNode)
    79  	peers := make([]*adapters.NodeConfig, nodeCount)
    80  	nodes := make([]*enode.Node, nodeCount)
    81  	for i := 0; i < nodeCount; i++ {
    82  		peers[i] = adapters.RandomNodeConfig()
    83  		peers[i].Services = []string{"mock"}
    84  		if _, err := net.NewNodeWithConfig(peers[i]); err != nil {
    85  			panic(fmt.Sprintf("error initializing peer %v: %v", peers[i].ID, err))
    86  		}
    87  		if err := net.Start(peers[i].ID); err != nil {
    88  			panic(fmt.Sprintf("error starting peer %v: %v", peers[i].ID, err))
    89  		}
    90  		nodes[i] = peers[i].Node()
    91  	}
    92  	events := make(chan *p2p.PeerEvent, 1000)
    93  	node.SubscribeEvents(events)
    94  	ps := &ProtocolSession{
    95  		Server:  node.Server(),
    96  		Nodes:   nodes,
    97  		adapter: adapter,
    98  		events:  events,
    99  	}
   100  	self := &ProtocolTester{
   101  		ProtocolSession: ps,
   102  		network:         net,
   103  	}
   104  
   105  	self.Connect(nodeConfig.ID, peers...)
   106  
   107  	return self
   108  }
   109  
   110  // Stop stops the p2p server
   111  func (t *ProtocolTester) Stop() {
   112  	t.Server.Stop()
   113  }
   114  
   115  // Connect brings up the remote peer node and connects it using the
   116  // p2p/simulations network connection with the in memory network adapter
   117  func (t *ProtocolTester) Connect(selfID enode.ID, peers ...*adapters.NodeConfig) {
   118  	for _, peer := range peers {
   119  		log.Trace(fmt.Sprintf("connect to %v", peer.ID))
   120  		if err := t.network.Connect(selfID, peer.ID); err != nil {
   121  			panic(fmt.Sprintf("error connecting to peer %v: %v", peer.ID, err))
   122  		}
   123  	}
   124  
   125  }
   126  
   127  // testNode wraps a protocol run function and implements the node.Service
   128  // interface
   129  type testNode struct {
   130  	run func(*p2p.Peer, p2p.MsgReadWriter) error
   131  }
   132  
   133  func (t *testNode) Protocols() []p2p.Protocol {
   134  	return []p2p.Protocol{{
   135  		Length: 100,
   136  		Run:    t.run,
   137  	}}
   138  }
   139  
   140  func (t *testNode) APIs() []rpc.API {
   141  	return nil
   142  }
   143  
   144  func (t *testNode) Start(server *p2p.Server) error {
   145  	return nil
   146  }
   147  
   148  func (t *testNode) Stop() error {
   149  	return nil
   150  }
   151  
   152  // mockNode is a testNode which doesn't actually run a protocol, instead
   153  // exposing channels so that tests can manually trigger and expect certain
   154  // messages
   155  type mockNode struct {
   156  	testNode
   157  
   158  	trigger  chan *Trigger
   159  	expect   chan []Expect
   160  	err      chan error
   161  	stop     chan struct{}
   162  	stopOnce sync.Once
   163  }
   164  
   165  func newMockNode() *mockNode {
   166  	mock := &mockNode{
   167  		trigger: make(chan *Trigger),
   168  		expect:  make(chan []Expect),
   169  		err:     make(chan error),
   170  		stop:    make(chan struct{}),
   171  	}
   172  	mock.testNode.run = mock.Run
   173  	return mock
   174  }
   175  
   176  // Run is a protocol run function which just loops waiting for tests to
   177  // instruct it to either trigger or expect a message from the peer
   178  func (m *mockNode) Run(peer *p2p.Peer, rw p2p.MsgReadWriter) error {
   179  	for {
   180  		select {
   181  		case trig := <-m.trigger:
   182  			wmsg := Wrap(trig.Msg)
   183  			m.err <- p2p.Send(rw, trig.Code, wmsg)
   184  		case exps := <-m.expect:
   185  			m.err <- expectMsgs(rw, exps)
   186  		case <-m.stop:
   187  			return nil
   188  		}
   189  	}
   190  }
   191  
   192  func (m *mockNode) Trigger(trig *Trigger) error {
   193  	m.trigger <- trig
   194  	return <-m.err
   195  }
   196  
   197  func (m *mockNode) Expect(exp ...Expect) error {
   198  	m.expect <- exp
   199  	return <-m.err
   200  }
   201  
   202  func (m *mockNode) Stop() error {
   203  	m.stopOnce.Do(func() { close(m.stop) })
   204  	return nil
   205  }
   206  
   207  func expectMsgs(rw p2p.MsgReadWriter, exps []Expect) error {
   208  	matched := make([]bool, len(exps))
   209  	for {
   210  		msg, err := rw.ReadMsg()
   211  		if err != nil {
   212  			if err == io.EOF {
   213  				break
   214  			}
   215  			return err
   216  		}
   217  		actualContent, err := ioutil.ReadAll(msg.Payload)
   218  		if err != nil {
   219  			return err
   220  		}
   221  		var found bool
   222  		for i, exp := range exps {
   223  			if exp.Code == msg.Code && bytes.Equal(actualContent, mustEncodeMsg(Wrap(exp.Msg))) {
   224  				if matched[i] {
   225  					return fmt.Errorf("message #%d received two times", i)
   226  				}
   227  				matched[i] = true
   228  				found = true
   229  				break
   230  			}
   231  		}
   232  		if !found {
   233  			expected := make([]string, 0)
   234  			for i, exp := range exps {
   235  				if matched[i] {
   236  					continue
   237  				}
   238  				expected = append(expected, fmt.Sprintf("code %d payload %x", exp.Code, mustEncodeMsg(Wrap(exp.Msg))))
   239  			}
   240  			return fmt.Errorf("unexpected message code %d payload %x, expected %s", msg.Code, actualContent, strings.Join(expected, " or "))
   241  		}
   242  		done := true
   243  		for _, m := range matched {
   244  			if !m {
   245  				done = false
   246  				break
   247  			}
   248  		}
   249  		if done {
   250  			return nil
   251  		}
   252  	}
   253  	for i, m := range matched {
   254  		if !m {
   255  			return fmt.Errorf("expected message #%d not received", i)
   256  		}
   257  	}
   258  	return nil
   259  }
   260  
   261  // mustEncodeMsg uses rlp to encode a message.
   262  // In case of error it panics.
   263  func mustEncodeMsg(msg interface{}) []byte {
   264  	contentEnc, err := rlp.EncodeToBytes(msg)
   265  	if err != nil {
   266  		panic("content encode error: " + err.Error())
   267  	}
   268  	return contentEnc
   269  }
   270  
   271  type WrappedMsg struct {
   272  	Context []byte
   273  	Size    uint32
   274  	Payload []byte
   275  }
   276  
   277  func Wrap(msg interface{}) interface{} {
   278  	data, _ := rlp.EncodeToBytes(msg)
   279  	return &WrappedMsg{
   280  		Size:    uint32(len(data)),
   281  		Payload: data,
   282  	}
   283  }