github.com/anthdm/go-ethereum@v1.8.4-0.20180412101906-60516c83b011/eth/handler_test.go (about)

     1  // Copyright 2015 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 eth
    18  
    19  import (
    20  	"math"
    21  	"math/big"
    22  	"math/rand"
    23  	"testing"
    24  	"time"
    25  
    26  	"github.com/ethereum/go-ethereum/common"
    27  	"github.com/ethereum/go-ethereum/consensus/ethash"
    28  	"github.com/ethereum/go-ethereum/core"
    29  	"github.com/ethereum/go-ethereum/core/state"
    30  	"github.com/ethereum/go-ethereum/core/types"
    31  	"github.com/ethereum/go-ethereum/core/vm"
    32  	"github.com/ethereum/go-ethereum/crypto"
    33  	"github.com/ethereum/go-ethereum/eth/downloader"
    34  	"github.com/ethereum/go-ethereum/ethdb"
    35  	"github.com/ethereum/go-ethereum/event"
    36  	"github.com/ethereum/go-ethereum/p2p"
    37  	"github.com/ethereum/go-ethereum/params"
    38  )
    39  
    40  // Tests that protocol versions and modes of operations are matched up properly.
    41  func TestProtocolCompatibility(t *testing.T) {
    42  	// Define the compatibility chart
    43  	tests := []struct {
    44  		version    uint
    45  		mode       downloader.SyncMode
    46  		compatible bool
    47  	}{
    48  		{61, downloader.FullSync, true}, {62, downloader.FullSync, true}, {63, downloader.FullSync, true},
    49  		{61, downloader.FastSync, false}, {62, downloader.FastSync, false}, {63, downloader.FastSync, true},
    50  	}
    51  	// Make sure anything we screw up is restored
    52  	backup := ProtocolVersions
    53  	defer func() { ProtocolVersions = backup }()
    54  
    55  	// Try all available compatibility configs and check for errors
    56  	for i, tt := range tests {
    57  		ProtocolVersions = []uint{tt.version}
    58  
    59  		pm, _, err := newTestProtocolManager(tt.mode, 0, nil, nil)
    60  		if pm != nil {
    61  			defer pm.Stop()
    62  		}
    63  		if (err == nil && !tt.compatible) || (err != nil && tt.compatible) {
    64  			t.Errorf("test %d: compatibility mismatch: have error %v, want compatibility %v", i, err, tt.compatible)
    65  		}
    66  	}
    67  }
    68  
    69  // Tests that block headers can be retrieved from a remote chain based on user queries.
    70  func TestGetBlockHeaders62(t *testing.T) { testGetBlockHeaders(t, 62) }
    71  func TestGetBlockHeaders63(t *testing.T) { testGetBlockHeaders(t, 63) }
    72  
    73  func testGetBlockHeaders(t *testing.T, protocol int) {
    74  	pm, _ := newTestProtocolManagerMust(t, downloader.FullSync, downloader.MaxHashFetch+15, nil, nil)
    75  	peer, _ := newTestPeer("peer", protocol, pm, true)
    76  	defer peer.close()
    77  
    78  	// Create a "random" unknown hash for testing
    79  	var unknown common.Hash
    80  	for i := range unknown {
    81  		unknown[i] = byte(i)
    82  	}
    83  	// Create a batch of tests for various scenarios
    84  	limit := uint64(downloader.MaxHeaderFetch)
    85  	tests := []struct {
    86  		query  *getBlockHeadersData // The query to execute for header retrieval
    87  		expect []common.Hash        // The hashes of the block whose headers are expected
    88  	}{
    89  		// A single random block should be retrievable by hash and number too
    90  		{
    91  			&getBlockHeadersData{Origin: hashOrNumber{Hash: pm.blockchain.GetBlockByNumber(limit / 2).Hash()}, Amount: 1},
    92  			[]common.Hash{pm.blockchain.GetBlockByNumber(limit / 2).Hash()},
    93  		}, {
    94  			&getBlockHeadersData{Origin: hashOrNumber{Number: limit / 2}, Amount: 1},
    95  			[]common.Hash{pm.blockchain.GetBlockByNumber(limit / 2).Hash()},
    96  		},
    97  		// Multiple headers should be retrievable in both directions
    98  		{
    99  			&getBlockHeadersData{Origin: hashOrNumber{Number: limit / 2}, Amount: 3},
   100  			[]common.Hash{
   101  				pm.blockchain.GetBlockByNumber(limit / 2).Hash(),
   102  				pm.blockchain.GetBlockByNumber(limit/2 + 1).Hash(),
   103  				pm.blockchain.GetBlockByNumber(limit/2 + 2).Hash(),
   104  			},
   105  		}, {
   106  			&getBlockHeadersData{Origin: hashOrNumber{Number: limit / 2}, Amount: 3, Reverse: true},
   107  			[]common.Hash{
   108  				pm.blockchain.GetBlockByNumber(limit / 2).Hash(),
   109  				pm.blockchain.GetBlockByNumber(limit/2 - 1).Hash(),
   110  				pm.blockchain.GetBlockByNumber(limit/2 - 2).Hash(),
   111  			},
   112  		},
   113  		// Multiple headers with skip lists should be retrievable
   114  		{
   115  			&getBlockHeadersData{Origin: hashOrNumber{Number: limit / 2}, Skip: 3, Amount: 3},
   116  			[]common.Hash{
   117  				pm.blockchain.GetBlockByNumber(limit / 2).Hash(),
   118  				pm.blockchain.GetBlockByNumber(limit/2 + 4).Hash(),
   119  				pm.blockchain.GetBlockByNumber(limit/2 + 8).Hash(),
   120  			},
   121  		}, {
   122  			&getBlockHeadersData{Origin: hashOrNumber{Number: limit / 2}, Skip: 3, Amount: 3, Reverse: true},
   123  			[]common.Hash{
   124  				pm.blockchain.GetBlockByNumber(limit / 2).Hash(),
   125  				pm.blockchain.GetBlockByNumber(limit/2 - 4).Hash(),
   126  				pm.blockchain.GetBlockByNumber(limit/2 - 8).Hash(),
   127  			},
   128  		},
   129  		// The chain endpoints should be retrievable
   130  		{
   131  			&getBlockHeadersData{Origin: hashOrNumber{Number: 0}, Amount: 1},
   132  			[]common.Hash{pm.blockchain.GetBlockByNumber(0).Hash()},
   133  		}, {
   134  			&getBlockHeadersData{Origin: hashOrNumber{Number: pm.blockchain.CurrentBlock().NumberU64()}, Amount: 1},
   135  			[]common.Hash{pm.blockchain.CurrentBlock().Hash()},
   136  		},
   137  		// Ensure protocol limits are honored
   138  		{
   139  			&getBlockHeadersData{Origin: hashOrNumber{Number: pm.blockchain.CurrentBlock().NumberU64() - 1}, Amount: limit + 10, Reverse: true},
   140  			pm.blockchain.GetBlockHashesFromHash(pm.blockchain.CurrentBlock().Hash(), limit),
   141  		},
   142  		// Check that requesting more than available is handled gracefully
   143  		{
   144  			&getBlockHeadersData{Origin: hashOrNumber{Number: pm.blockchain.CurrentBlock().NumberU64() - 4}, Skip: 3, Amount: 3},
   145  			[]common.Hash{
   146  				pm.blockchain.GetBlockByNumber(pm.blockchain.CurrentBlock().NumberU64() - 4).Hash(),
   147  				pm.blockchain.GetBlockByNumber(pm.blockchain.CurrentBlock().NumberU64()).Hash(),
   148  			},
   149  		}, {
   150  			&getBlockHeadersData{Origin: hashOrNumber{Number: 4}, Skip: 3, Amount: 3, Reverse: true},
   151  			[]common.Hash{
   152  				pm.blockchain.GetBlockByNumber(4).Hash(),
   153  				pm.blockchain.GetBlockByNumber(0).Hash(),
   154  			},
   155  		},
   156  		// Check that requesting more than available is handled gracefully, even if mid skip
   157  		{
   158  			&getBlockHeadersData{Origin: hashOrNumber{Number: pm.blockchain.CurrentBlock().NumberU64() - 4}, Skip: 2, Amount: 3},
   159  			[]common.Hash{
   160  				pm.blockchain.GetBlockByNumber(pm.blockchain.CurrentBlock().NumberU64() - 4).Hash(),
   161  				pm.blockchain.GetBlockByNumber(pm.blockchain.CurrentBlock().NumberU64() - 1).Hash(),
   162  			},
   163  		}, {
   164  			&getBlockHeadersData{Origin: hashOrNumber{Number: 4}, Skip: 2, Amount: 3, Reverse: true},
   165  			[]common.Hash{
   166  				pm.blockchain.GetBlockByNumber(4).Hash(),
   167  				pm.blockchain.GetBlockByNumber(1).Hash(),
   168  			},
   169  		},
   170  		// Check a corner case where requesting more can iterate past the endpoints
   171  		{
   172  			&getBlockHeadersData{Origin: hashOrNumber{Number: 2}, Amount: 5, Reverse: true},
   173  			[]common.Hash{
   174  				pm.blockchain.GetBlockByNumber(2).Hash(),
   175  				pm.blockchain.GetBlockByNumber(1).Hash(),
   176  				pm.blockchain.GetBlockByNumber(0).Hash(),
   177  			},
   178  		},
   179  		// Check a corner case where skipping overflow loops back into the chain start
   180  		{
   181  			&getBlockHeadersData{Origin: hashOrNumber{Hash: pm.blockchain.GetBlockByNumber(3).Hash()}, Amount: 2, Reverse: false, Skip: math.MaxUint64 - 1},
   182  			[]common.Hash{
   183  				pm.blockchain.GetBlockByNumber(3).Hash(),
   184  			},
   185  		},
   186  		// Check a corner case where skipping overflow loops back to the same header
   187  		{
   188  			&getBlockHeadersData{Origin: hashOrNumber{Hash: pm.blockchain.GetBlockByNumber(1).Hash()}, Amount: 2, Reverse: false, Skip: math.MaxUint64},
   189  			[]common.Hash{
   190  				pm.blockchain.GetBlockByNumber(1).Hash(),
   191  			},
   192  		},
   193  		// Check that non existing headers aren't returned
   194  		{
   195  			&getBlockHeadersData{Origin: hashOrNumber{Hash: unknown}, Amount: 1},
   196  			[]common.Hash{},
   197  		}, {
   198  			&getBlockHeadersData{Origin: hashOrNumber{Number: pm.blockchain.CurrentBlock().NumberU64() + 1}, Amount: 1},
   199  			[]common.Hash{},
   200  		},
   201  	}
   202  	// Run each of the tests and verify the results against the chain
   203  	for i, tt := range tests {
   204  		// Collect the headers to expect in the response
   205  		headers := []*types.Header{}
   206  		for _, hash := range tt.expect {
   207  			headers = append(headers, pm.blockchain.GetBlockByHash(hash).Header())
   208  		}
   209  		// Send the hash request and verify the response
   210  		p2p.Send(peer.app, 0x03, tt.query)
   211  		if err := p2p.ExpectMsg(peer.app, 0x04, headers); err != nil {
   212  			t.Errorf("test %d: headers mismatch: %v", i, err)
   213  		}
   214  		// If the test used number origins, repeat with hashes as the too
   215  		if tt.query.Origin.Hash == (common.Hash{}) {
   216  			if origin := pm.blockchain.GetBlockByNumber(tt.query.Origin.Number); origin != nil {
   217  				tt.query.Origin.Hash, tt.query.Origin.Number = origin.Hash(), 0
   218  
   219  				p2p.Send(peer.app, 0x03, tt.query)
   220  				if err := p2p.ExpectMsg(peer.app, 0x04, headers); err != nil {
   221  					t.Errorf("test %d: headers mismatch: %v", i, err)
   222  				}
   223  			}
   224  		}
   225  	}
   226  }
   227  
   228  // Tests that block contents can be retrieved from a remote chain based on their hashes.
   229  func TestGetBlockBodies62(t *testing.T) { testGetBlockBodies(t, 62) }
   230  func TestGetBlockBodies63(t *testing.T) { testGetBlockBodies(t, 63) }
   231  
   232  func testGetBlockBodies(t *testing.T, protocol int) {
   233  	pm, _ := newTestProtocolManagerMust(t, downloader.FullSync, downloader.MaxBlockFetch+15, nil, nil)
   234  	peer, _ := newTestPeer("peer", protocol, pm, true)
   235  	defer peer.close()
   236  
   237  	// Create a batch of tests for various scenarios
   238  	limit := downloader.MaxBlockFetch
   239  	tests := []struct {
   240  		random    int           // Number of blocks to fetch randomly from the chain
   241  		explicit  []common.Hash // Explicitly requested blocks
   242  		available []bool        // Availability of explicitly requested blocks
   243  		expected  int           // Total number of existing blocks to expect
   244  	}{
   245  		{1, nil, nil, 1},                                                         // A single random block should be retrievable
   246  		{10, nil, nil, 10},                                                       // Multiple random blocks should be retrievable
   247  		{limit, nil, nil, limit},                                                 // The maximum possible blocks should be retrievable
   248  		{limit + 1, nil, nil, limit},                                             // No more than the possible block count should be returned
   249  		{0, []common.Hash{pm.blockchain.Genesis().Hash()}, []bool{true}, 1},      // The genesis block should be retrievable
   250  		{0, []common.Hash{pm.blockchain.CurrentBlock().Hash()}, []bool{true}, 1}, // The chains head block should be retrievable
   251  		{0, []common.Hash{{}}, []bool{false}, 0},                                 // A non existent block should not be returned
   252  
   253  		// Existing and non-existing blocks interleaved should not cause problems
   254  		{0, []common.Hash{
   255  			{},
   256  			pm.blockchain.GetBlockByNumber(1).Hash(),
   257  			{},
   258  			pm.blockchain.GetBlockByNumber(10).Hash(),
   259  			{},
   260  			pm.blockchain.GetBlockByNumber(100).Hash(),
   261  			{},
   262  		}, []bool{false, true, false, true, false, true, false}, 3},
   263  	}
   264  	// Run each of the tests and verify the results against the chain
   265  	for i, tt := range tests {
   266  		// Collect the hashes to request, and the response to expect
   267  		hashes, seen := []common.Hash{}, make(map[int64]bool)
   268  		bodies := []*blockBody{}
   269  
   270  		for j := 0; j < tt.random; j++ {
   271  			for {
   272  				num := rand.Int63n(int64(pm.blockchain.CurrentBlock().NumberU64()))
   273  				if !seen[num] {
   274  					seen[num] = true
   275  
   276  					block := pm.blockchain.GetBlockByNumber(uint64(num))
   277  					hashes = append(hashes, block.Hash())
   278  					if len(bodies) < tt.expected {
   279  						bodies = append(bodies, &blockBody{Transactions: block.Transactions(), Uncles: block.Uncles()})
   280  					}
   281  					break
   282  				}
   283  			}
   284  		}
   285  		for j, hash := range tt.explicit {
   286  			hashes = append(hashes, hash)
   287  			if tt.available[j] && len(bodies) < tt.expected {
   288  				block := pm.blockchain.GetBlockByHash(hash)
   289  				bodies = append(bodies, &blockBody{Transactions: block.Transactions(), Uncles: block.Uncles()})
   290  			}
   291  		}
   292  		// Send the hash request and verify the response
   293  		p2p.Send(peer.app, 0x05, hashes)
   294  		if err := p2p.ExpectMsg(peer.app, 0x06, bodies); err != nil {
   295  			t.Errorf("test %d: bodies mismatch: %v", i, err)
   296  		}
   297  	}
   298  }
   299  
   300  // Tests that the node state database can be retrieved based on hashes.
   301  func TestGetNodeData63(t *testing.T) { testGetNodeData(t, 63) }
   302  
   303  func testGetNodeData(t *testing.T, protocol int) {
   304  	// Define three accounts to simulate transactions with
   305  	acc1Key, _ := crypto.HexToECDSA("8a1f9a8f95be41cd7ccb6168179afb4504aefe388d1e14474d32c45c72ce7b7a")
   306  	acc2Key, _ := crypto.HexToECDSA("49a7b37aa6f6645917e7b807e9d1c00d4fa71f18343b0d4122a4d2df64dd6fee")
   307  	acc1Addr := crypto.PubkeyToAddress(acc1Key.PublicKey)
   308  	acc2Addr := crypto.PubkeyToAddress(acc2Key.PublicKey)
   309  
   310  	signer := types.HomesteadSigner{}
   311  	// Create a chain generator with some simple transactions (blatantly stolen from @fjl/chain_markets_test)
   312  	generator := func(i int, block *core.BlockGen) {
   313  		switch i {
   314  		case 0:
   315  			// In block 1, the test bank sends account #1 some ether.
   316  			tx, _ := types.SignTx(types.NewTransaction(block.TxNonce(testBank), acc1Addr, big.NewInt(10000), params.TxGas, nil, nil), signer, testBankKey)
   317  			block.AddTx(tx)
   318  		case 1:
   319  			// In block 2, the test bank sends some more ether to account #1.
   320  			// acc1Addr passes it on to account #2.
   321  			tx1, _ := types.SignTx(types.NewTransaction(block.TxNonce(testBank), acc1Addr, big.NewInt(1000), params.TxGas, nil, nil), signer, testBankKey)
   322  			tx2, _ := types.SignTx(types.NewTransaction(block.TxNonce(acc1Addr), acc2Addr, big.NewInt(1000), params.TxGas, nil, nil), signer, acc1Key)
   323  			block.AddTx(tx1)
   324  			block.AddTx(tx2)
   325  		case 2:
   326  			// Block 3 is empty but was mined by account #2.
   327  			block.SetCoinbase(acc2Addr)
   328  			block.SetExtra([]byte("yeehaw"))
   329  		case 3:
   330  			// Block 4 includes blocks 2 and 3 as uncle headers (with modified extra data).
   331  			b2 := block.PrevBlock(1).Header()
   332  			b2.Extra = []byte("foo")
   333  			block.AddUncle(b2)
   334  			b3 := block.PrevBlock(2).Header()
   335  			b3.Extra = []byte("foo")
   336  			block.AddUncle(b3)
   337  		}
   338  	}
   339  	// Assemble the test environment
   340  	pm, db := newTestProtocolManagerMust(t, downloader.FullSync, 4, generator, nil)
   341  	peer, _ := newTestPeer("peer", protocol, pm, true)
   342  	defer peer.close()
   343  
   344  	// Fetch for now the entire chain db
   345  	hashes := []common.Hash{}
   346  	for _, key := range db.Keys() {
   347  		if len(key) == len(common.Hash{}) {
   348  			hashes = append(hashes, common.BytesToHash(key))
   349  		}
   350  	}
   351  	p2p.Send(peer.app, 0x0d, hashes)
   352  	msg, err := peer.app.ReadMsg()
   353  	if err != nil {
   354  		t.Fatalf("failed to read node data response: %v", err)
   355  	}
   356  	if msg.Code != 0x0e {
   357  		t.Fatalf("response packet code mismatch: have %x, want %x", msg.Code, 0x0c)
   358  	}
   359  	var data [][]byte
   360  	if err := msg.Decode(&data); err != nil {
   361  		t.Fatalf("failed to decode response node data: %v", err)
   362  	}
   363  	// Verify that all hashes correspond to the requested data, and reconstruct a state tree
   364  	for i, want := range hashes {
   365  		if hash := crypto.Keccak256Hash(data[i]); hash != want {
   366  			t.Errorf("data hash mismatch: have %x, want %x", hash, want)
   367  		}
   368  	}
   369  	statedb, _ := ethdb.NewMemDatabase()
   370  	for i := 0; i < len(data); i++ {
   371  		statedb.Put(hashes[i].Bytes(), data[i])
   372  	}
   373  	accounts := []common.Address{testBank, acc1Addr, acc2Addr}
   374  	for i := uint64(0); i <= pm.blockchain.CurrentBlock().NumberU64(); i++ {
   375  		trie, _ := state.New(pm.blockchain.GetBlockByNumber(i).Root(), state.NewDatabase(statedb))
   376  
   377  		for j, acc := range accounts {
   378  			state, _ := pm.blockchain.State()
   379  			bw := state.GetBalance(acc)
   380  			bh := trie.GetBalance(acc)
   381  
   382  			if (bw != nil && bh == nil) || (bw == nil && bh != nil) {
   383  				t.Errorf("test %d, account %d: balance mismatch: have %v, want %v", i, j, bh, bw)
   384  			}
   385  			if bw != nil && bh != nil && bw.Cmp(bw) != 0 {
   386  				t.Errorf("test %d, account %d: balance mismatch: have %v, want %v", i, j, bh, bw)
   387  			}
   388  		}
   389  	}
   390  }
   391  
   392  // Tests that the transaction receipts can be retrieved based on hashes.
   393  func TestGetReceipt63(t *testing.T) { testGetReceipt(t, 63) }
   394  
   395  func testGetReceipt(t *testing.T, protocol int) {
   396  	// Define three accounts to simulate transactions with
   397  	acc1Key, _ := crypto.HexToECDSA("8a1f9a8f95be41cd7ccb6168179afb4504aefe388d1e14474d32c45c72ce7b7a")
   398  	acc2Key, _ := crypto.HexToECDSA("49a7b37aa6f6645917e7b807e9d1c00d4fa71f18343b0d4122a4d2df64dd6fee")
   399  	acc1Addr := crypto.PubkeyToAddress(acc1Key.PublicKey)
   400  	acc2Addr := crypto.PubkeyToAddress(acc2Key.PublicKey)
   401  
   402  	signer := types.HomesteadSigner{}
   403  	// Create a chain generator with some simple transactions (blatantly stolen from @fjl/chain_markets_test)
   404  	generator := func(i int, block *core.BlockGen) {
   405  		switch i {
   406  		case 0:
   407  			// In block 1, the test bank sends account #1 some ether.
   408  			tx, _ := types.SignTx(types.NewTransaction(block.TxNonce(testBank), acc1Addr, big.NewInt(10000), params.TxGas, nil, nil), signer, testBankKey)
   409  			block.AddTx(tx)
   410  		case 1:
   411  			// In block 2, the test bank sends some more ether to account #1.
   412  			// acc1Addr passes it on to account #2.
   413  			tx1, _ := types.SignTx(types.NewTransaction(block.TxNonce(testBank), acc1Addr, big.NewInt(1000), params.TxGas, nil, nil), signer, testBankKey)
   414  			tx2, _ := types.SignTx(types.NewTransaction(block.TxNonce(acc1Addr), acc2Addr, big.NewInt(1000), params.TxGas, nil, nil), signer, acc1Key)
   415  			block.AddTx(tx1)
   416  			block.AddTx(tx2)
   417  		case 2:
   418  			// Block 3 is empty but was mined by account #2.
   419  			block.SetCoinbase(acc2Addr)
   420  			block.SetExtra([]byte("yeehaw"))
   421  		case 3:
   422  			// Block 4 includes blocks 2 and 3 as uncle headers (with modified extra data).
   423  			b2 := block.PrevBlock(1).Header()
   424  			b2.Extra = []byte("foo")
   425  			block.AddUncle(b2)
   426  			b3 := block.PrevBlock(2).Header()
   427  			b3.Extra = []byte("foo")
   428  			block.AddUncle(b3)
   429  		}
   430  	}
   431  	// Assemble the test environment
   432  	pm, _ := newTestProtocolManagerMust(t, downloader.FullSync, 4, generator, nil)
   433  	peer, _ := newTestPeer("peer", protocol, pm, true)
   434  	defer peer.close()
   435  
   436  	// Collect the hashes to request, and the response to expect
   437  	hashes, receipts := []common.Hash{}, []types.Receipts{}
   438  	for i := uint64(0); i <= pm.blockchain.CurrentBlock().NumberU64(); i++ {
   439  		block := pm.blockchain.GetBlockByNumber(i)
   440  
   441  		hashes = append(hashes, block.Hash())
   442  		receipts = append(receipts, pm.blockchain.GetReceiptsByHash(block.Hash()))
   443  	}
   444  	// Send the hash request and verify the response
   445  	p2p.Send(peer.app, 0x0f, hashes)
   446  	if err := p2p.ExpectMsg(peer.app, 0x10, receipts); err != nil {
   447  		t.Errorf("receipts mismatch: %v", err)
   448  	}
   449  }
   450  
   451  // Tests that post eth protocol handshake, DAO fork-enabled clients also execute
   452  // a DAO "challenge" verifying each others' DAO fork headers to ensure they're on
   453  // compatible chains.
   454  func TestDAOChallengeNoVsNo(t *testing.T)       { testDAOChallenge(t, false, false, false) }
   455  func TestDAOChallengeNoVsPro(t *testing.T)      { testDAOChallenge(t, false, true, false) }
   456  func TestDAOChallengeProVsNo(t *testing.T)      { testDAOChallenge(t, true, false, false) }
   457  func TestDAOChallengeProVsPro(t *testing.T)     { testDAOChallenge(t, true, true, false) }
   458  func TestDAOChallengeNoVsTimeout(t *testing.T)  { testDAOChallenge(t, false, false, true) }
   459  func TestDAOChallengeProVsTimeout(t *testing.T) { testDAOChallenge(t, true, true, true) }
   460  
   461  func testDAOChallenge(t *testing.T, localForked, remoteForked bool, timeout bool) {
   462  	// Reduce the DAO handshake challenge timeout
   463  	if timeout {
   464  		defer func(old time.Duration) { daoChallengeTimeout = old }(daoChallengeTimeout)
   465  		daoChallengeTimeout = 500 * time.Millisecond
   466  	}
   467  	// Create a DAO aware protocol manager
   468  	var (
   469  		evmux         = new(event.TypeMux)
   470  		pow           = ethash.NewFaker()
   471  		db, _         = ethdb.NewMemDatabase()
   472  		config        = &params.ChainConfig{DAOForkBlock: big.NewInt(1), DAOForkSupport: localForked}
   473  		gspec         = &core.Genesis{Config: config}
   474  		genesis       = gspec.MustCommit(db)
   475  		blockchain, _ = core.NewBlockChain(db, nil, config, pow, vm.Config{})
   476  	)
   477  	pm, err := NewProtocolManager(config, downloader.FullSync, DefaultConfig.NetworkId, evmux, new(testTxPool), pow, blockchain, db)
   478  	if err != nil {
   479  		t.Fatalf("failed to start test protocol manager: %v", err)
   480  	}
   481  	pm.Start(1000)
   482  	defer pm.Stop()
   483  
   484  	// Connect a new peer and check that we receive the DAO challenge
   485  	peer, _ := newTestPeer("peer", eth63, pm, true)
   486  	defer peer.close()
   487  
   488  	challenge := &getBlockHeadersData{
   489  		Origin:  hashOrNumber{Number: config.DAOForkBlock.Uint64()},
   490  		Amount:  1,
   491  		Skip:    0,
   492  		Reverse: false,
   493  	}
   494  	if err := p2p.ExpectMsg(peer.app, GetBlockHeadersMsg, challenge); err != nil {
   495  		t.Fatalf("challenge mismatch: %v", err)
   496  	}
   497  	// Create a block to reply to the challenge if no timeout is simulated
   498  	if !timeout {
   499  		blocks, _ := core.GenerateChain(&params.ChainConfig{}, genesis, ethash.NewFaker(), db, 1, func(i int, block *core.BlockGen) {
   500  			if remoteForked {
   501  				block.SetExtra(params.DAOForkBlockExtra)
   502  			}
   503  		})
   504  		if err := p2p.Send(peer.app, BlockHeadersMsg, []*types.Header{blocks[0].Header()}); err != nil {
   505  			t.Fatalf("failed to answer challenge: %v", err)
   506  		}
   507  		time.Sleep(100 * time.Millisecond) // Sleep to avoid the verification racing with the drops
   508  	} else {
   509  		// Otherwise wait until the test timeout passes
   510  		time.Sleep(daoChallengeTimeout + 500*time.Millisecond)
   511  	}
   512  	// Verify that depending on fork side, the remote peer is maintained or dropped
   513  	if localForked == remoteForked && !timeout {
   514  		if peers := pm.peers.Len(); peers != 1 {
   515  			t.Fatalf("peer count mismatch: have %d, want %d", peers, 1)
   516  		}
   517  	} else {
   518  		if peers := pm.peers.Len(); peers != 0 {
   519  			t.Fatalf("peer count mismatch: have %d, want %d", peers, 0)
   520  		}
   521  	}
   522  }