1 files changed, 218 insertions, 0 deletions
diff --git a/eth/protocol_test.go b/eth/protocol_test.go
new file mode 100644
index 000000000..8ca6d1be6
--- /dev/null
+++ b/eth/protocol_test.go
@@ -0,0 +1,218 @@
+package eth
+
+import (
+	"log"
+	"math/big"
+	"os"
+	"testing"
+	"time"
+
+	"github.com/ethereum/go-ethereum/common"
+	"github.com/ethereum/go-ethereum/core/types"
+	"github.com/ethereum/go-ethereum/crypto"
+	"github.com/ethereum/go-ethereum/errs"
+	ethlogger "github.com/ethereum/go-ethereum/logger"
+	"github.com/ethereum/go-ethereum/p2p"
+	"github.com/ethereum/go-ethereum/p2p/discover"
+)
+
+var logsys = ethlogger.NewStdLogSystem(os.Stdout, log.LstdFlags, ethlogger.LogLevel(ethlogger.DebugDetailLevel))
+
+var ini = false
+
+func logInit() {
+	if !ini {
+		ethlogger.AddLogSystem(logsys)
+		ini = true
+	}
+}
+
+type testTxPool struct {
+	getTransactions func() []*types.Transaction
+	addTransactions func(txs []*types.Transaction)
+}
+
+type testChainManager struct {
+	getBlockHashes func(hash common.Hash, amount uint64) (hashes []common.Hash)
+	getBlock       func(hash common.Hash) *types.Block
+	status         func() (td *big.Int, currentBlock common.Hash, genesisBlock common.Hash)
+}
+
+type testBlockPool struct {
+	addBlockHashes func(next func() (common.Hash, bool), peerId string)
+	addBlock       func(block *types.Block, peerId string) (err error)
+	addPeer        func(td *big.Int, currentBlock common.Hash, peerId string, requestHashes func(common.Hash) error, requestBlocks func([]common.Hash) error, peerError func(*errs.Error)) (best bool, suspended bool)
+	removePeer     func(peerId string)
+}
+
+func (self *testTxPool) AddTransactions(txs []*types.Transaction) {
+	if self.addTransactions != nil {
+		self.addTransactions(txs)
+	}
+}
+
+func (self *testTxPool) GetTransactions() types.Transactions { return nil }
+
+func (self *testChainManager) GetBlockHashesFromHash(hash common.Hash, amount uint64) (hashes []common.Hash) {
+	if self.getBlockHashes != nil {
+		hashes = self.getBlockHashes(hash, amount)
+	}
+	return
+}
+
+func (self *testChainManager) Status() (td *big.Int, currentBlock common.Hash, genesisBlock common.Hash) {
+	if self.status != nil {
+		td, currentBlock, genesisBlock = self.status()
+	}
+	return
+}
+
+func (self *testChainManager) GetBlock(hash common.Hash) (block *types.Block) {
+	if self.getBlock != nil {
+		block = self.getBlock(hash)
+	}
+	return
+}
+
+func (self *testBlockPool) AddBlockHashes(next func() (common.Hash, bool), peerId string) {
+	if self.addBlockHashes != nil {
+		self.addBlockHashes(next, peerId)
+	}
+}
+
+func (self *testBlockPool) AddBlock(block *types.Block, peerId string) {
+	if self.addBlock != nil {
+		self.addBlock(block, peerId)
+	}
+}
+
+func (self *testBlockPool) AddPeer(td *big.Int, currentBlock common.Hash, peerId string, requestBlockHashes func(common.Hash) error, requestBlocks func([]common.Hash) error, peerError func(*errs.Error)) (best bool, suspended bool) {
+	if self.addPeer != nil {
+		best, suspended = self.addPeer(td, currentBlock, peerId, requestBlockHashes, requestBlocks, peerError)
+	}
+	return
+}
+
+func (self *testBlockPool) RemovePeer(peerId string) {
+	if self.removePeer != nil {
+		self.removePeer(peerId)
+	}
+}
+
+func testPeer() *p2p.Peer {
+	var id discover.NodeID
+	pk := crypto.GenerateNewKeyPair().PublicKey
+	copy(id[:], pk)
+	return p2p.NewPeer(id, "test peer", []p2p.Cap{})
+}
+
+type ethProtocolTester struct {
+	p2p.MsgReadWriter // writing to the tester feeds the protocol
+
+	quit         chan error
+	pipe         *p2p.MsgPipeRW    // the protocol read/writes on this end
+	txPool       *testTxPool       // txPool
+	chainManager *testChainManager // chainManager
+	blockPool    *testBlockPool    // blockPool
+	t            *testing.T
+}
+
+func newEth(t *testing.T) *ethProtocolTester {
+	p1, p2 := p2p.MsgPipe()
+	return &ethProtocolTester{
+		MsgReadWriter: p1,
+		quit:          make(chan error, 1),
+		pipe:          p2,
+		txPool:        &testTxPool{},
+		chainManager:  &testChainManager{},
+		blockPool:     &testBlockPool{},
+		t:             t,
+	}
+}
+
+func (self *ethProtocolTester) reset() {
+	self.pipe.Close()
+
+	p1, p2 := p2p.MsgPipe()
+	self.MsgReadWriter = p1
+	self.pipe = p2
+	self.quit = make(chan error, 1)
+}
+
+func (self *ethProtocolTester) checkError(expCode int, delay time.Duration) (err error) {
+	var timer = time.After(delay)
+	select {
+	case err = <-self.quit:
+	case <-timer:
+		self.t.Errorf("no error after %v, expected %v", delay, expCode)
+		return
+	}
+	perr, ok := err.(*errs.Error)
+	if ok && perr != nil {
+		if code := perr.Code; code != expCode {
+			self.t.Errorf("expected protocol error (code %v), got %v (%v)", expCode, code, err)
+		}
+	} else {
+		self.t.Errorf("expected protocol error (code %v), got %v", expCode, err)
+	}
+	return
+}
+
+func (self *ethProtocolTester) run() {
+	err := runEthProtocol(ProtocolVersion, NetworkId, self.txPool, self.chainManager, self.blockPool, testPeer(), self.pipe)
+	self.quit <- err
+}
+
+func TestStatusMsgErrors(t *testing.T) {
+	logInit()
+	eth := newEth(t)
+	td := common.Big1
+	currentBlock := common.Hash{1}
+	genesis := common.Hash{2}
+	eth.chainManager.status = func() (*big.Int, common.Hash, common.Hash) { return td, currentBlock, genesis }
+	go eth.run()
+
+	tests := []struct {
+		code          uint64
+		data          interface{}
+		wantErrorCode int
+	}{
+		{
+			code: TxMsg, data: []interface{}{},
+			wantErrorCode: ErrNoStatusMsg,
+		},
+		{
+			code: StatusMsg, data: statusMsgData{10, NetworkId, td, currentBlock, genesis},
+			wantErrorCode: ErrProtocolVersionMismatch,
+		},
+		{
+			code: StatusMsg, data: statusMsgData{ProtocolVersion, 999, td, currentBlock, genesis},
+			wantErrorCode: ErrNetworkIdMismatch,
+		},
+		{
+			code: StatusMsg, data: statusMsgData{ProtocolVersion, NetworkId, td, currentBlock, common.Hash{3}},
+			wantErrorCode: ErrGenesisBlockMismatch,
+		},
+	}
+	for _, test := range tests {
+		// first outgoing msg should be StatusMsg.
+		err := p2p.ExpectMsg(eth, StatusMsg, &statusMsgData{
+			ProtocolVersion: ProtocolVersion,
+			NetworkId:       NetworkId,
+			TD:              td,
+			CurrentBlock:    currentBlock,
+			GenesisBlock:    genesis,
+		})
+		if err != nil {
+			t.Fatalf("incorrect outgoing status: %v", err)
+		}
+
+		// the send call might hang until reset because
+		// the protocol might not read the payload.
+		go p2p.Send(eth, test.code, test.data)
+		eth.checkError(test.wantErrorCode, 1*time.Second)
+
+		eth.reset()
+		go eth.run()
+	}
+}