go-ethereum/p2p/testing/protocoltester.go

283 lines
7.2 KiB
Go

// Copyright 2017 The go-ethereum Authors
// This file is part of the go-ethereum library.
//
// The go-ethereum library is free software: you can redistribute it and/or modify
// it under the terms of the GNU Lesser General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// The go-ethereum library is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
// GNU Lesser General Public License for more details.
//
// You should have received a copy of the GNU Lesser General Public License
// along with the go-ethereum library. If not, see <http://www.gnu.org/licenses/>.
/*
the p2p/testing package provides a unit test scheme to check simple
protocol message exchanges with one pivot node and a number of dummy peers
The pivot test node runs a node.Service, the dummy peers run a mock node
that can be used to send and receive messages
*/
package testing
import (
"bytes"
"fmt"
"io"
"io/ioutil"
"strings"
"sync"
"github.com/ethereum/go-ethereum/log"
"github.com/ethereum/go-ethereum/node"
"github.com/ethereum/go-ethereum/p2p"
"github.com/ethereum/go-ethereum/p2p/enode"
"github.com/ethereum/go-ethereum/p2p/simulations"
"github.com/ethereum/go-ethereum/p2p/simulations/adapters"
"github.com/ethereum/go-ethereum/rlp"
"github.com/ethereum/go-ethereum/rpc"
)
// ProtocolTester is the tester environment used for unit testing protocol
// message exchanges. It uses p2p/simulations framework
type ProtocolTester struct {
*ProtocolSession
network *simulations.Network
}
// NewProtocolTester constructs a new ProtocolTester
// it takes as argument the pivot node id, the number of dummy peers and the
// protocol run function called on a peer connection by the p2p server
func NewProtocolTester(id enode.ID, nodeCount int, run func(*p2p.Peer, p2p.MsgReadWriter) error) *ProtocolTester {
services := adapters.Services{
"test": func(ctx *adapters.ServiceContext) (node.Service, error) {
return &testNode{run}, nil
},
"mock": func(ctx *adapters.ServiceContext) (node.Service, error) {
return newMockNode(), nil
},
}
adapter := adapters.NewSimAdapter(services)
net := simulations.NewNetwork(adapter, &simulations.NetworkConfig{})
if _, err := net.NewNodeWithConfig(&adapters.NodeConfig{
ID: id,
EnableMsgEvents: true,
Services: []string{"test"},
}); err != nil {
panic(err.Error())
}
if err := net.Start(id); err != nil {
panic(err.Error())
}
node := net.GetNode(id).Node.(*adapters.SimNode)
peers := make([]*adapters.NodeConfig, nodeCount)
nodes := make([]*enode.Node, nodeCount)
for i := 0; i < nodeCount; i++ {
peers[i] = adapters.RandomNodeConfig()
peers[i].Services = []string{"mock"}
nodes[i] = peers[i].Node()
}
events := make(chan *p2p.PeerEvent, 1000)
node.SubscribeEvents(events)
ps := &ProtocolSession{
Server: node.Server(),
Nodes: nodes,
adapter: adapter,
events: events,
}
self := &ProtocolTester{
ProtocolSession: ps,
network: net,
}
self.Connect(id, peers...)
return self
}
// Stop stops the p2p server
func (t *ProtocolTester) Stop() {
t.Server.Stop()
}
// Connect brings up the remote peer node and connects it using the
// p2p/simulations network connection with the in memory network adapter
func (t *ProtocolTester) Connect(selfID enode.ID, peers ...*adapters.NodeConfig) {
for _, peer := range peers {
log.Trace(fmt.Sprintf("start node %v", peer.ID))
if _, err := t.network.NewNodeWithConfig(peer); err != nil {
panic(fmt.Sprintf("error starting peer %v: %v", peer.ID, err))
}
if err := t.network.Start(peer.ID); err != nil {
panic(fmt.Sprintf("error starting peer %v: %v", peer.ID, err))
}
log.Trace(fmt.Sprintf("connect to %v", peer.ID))
if err := t.network.Connect(selfID, peer.ID); err != nil {
panic(fmt.Sprintf("error connecting to peer %v: %v", peer.ID, err))
}
}
}
// testNode wraps a protocol run function and implements the node.Service
// interface
type testNode struct {
run func(*p2p.Peer, p2p.MsgReadWriter) error
}
func (t *testNode) Protocols() []p2p.Protocol {
return []p2p.Protocol{{
Length: 100,
Run: t.run,
}}
}
func (t *testNode) APIs() []rpc.API {
return nil
}
func (t *testNode) Start(server *p2p.Server) error {
return nil
}
func (t *testNode) Stop() error {
return nil
}
// mockNode is a testNode which doesn't actually run a protocol, instead
// exposing channels so that tests can manually trigger and expect certain
// messages
type mockNode struct {
testNode
trigger chan *Trigger
expect chan []Expect
err chan error
stop chan struct{}
stopOnce sync.Once
}
func newMockNode() *mockNode {
mock := &mockNode{
trigger: make(chan *Trigger),
expect: make(chan []Expect),
err: make(chan error),
stop: make(chan struct{}),
}
mock.testNode.run = mock.Run
return mock
}
// Run is a protocol run function which just loops waiting for tests to
// instruct it to either trigger or expect a message from the peer
func (m *mockNode) Run(peer *p2p.Peer, rw p2p.MsgReadWriter) error {
for {
select {
case trig := <-m.trigger:
wmsg := Wrap(trig.Msg)
m.err <- p2p.Send(rw, trig.Code, wmsg)
case exps := <-m.expect:
m.err <- expectMsgs(rw, exps)
case <-m.stop:
return nil
}
}
}
func (m *mockNode) Trigger(trig *Trigger) error {
m.trigger <- trig
return <-m.err
}
func (m *mockNode) Expect(exp ...Expect) error {
m.expect <- exp
return <-m.err
}
func (m *mockNode) Stop() error {
m.stopOnce.Do(func() { close(m.stop) })
return nil
}
func expectMsgs(rw p2p.MsgReadWriter, exps []Expect) error {
matched := make([]bool, len(exps))
for {
msg, err := rw.ReadMsg()
if err != nil {
if err == io.EOF {
break
}
return err
}
actualContent, err := ioutil.ReadAll(msg.Payload)
if err != nil {
return err
}
var found bool
for i, exp := range exps {
if exp.Code == msg.Code && bytes.Equal(actualContent, mustEncodeMsg(Wrap(exp.Msg))) {
if matched[i] {
return fmt.Errorf("message #%d received two times", i)
}
matched[i] = true
found = true
break
}
}
if !found {
expected := make([]string, 0)
for i, exp := range exps {
if matched[i] {
continue
}
expected = append(expected, fmt.Sprintf("code %d payload %x", exp.Code, mustEncodeMsg(Wrap(exp.Msg))))
}
return fmt.Errorf("unexpected message code %d payload %x, expected %s", msg.Code, actualContent, strings.Join(expected, " or "))
}
done := true
for _, m := range matched {
if !m {
done = false
break
}
}
if done {
return nil
}
}
for i, m := range matched {
if !m {
return fmt.Errorf("expected message #%d not received", i)
}
}
return nil
}
// mustEncodeMsg uses rlp to encode a message.
// In case of error it panics.
func mustEncodeMsg(msg interface{}) []byte {
contentEnc, err := rlp.EncodeToBytes(msg)
if err != nil {
panic("content encode error: " + err.Error())
}
return contentEnc
}
type WrappedMsg struct {
Context []byte
Size uint32
Payload []byte
}
func Wrap(msg interface{}) interface{} {
data, _ := rlp.EncodeToBytes(msg)
return &WrappedMsg{
Size: uint32(len(data)),
Payload: data,
}
}