github.com/prysmaticlabs/prysm@v1.4.4/tools/beacon-fuzz/main.go (about)

     1  package main
     2  
     3  import (
     4  	"bytes"
     5  	"flag"
     6  	"fmt"
     7  	"io/ioutil"
     8  	"os"
     9  	"path/filepath"
    10  	"strconv"
    11  	"text/template"
    12  
    13  	"github.com/prysmaticlabs/prysm/shared/fileutil"
    14  )
    15  
    16  var (
    17  	output = flag.String("output", "", "Output filepath for generated states file.")
    18  )
    19  
    20  const tpl = `// Code generated by //tools/beacon-fuzz:beacon-fuzz. DO NOT EDIT.
    21  package {{.Package}}
    22  
    23  // generateStates is a map of generated states from ssz.
    24  var generatedStates = map[uint16]string{
    25  	{{.MapStr}}
    26  }
    27  `
    28  
    29  // This program generates a map of ID(uint16) -> hex encoded strings of SSZ binary data. This tool
    30  // exists to facilitate running beacon-fuzz targets within the constraints of fuzzit. I.e. fuzz
    31  // only loads the corpus to the file system without the beacon states. An alternative approach would
    32  // be to create a docker image where the state files are available or changing the corpus seed data
    33  // to contain the beacon state itself.
    34  func main() {
    35  	flag.Parse()
    36  	if *output == "" {
    37  		panic("Missing output. Usage: beacon-fuzz --output=out.go path/to/state/0 path/to/state/1 ...")
    38  	}
    39  	statePaths := os.Args[2:]
    40  
    41  	if len(statePaths) > 15 {
    42  		statePaths = statePaths[:15]
    43  	}
    44  
    45  	m := make(map[int][]byte, len(statePaths))
    46  	for _, p := range statePaths {
    47  		ID, err := strconv.Atoi(filepath.Base(p))
    48  		if err != nil {
    49  			panic(fmt.Sprintf("%s does not end in an integer for the filename.", p))
    50  		}
    51  		b, err := ioutil.ReadFile(p)
    52  		if err != nil {
    53  			panic(err)
    54  		}
    55  		m[ID] = b
    56  	}
    57  
    58  	res := execTmpl(tpl, input{Package: "testing", MapStr: sszBytesToMapStr(m)})
    59  	if err := fileutil.WriteFile(*output, res.Bytes()); err != nil {
    60  		panic(err)
    61  	}
    62  }
    63  
    64  func sszBytesToMapStr(ss map[int][]byte) string {
    65  	dst := ""
    66  	for i, s := range ss {
    67  		dst += fmt.Sprintf("%d: \"%x\",", i, s)
    68  	}
    69  	return dst
    70  }
    71  
    72  type input struct {
    73  	Package string
    74  	States  map[int][]byte
    75  	MapStr  string
    76  }
    77  
    78  func execTmpl(tpl string, input interface{}) *bytes.Buffer {
    79  	tmpl, err := template.New("template").Parse(tpl)
    80  	if err != nil {
    81  		panic(err)
    82  	}
    83  	buf := new(bytes.Buffer)
    84  	if err = tmpl.Execute(buf, input); err != nil {
    85  		panic(err)
    86  	}
    87  	return buf
    88  }