github.com/ethereum/go-ethereum@v1.16.1/eth/protocols/snap/handler_fuzzing_test.go (about)

     1  // Copyright 2021 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 snap
    18  
    19  import (
    20  	"bytes"
    21  	"encoding/binary"
    22  	"fmt"
    23  	"math/big"
    24  	"testing"
    25  	"time"
    26  
    27  	"github.com/ethereum/go-ethereum/common"
    28  	"github.com/ethereum/go-ethereum/consensus/ethash"
    29  	"github.com/ethereum/go-ethereum/core"
    30  	"github.com/ethereum/go-ethereum/core/rawdb"
    31  	"github.com/ethereum/go-ethereum/core/types"
    32  	"github.com/ethereum/go-ethereum/p2p"
    33  	"github.com/ethereum/go-ethereum/p2p/enode"
    34  	"github.com/ethereum/go-ethereum/params"
    35  	"github.com/ethereum/go-ethereum/rlp"
    36  	fuzz "github.com/google/gofuzz"
    37  )
    38  
    39  func FuzzARange(f *testing.F) {
    40  	f.Fuzz(func(t *testing.T, data []byte) {
    41  		doFuzz(data, &GetAccountRangePacket{}, GetAccountRangeMsg)
    42  	})
    43  }
    44  
    45  func FuzzSRange(f *testing.F) {
    46  	f.Fuzz(func(t *testing.T, data []byte) {
    47  		doFuzz(data, &GetStorageRangesPacket{}, GetStorageRangesMsg)
    48  	})
    49  }
    50  
    51  func FuzzByteCodes(f *testing.F) {
    52  	f.Fuzz(func(t *testing.T, data []byte) {
    53  		doFuzz(data, &GetByteCodesPacket{}, GetByteCodesMsg)
    54  	})
    55  }
    56  
    57  func FuzzTrieNodes(f *testing.F) {
    58  	f.Fuzz(func(t *testing.T, data []byte) {
    59  		doFuzz(data, &GetTrieNodesPacket{}, GetTrieNodesMsg)
    60  	})
    61  }
    62  
    63  func doFuzz(input []byte, obj interface{}, code int) {
    64  	bc := getChain()
    65  	defer bc.Stop()
    66  	fuzz.NewFromGoFuzz(input).Fuzz(obj)
    67  	var data []byte
    68  	switch p := obj.(type) {
    69  	case *GetTrieNodesPacket:
    70  		p.Root = trieRoot
    71  		data, _ = rlp.EncodeToBytes(obj)
    72  	default:
    73  		data, _ = rlp.EncodeToBytes(obj)
    74  	}
    75  	cli := &dummyRW{
    76  		code: uint64(code),
    77  		data: data,
    78  	}
    79  	peer := NewFakePeer(65, "gazonk01", cli)
    80  	err := HandleMessage(&dummyBackend{bc}, peer)
    81  	switch {
    82  	case err == nil && cli.writeCount != 1:
    83  		panic(fmt.Sprintf("Expected 1 response, got %d", cli.writeCount))
    84  	case err != nil && cli.writeCount != 0:
    85  		panic(fmt.Sprintf("Expected 0 response, got %d", cli.writeCount))
    86  	}
    87  }
    88  
    89  var trieRoot common.Hash
    90  
    91  func getChain() *core.BlockChain {
    92  	ga := make(types.GenesisAlloc, 1000)
    93  	var a = make([]byte, 20)
    94  	var mkStorage = func(k, v int) (common.Hash, common.Hash) {
    95  		var kB = make([]byte, 32)
    96  		var vB = make([]byte, 32)
    97  		binary.LittleEndian.PutUint64(kB, uint64(k))
    98  		binary.LittleEndian.PutUint64(vB, uint64(v))
    99  		return common.BytesToHash(kB), common.BytesToHash(vB)
   100  	}
   101  	storage := make(map[common.Hash]common.Hash)
   102  	for i := 0; i < 10; i++ {
   103  		k, v := mkStorage(i, i)
   104  		storage[k] = v
   105  	}
   106  	for i := 0; i < 1000; i++ {
   107  		binary.LittleEndian.PutUint64(a, uint64(i+0xff))
   108  		acc := types.Account{Balance: big.NewInt(int64(i))}
   109  		if i%2 == 1 {
   110  			acc.Storage = storage
   111  		}
   112  		ga[common.BytesToAddress(a)] = acc
   113  	}
   114  	gspec := &core.Genesis{
   115  		Config: params.TestChainConfig,
   116  		Alloc:  ga,
   117  	}
   118  	_, blocks, _ := core.GenerateChainWithGenesis(gspec, ethash.NewFaker(), 2, func(i int, gen *core.BlockGen) {})
   119  	options := &core.BlockChainConfig{
   120  		TrieCleanLimit: 0,
   121  		TrieDirtyLimit: 0,
   122  		TrieTimeLimit:  5 * time.Minute,
   123  		NoPrefetch:     true,
   124  		SnapshotLimit:  100,
   125  		SnapshotWait:   true,
   126  	}
   127  	trieRoot = blocks[len(blocks)-1].Root()
   128  	bc, _ := core.NewBlockChain(rawdb.NewMemoryDatabase(), gspec, ethash.NewFaker(), options)
   129  	if _, err := bc.InsertChain(blocks); err != nil {
   130  		panic(err)
   131  	}
   132  	return bc
   133  }
   134  
   135  type dummyBackend struct {
   136  	chain *core.BlockChain
   137  }
   138  
   139  func (d *dummyBackend) Chain() *core.BlockChain       { return d.chain }
   140  func (d *dummyBackend) RunPeer(*Peer, Handler) error  { return nil }
   141  func (d *dummyBackend) PeerInfo(enode.ID) interface{} { return "Foo" }
   142  func (d *dummyBackend) Handle(*Peer, Packet) error    { return nil }
   143  
   144  type dummyRW struct {
   145  	code       uint64
   146  	data       []byte
   147  	writeCount int
   148  }
   149  
   150  func (d *dummyRW) ReadMsg() (p2p.Msg, error) {
   151  	return p2p.Msg{
   152  		Code:       d.code,
   153  		Payload:    bytes.NewReader(d.data),
   154  		ReceivedAt: time.Now(),
   155  		Size:       uint32(len(d.data)),
   156  	}, nil
   157  }
   158  
   159  func (d *dummyRW) WriteMsg(msg p2p.Msg) error {
   160  	d.writeCount++
   161  	return nil
   162  }