cmd/devp2p/internal/ethtest: improve eth test suite (#21615)

This fixes issues with the protocol handshake and status exchange
and adds support for responding to GetBlockHeaders requests.

cmd/devp2p/internal/ethtest/chain.go

@@ -1,3 +1,19 @@
+// Copyright 2020 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/>.
+
 package ethtest
 
 import (
@@ -68,6 +84,43 @@ func (c *Chain) Head() *types.Block {
 	return c.blocks[c.Len()-1]
 }
 
+func (c *Chain) GetHeaders(req GetBlockHeaders) (BlockHeaders, error) {
+	if req.Amount < 1 {
+		return nil, fmt.Errorf("no block headers requested")
+	}
+
+	headers := make(BlockHeaders, req.Amount)
+	var blockNumber uint64
+
+	// range over blocks to check if our chain has the requested header
+	for _, block := range c.blocks {
+		if block.Hash() == req.Origin.Hash || block.Number().Uint64() == req.Origin.Number {
+			headers[0] = block.Header()
+			blockNumber = block.Number().Uint64()
+		}
+	}
+	if headers[0] == nil {
+		return nil, fmt.Errorf("no headers found for given origin number %v, hash %v", req.Origin.Number, req.Origin.Hash)
+	}
+
+	if req.Reverse {
+		for i := 1; i < int(req.Amount); i++ {
+			blockNumber -= (1 - req.Skip)
+			headers[i] = c.blocks[blockNumber].Header()
+
+		}
+
+		return headers, nil
+	}
+
+	for i := 1; i < int(req.Amount); i++ {
+		blockNumber += (1 + req.Skip)
+		headers[i] = c.blocks[blockNumber].Header()
+	}
+
+	return headers, nil
+}
+
 // loadChain takes the given chain.rlp file, and decodes and returns
 // the blocks from the file.
 func loadChain(chainfile string, genesis string) (*Chain, error) {

cmd/devp2p/internal/ethtest/chain_test.go

@@ -0,0 +1,150 @@
+// Copyright 2020 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/>.
+
+package ethtest
+
+import (
+	"path/filepath"
+	"strconv"
+	"testing"
+
+	"github.com/ethereum/go-ethereum/p2p"
+	"github.com/stretchr/testify/assert"
+)
+
+// TestEthProtocolNegotiation tests whether the test suite
+// can negotiate the highest eth protocol in a status message exchange
+func TestEthProtocolNegotiation(t *testing.T) {
+	var tests = []struct {
+		conn     *Conn
+		caps     []p2p.Cap
+		expected uint32
+	}{
+		{
+			conn: &Conn{},
+			caps: []p2p.Cap{
+				{Name: "eth", Version: 63},
+				{Name: "eth", Version: 64},
+				{Name: "eth", Version: 65},
+			},
+			expected: uint32(65),
+		},
+		{
+			conn: &Conn{},
+			caps: []p2p.Cap{
+				{Name: "eth", Version: 0},
+				{Name: "eth", Version: 89},
+				{Name: "eth", Version: 65},
+			},
+			expected: uint32(65),
+		},
+		{
+			conn: &Conn{},
+			caps: []p2p.Cap{
+				{Name: "eth", Version: 63},
+				{Name: "eth", Version: 64},
+				{Name: "wrongProto", Version: 65},
+			},
+			expected: uint32(64),
+		},
+	}
+
+	for i, tt := range tests {
+		t.Run(strconv.Itoa(i), func(t *testing.T) {
+			tt.conn.negotiateEthProtocol(tt.caps)
+			assert.Equal(t, tt.expected, uint32(tt.conn.ethProtocolVersion))
+		})
+	}
+}
+
+// TestChain_GetHeaders tests whether the test suite can correctly
+// respond to a GetBlockHeaders request from a node.
+func TestChain_GetHeaders(t *testing.T) {
+	chainFile, err := filepath.Abs("./testdata/chain.rlp.gz")
+	if err != nil {
+		t.Fatal(err)
+	}
+	genesisFile, err := filepath.Abs("./testdata/genesis.json")
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	chain, err := loadChain(chainFile, genesisFile)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	var tests = []struct {
+		req      GetBlockHeaders
+		expected BlockHeaders
+	}{
+		{
+			req: GetBlockHeaders{
+				Origin: hashOrNumber{
+					Number: uint64(2),
+				},
+				Amount:  uint64(5),
+				Skip:    1,
+				Reverse: false,
+			},
+			expected: BlockHeaders{
+				chain.blocks[2].Header(),
+				chain.blocks[4].Header(),
+				chain.blocks[6].Header(),
+				chain.blocks[8].Header(),
+				chain.blocks[10].Header(),
+			},
+		},
+		{
+			req: GetBlockHeaders{
+				Origin: hashOrNumber{
+					Number: uint64(chain.Len() - 1),
+				},
+				Amount:  uint64(3),
+				Skip:    0,
+				Reverse: true,
+			},
+			expected: BlockHeaders{
+				chain.blocks[chain.Len()-1].Header(),
+				chain.blocks[chain.Len()-2].Header(),
+				chain.blocks[chain.Len()-3].Header(),
+			},
+		},
+		{
+			req: GetBlockHeaders{
+				Origin: hashOrNumber{
+					Hash: chain.Head().Hash(),
+				},
+				Amount:  uint64(1),
+				Skip:    0,
+				Reverse: false,
+			},
+			expected: BlockHeaders{
+				chain.Head().Header(),
+			},
+		},
+	}
+
+	for i, tt := range tests {
+		t.Run(strconv.Itoa(i), func(t *testing.T) {
+			headers, err := chain.GetHeaders(tt.req)
+			if err != nil {
+				t.Fatal(err)
+			}
+			assert.Equal(t, headers, tt.expected)
+		})
+	}
+}

cmd/devp2p/internal/ethtest/suite.go

@@ -1,19 +1,29 @@
+// Copyright 2020 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/>.
+
 package ethtest
 
 import (
-	"crypto/ecdsa"
 	"fmt"
 	"net"
-	"reflect"
-	"time"
 
-	"github.com/ethereum/go-ethereum/core/types"
 	"github.com/ethereum/go-ethereum/crypto"
 	"github.com/ethereum/go-ethereum/internal/utesting"
-	"github.com/ethereum/go-ethereum/p2p"
 	"github.com/ethereum/go-ethereum/p2p/enode"
 	"github.com/ethereum/go-ethereum/p2p/rlpx"
-	"github.com/ethereum/go-ethereum/rlp"
 	"github.com/stretchr/testify/assert"
 )
 
@@ -26,137 +36,6 @@ type Suite struct {
 	fullChain *Chain
 }
 
-type Conn struct {
-	*rlpx.Conn
-	ourKey *ecdsa.PrivateKey
-}
-
-func (c *Conn) Read() Message {
-	code, rawData, _, err := c.Conn.Read()
-	if err != nil {
-		return &Error{fmt.Errorf("could not read from connection: %v", err)}
-	}
-
-	var msg Message
-	switch int(code) {
-	case (Hello{}).Code():
-		msg = new(Hello)
-	case (Disconnect{}).Code():
-		msg = new(Disconnect)
-	case (Status{}).Code():
-		msg = new(Status)
-	case (GetBlockHeaders{}).Code():
-		msg = new(GetBlockHeaders)
-	case (BlockHeaders{}).Code():
-		msg = new(BlockHeaders)
-	case (GetBlockBodies{}).Code():
-		msg = new(GetBlockBodies)
-	case (BlockBodies{}).Code():
-		msg = new(BlockBodies)
-	case (NewBlock{}).Code():
-		msg = new(NewBlock)
-	case (NewBlockHashes{}).Code():
-		msg = new(NewBlockHashes)
-	default:
-		return &Error{fmt.Errorf("invalid message code: %d", code)}
-	}
-
-	if err := rlp.DecodeBytes(rawData, msg); err != nil {
-		return &Error{fmt.Errorf("could not rlp decode message: %v", err)}
-	}
-
-	return msg
-}
-
-func (c *Conn) Write(msg Message) error {
-	payload, err := rlp.EncodeToBytes(msg)
-	if err != nil {
-		return err
-	}
-	_, err = c.Conn.Write(uint64(msg.Code()), payload)
-	return err
-
-}
-
-// handshake checks to make sure a `HELLO` is received.
-func (c *Conn) handshake(t *utesting.T) Message {
-	// write protoHandshake to client
-	pub0 := crypto.FromECDSAPub(&c.ourKey.PublicKey)[1:]
-	ourHandshake := &Hello{
-		Version: 5,
-		Caps:    []p2p.Cap{{Name: "eth", Version: 64}, {Name: "eth", Version: 65}},
-		ID:      pub0,
-	}
-	if err := c.Write(ourHandshake); err != nil {
-		t.Fatalf("could not write to connection: %v", err)
-	}
-	// read protoHandshake from client
-	switch msg := c.Read().(type) {
-	case *Hello:
-		return msg
-	default:
-		t.Fatalf("bad handshake: %v", msg)
-		return nil
-	}
-}
-
-// statusExchange performs a `Status` message exchange with the given
-// node.
-func (c *Conn) statusExchange(t *utesting.T, chain *Chain) Message {
-	// read status message from client
-	var message Message
-	switch msg := c.Read().(type) {
-	case *Status:
-		if msg.Head != chain.blocks[chain.Len()-1].Hash() {
-			t.Fatalf("wrong head in status: %v", msg.Head)
-		}
-		if msg.TD.Cmp(chain.TD(chain.Len())) != 0 {
-			t.Fatalf("wrong TD in status: %v", msg.TD)
-		}
-		if !reflect.DeepEqual(msg.ForkID, chain.ForkID()) {
-			t.Fatalf("wrong fork ID in status: %v", msg.ForkID)
-		}
-		message = msg
-	default:
-		t.Fatalf("bad status message: %v", msg)
-	}
-	// write status message to client
-	status := Status{
-		ProtocolVersion: 64,
-		NetworkID:       1,
-		TD:              chain.TD(chain.Len()),
-		Head:            chain.blocks[chain.Len()-1].Hash(),
-		Genesis:         chain.blocks[0].Hash(),
-		ForkID:          chain.ForkID(),
-	}
-	if err := c.Write(status); err != nil {
-		t.Fatalf("could not write to connection: %v", err)
-	}
-
-	return message
-}
-
-// waitForBlock waits for confirmation from the client that it has
-// imported the given block.
-func (c *Conn) waitForBlock(block *types.Block) error {
-	for {
-		req := &GetBlockHeaders{Origin: hashOrNumber{Hash: block.Hash()}, Amount: 1}
-		if err := c.Write(req); err != nil {
-			return err
-		}
-
-		switch msg := c.Read().(type) {
-		case *BlockHeaders:
-			if len(*msg) > 0 {
-				return nil
-			}
-			time.Sleep(100 * time.Millisecond)
-		default:
-			return fmt.Errorf("invalid message: %v", msg)
-		}
-	}
-}
-
 // NewSuite creates and returns a new eth-test suite that can
 // be used to test the given node against the given blockchain
 // data.
@@ -196,7 +75,7 @@ func (s *Suite) TestStatus(t *utesting.T) {
 	case *Status:
 		t.Logf("%+v\n", msg)
 	default:
-		t.Fatalf("error: %v", msg)
+		t.Fatalf("unexpected: %#v", msg)
 	}
 }
 
@@ -225,7 +104,7 @@ func (s *Suite) TestGetBlockHeaders(t *utesting.T) {
 		t.Fatalf("could not write to connection: %v", err)
 	}
 
-	switch msg := conn.Read().(type) {
+	switch msg := conn.ReadAndServe(s.chain).(type) {
 	case *BlockHeaders:
 		headers := msg
 		for _, header := range *headers {
@@ -234,7 +113,7 @@ func (s *Suite) TestGetBlockHeaders(t *utesting.T) {
 			t.Logf("\nHEADER FOR BLOCK NUMBER %d: %+v\n", header.Number, header)
 		}
 	default:
-		t.Fatalf("error: %v", msg)
+		t.Fatalf("unexpected: %#v", msg)
 	}
 }
 
@@ -254,14 +133,14 @@ func (s *Suite) TestGetBlockBodies(t *utesting.T) {
 		t.Fatalf("could not write to connection: %v", err)
 	}
 
-	switch msg := conn.Read().(type) {
+	switch msg := conn.ReadAndServe(s.chain).(type) {
 	case *BlockBodies:
 		bodies := msg
 		for _, body := range *bodies {
 			t.Logf("\nBODY: %+v\n", body)
 		}
 	default:
-		t.Fatalf("error: %v", msg)
+		t.Fatalf("unexpected: %#v", msg)
 	}
 }
 
@@ -294,7 +173,7 @@ func (s *Suite) TestBroadcast(t *utesting.T) {
 		t.Fatalf("could not write to connection: %v", err)
 	}
 
-	switch msg := receiveConn.Read().(type) {
+	switch msg := receiveConn.ReadAndServe(s.chain).(type) {
 	case *NewBlock:
 		assert.Equal(t, blockAnnouncement.Block.Header(), msg.Block.Header(),
 			"wrong block header in announcement")
@@ -305,7 +184,7 @@ func (s *Suite) TestBroadcast(t *utesting.T) {
 		assert.Equal(t, blockAnnouncement.Block.Hash(), hashes[0].Hash,
 			"wrong block hash in announcement")
 	default:
-		t.Fatal(msg)
+		t.Fatalf("unexpected: %#v", msg)
 	}
 	// update test suite chain
 	s.chain.blocks = append(s.chain.blocks, s.fullChain.blocks[1000])

cmd/devp2p/internal/ethtest/testdata/genesis.json

@@ -0,0 +1,26 @@
+{
+    "config": {
+        "chainId": 1,
+        "homesteadBlock": 0,
+        "eip150Block": 0,
+        "eip155Block": 0,
+        "eip158Block": 0,
+        "byzantiumBlock": 0,
+        "ethash": {}
+    },
+    "nonce": "0xdeadbeefdeadbeef",
+    "timestamp": "0x0",
+    "extraData": "0x0000000000000000000000000000000000000000000000000000000000000000",
+    "gasLimit": "0x8000000",
+    "difficulty": "0x10",
+    "mixHash": "0x0000000000000000000000000000000000000000000000000000000000000000",
+    "coinbase": "0x0000000000000000000000000000000000000000",
+    "alloc": {
+        "71562b71999873db5b286df957af199ec94617f7": {
+            "balance": "0xf4240"
+        }
+    },
+    "number": "0x0",
+    "gasUsed": "0x0",
+    "parentHash": "0x0000000000000000000000000000000000000000000000000000000000000000"
+}

cmd/devp2p/internal/ethtest/types.go

@@ -1,14 +1,36 @@
+// Copyright 2020 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/>.
+
 package ethtest
 
 import (
+	"crypto/ecdsa"
 	"fmt"
 	"io"
 	"math/big"
+	"reflect"
+	"time"
 
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/core/forkid"
 	"github.com/ethereum/go-ethereum/core/types"
+	"github.com/ethereum/go-ethereum/crypto"
+	"github.com/ethereum/go-ethereum/internal/utesting"
 	"github.com/ethereum/go-ethereum/p2p"
+	"github.com/ethereum/go-ethereum/p2p/rlpx"
 	"github.com/ethereum/go-ethereum/rlp"
 )
 
@@ -23,6 +45,7 @@ type Error struct {
 func (e *Error) Unwrap() error    { return e.err }
 func (e *Error) Error() string    { return e.err.Error() }
 func (e *Error) Code() int        { return -1 }
+func (e *Error) GoString() string { return e.Error() }
 
 // Hello is the RLP structure of the protocol handshake.
 type Hello struct {
@@ -45,6 +68,14 @@ type Disconnect struct {
 
 func (d Disconnect) Code() int { return 0x01 }
 
+type Ping struct{}
+
+func (p Ping) Code() int { return 0x02 }
+
+type Pong struct{}
+
+func (p Pong) Code() int { return 0x03 }
+
 // Status is the network packet for the status message for eth/64 and later.
 type Status struct {
 	ProtocolVersion uint32
@@ -132,3 +163,204 @@ func (gbb GetBlockBodies) Code() int { return 21 }
 type BlockBodies []*types.Body
 
 func (bb BlockBodies) Code() int { return 22 }
+
+// Conn represents an individual connection with a peer
+type Conn struct {
+	*rlpx.Conn
+	ourKey             *ecdsa.PrivateKey
+	ethProtocolVersion uint
+}
+
+func (c *Conn) Read() Message {
+	code, rawData, _, err := c.Conn.Read()
+	if err != nil {
+		return &Error{fmt.Errorf("could not read from connection: %v", err)}
+	}
+
+	var msg Message
+	switch int(code) {
+	case (Hello{}).Code():
+		msg = new(Hello)
+	case (Ping{}).Code():
+		msg = new(Ping)
+	case (Pong{}).Code():
+		msg = new(Pong)
+	case (Disconnect{}).Code():
+		msg = new(Disconnect)
+	case (Status{}).Code():
+		msg = new(Status)
+	case (GetBlockHeaders{}).Code():
+		msg = new(GetBlockHeaders)
+	case (BlockHeaders{}).Code():
+		msg = new(BlockHeaders)
+	case (GetBlockBodies{}).Code():
+		msg = new(GetBlockBodies)
+	case (BlockBodies{}).Code():
+		msg = new(BlockBodies)
+	case (NewBlock{}).Code():
+		msg = new(NewBlock)
+	case (NewBlockHashes{}).Code():
+		msg = new(NewBlockHashes)
+	default:
+		return &Error{fmt.Errorf("invalid message code: %d", code)}
+	}
+
+	if err := rlp.DecodeBytes(rawData, msg); err != nil {
+		return &Error{fmt.Errorf("could not rlp decode message: %v", err)}
+	}
+
+	return msg
+}
+
+// ReadAndServe serves GetBlockHeaders requests while waiting
+// on another message from the node.
+func (c *Conn) ReadAndServe(chain *Chain) Message {
+	for {
+		switch msg := c.Read().(type) {
+		case *Ping:
+			c.Write(&Pong{})
+		case *GetBlockHeaders:
+			req := *msg
+			headers, err := chain.GetHeaders(req)
+			if err != nil {
+				return &Error{fmt.Errorf("could not get headers for inbound header request: %v", err)}
+			}
+
+			if err := c.Write(headers); err != nil {
+				return &Error{fmt.Errorf("could not write to connection: %v", err)}
+			}
+		default:
+			return msg
+		}
+	}
+}
+
+func (c *Conn) Write(msg Message) error {
+	payload, err := rlp.EncodeToBytes(msg)
+	if err != nil {
+		return err
+	}
+	_, err = c.Conn.Write(uint64(msg.Code()), payload)
+	return err
+
+}
+
+// handshake checks to make sure a `HELLO` is received.
+func (c *Conn) handshake(t *utesting.T) Message {
+	// write protoHandshake to client
+	pub0 := crypto.FromECDSAPub(&c.ourKey.PublicKey)[1:]
+	ourHandshake := &Hello{
+		Version: 5,
+		Caps: []p2p.Cap{
+			{Name: "eth", Version: 64},
+			{Name: "eth", Version: 65},
+		},
+		ID: pub0,
+	}
+	if err := c.Write(ourHandshake); err != nil {
+		t.Fatalf("could not write to connection: %v", err)
+	}
+	// read protoHandshake from client
+	switch msg := c.Read().(type) {
+	case *Hello:
+		// set snappy if version is at least 5
+		if msg.Version >= 5 {
+			c.SetSnappy(true)
+		}
+
+		c.negotiateEthProtocol(msg.Caps)
+		if c.ethProtocolVersion == 0 {
+			t.Fatalf("unexpected eth protocol version")
+		}
+		return msg
+	default:
+		t.Fatalf("bad handshake: %#v", msg)
+		return nil
+	}
+}
+
+// negotiateEthProtocol sets the Conn's eth protocol version
+// to highest advertised capability from peer
+func (c *Conn) negotiateEthProtocol(caps []p2p.Cap) {
+	var highestEthVersion uint
+	for _, capability := range caps {
+		if capability.Name != "eth" {
+			continue
+		}
+		if capability.Version > highestEthVersion && capability.Version <= 65 {
+			highestEthVersion = capability.Version
+		}
+	}
+	c.ethProtocolVersion = highestEthVersion
+}
+
+// statusExchange performs a `Status` message exchange with the given
+// node.
+func (c *Conn) statusExchange(t *utesting.T, chain *Chain) Message {
+	// read status message from client
+	var message Message
+
+loop:
+	for {
+		switch msg := c.Read().(type) {
+		case *Status:
+			if msg.Head != chain.blocks[chain.Len()-1].Hash() {
+				t.Fatalf("wrong head in status: %v", msg.Head)
+			}
+			if msg.TD.Cmp(chain.TD(chain.Len())) != 0 {
+				t.Fatalf("wrong TD in status: %v", msg.TD)
+			}
+			if !reflect.DeepEqual(msg.ForkID, chain.ForkID()) {
+				t.Fatalf("wrong fork ID in status: %v", msg.ForkID)
+			}
+			message = msg
+			break loop
+		case *Disconnect:
+			t.Fatalf("disconnect received: %v", msg.Reason)
+		case *Ping:
+			c.Write(&Pong{}) // TODO (renaynay): in the future, this should be an error
+			// (PINGs should not be a response upon fresh connection)
+		default:
+			t.Fatalf("bad status message: %#v", msg)
+		}
+	}
+	// make sure eth protocol version is set for negotiation
+	if c.ethProtocolVersion == 0 {
+		t.Fatalf("eth protocol version must be set in Conn")
+	}
+	// write status message to client
+	status := Status{
+		ProtocolVersion: uint32(c.ethProtocolVersion),
+		NetworkID:       1,
+		TD:              chain.TD(chain.Len()),
+		Head:            chain.blocks[chain.Len()-1].Hash(),
+		Genesis:         chain.blocks[0].Hash(),
+		ForkID:          chain.ForkID(),
+	}
+	if err := c.Write(status); err != nil {
+		t.Fatalf("could not write to connection: %v", err)
+	}
+
+	return message
+}
+
+// waitForBlock waits for confirmation from the client that it has
+// imported the given block.
+func (c *Conn) waitForBlock(block *types.Block) error {
+	for {
+		req := &GetBlockHeaders{Origin: hashOrNumber{Hash: block.Hash()}, Amount: 1}
+		if err := c.Write(req); err != nil {
+			return err
+		}
+
+		switch msg := c.Read().(type) {
+		case *BlockHeaders:
+			if len(*msg) > 0 {
+				return nil
+			}
+			time.Sleep(100 * time.Millisecond)
+		default:
+			return fmt.Errorf("invalid message: %v", msg)
+		}
+	}
+}