convert_grep_to_fuzzer.go

/*
Copyright 2021 The Vitess Authors.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
    http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

/*
This file is used to autogenerate a fuzzer for Vitess' API marshaling routines.

To use this file:
git clone https://github.com/vitessio/vitess
cd vitess/go/vt
grep -r ') Unmarshal' .>>/tmp/marshal_targets.txt
cd ../test/fuzzing/autogenerate
go run convert_grep_to_fuzzer.go
# The fuzzer has now been built.
# To instruct OSS-fuzz to build the autogenerated fuzzer:
mv api_marshal_fuzzer.go ../
compile_go_fuzzer vitess.io/vitess/go/test/fuzzing FuzzAPIMarshal api_marshal_fuzzer
*/

package main

import (
	"bufio"
	"fmt"
	"os"
	"strings"
)

// importPathShort is a map to convert import paths to import aliases
var importPathShort = map[string]string{
	"vitess.io/vitess/go/vt/proto/query":             "query",
	"vitess.io/vitess/go/vt/proto/tabletmanagerdata": "tabletmanagerdata",
	"vitess.io/vitess/go/vt/proto/binlogdata":        "binlogdata",
	"vitess.io/vitess/go/vt/proto/vtgate":            "vtgate",
	"vitess.io/vitess/go/vt/proto/automation":        "automation",
	"vitess.io/vitess/go/vt/proto/tableacl":          "tableacl",
	"vitess.io/vitess/go/vt/proto/logutil":           "logutil",
	"vitess.io/vitess/go/vt/proto/vttest":            "vttest",
	"vitess.io/vitess/go/vt/proto/vttime":            "vttime",
	"vitess.io/vitess/go/vt/proto/workflow":          "workflow",
	"vitess.io/vitess/go/vt/proto/replicationdata":   "replicationdata",
	"vitess.io/vitess/go/vt/proto/vtctldata":         "vtctldata",
	"vitess.io/vitess/go/vt/proto/vtrpc":             "vtrpc",
	"vitess.io/vitess/go/vt/proto/vschema":           "vschema",
	"vitess.io/vitess/go/vt/proto/mysqlctl":          "mysqlctl",
	"vitess.io/vitess/go/vt/proto/vtadmin":           "vtadmin",
	"vitess.io/vitess/go/vt/proto/throttlerdata":     "throttlerdata",
	"vitess.io/vitess/go/vt/proto/topodata":          "topodata",
}

// pathToImportPath is a map to convert filepaths to import paths.
// the filepaths are available in the grep results.
var pathToImportPath = map[string]string{
	"./proto/query/query_vtproto.pb.go":                         "vitess.io/vitess/go/vt/proto/query",
	"./proto/tabletmanagerdata/tabletmanagerdata_vtproto.pb.go": "vitess.io/vitess/go/vt/proto/tabletmanagerdata",
	"./proto/binlogdata/binlogdata_vtproto.pb.go":               "vitess.io/vitess/go/vt/proto/binlogdata",
	"./proto/vtgate/vtgate_vtproto.pb.go":                       "vitess.io/vitess/go/vt/proto/vtgate",
	"./proto/automation/automation_vtproto.pb.go":               "vitess.io/vitess/go/vt/proto/automation",
	"./proto/tableacl/tableacl_vtproto.pb.go":                   "vitess.io/vitess/go/vt/proto/tableacl",
	"./proto/logutil/logutil_vtproto.pb.go":                     "vitess.io/vitess/go/vt/proto/logutil",
	"./proto/vttest/vttest_vtproto.pb.go":                       "vitess.io/vitess/go/vt/proto/vttest",
	"./proto/vttime/vttime_vtproto.pb.go":                       "vitess.io/vitess/go/vt/proto/vttime",
	"./proto/workflow/workflow_vtproto.pb.go":                   "vitess.io/vitess/go/vt/proto/workflow",
	"./proto/replicationdata/replicationdata_vtproto.pb.go":     "vitess.io/vitess/go/vt/proto/replicationdata",
	"./proto/vtctldata/vtctldata_vtproto.pb.go":                 "vitess.io/vitess/go/vt/proto/vtctldata",
	"./proto/vtrpc/vtrpc_vtproto.pb.go":                         "vitess.io/vitess/go/vt/proto/vtrpc",
	"./proto/vschema/vschema_vtproto.pb.go":                     "vitess.io/vitess/go/vt/proto/vschema",
	"./proto/mysqlctl/mysqlctl_vtproto.pb.go":                   "vitess.io/vitess/go/vt/proto/mysqlctl",
	"./proto/vtadmin/vtadmin_vtproto.pb.go":                     "vitess.io/vitess/go/vt/proto/vtadmin",
	"./proto/throttlerdata/throttlerdata_vtproto.pb.go":         "vitess.io/vitess/go/vt/proto/throttlerdata",
	"./proto/topodata/topodata_vtproto.pb.go":                   "vitess.io/vitess/go/vt/proto/topodata",
}

// contains checks if a string is present in a string slice
func contains(s []string, e string) bool {
	for _, a := range s {
		if a == e {
			return true
		}
	}
	return false
}

/*
createFunctionCall creates a function call

This is used to create the calls to the internal harnesses:
switch funcOp {
case 0:

	_ = FuzzqueryTarget(data2) <- created by createFunctionCall

case 1:

	_ = FuzzqueryVTGateCallerID(data2) <- created by createFunctionCall

case 2:

	        _ = FuzzqueryEventToken(data2) <- created by createFunctionCall
	}
*/
func createFunctionCall(shortName, structName string) string {
	return fmt.Sprintf("\t_ = Fuzz%s%s(data2)\n", shortName, structName)
}

/*
createHarness creates an internal harness.
An example of a generated harness:

	func FuzztabletmanagerdataTableDefinition(data []byte) error {
	        f := fuzz.NewConsumer(data)
	        s := &tabletmanagerdata.TableDefinition{}
	        err := f.GenerateStruct(s)
	        if err != nil {
	                return err
	        }
	        b, err := s.MarshalVT()
	        if err != nil {
	                return err
	        }
	        s2 := &tabletmanagerdata.TableDefinition{}
	        err = s2.UnmarshalVT(b)
	        if err != nil {
	                return err
	        }
	        newBytes, err := f.GetBytes()
	        if err != nil {
	                return err
	        }
	        s3 := &tabletmanagerdata.TableDefinition{}
	        err = s3.UnmarshalVT(newBytes)
	        return err
	}
*/
func createHarness(shortName, structName string) string {
	var harnessString strings.Builder
	harnessString.WriteString(fmt.Sprintf("\n\nfunc Fuzz%s%s(data []byte) error {\n", shortName, structName))
	harnessString.WriteString("\tf := fuzz.NewConsumer(data)\n")
	harnessString.WriteString(fmt.Sprintf("\ts := &%s.%s{}\n", shortName, structName))
	harnessString.WriteString("\terr := f.GenerateStruct(s)\n")
	harnessString.WriteString("\tif err != nil {\n")
	harnessString.WriteString("\t\treturn err\n")
	harnessString.WriteString("\t}\n")
	harnessString.WriteString("\tb, err := s.MarshalVT()\n")
	harnessString.WriteString("\tif err != nil {\n")
	harnessString.WriteString("\t\treturn err\n")
	harnessString.WriteString("\t}\n")
	harnessString.WriteString(fmt.Sprintf("\ts2 := &%s.%s{}\n", shortName, structName))
	harnessString.WriteString("\terr = s2.UnmarshalVT(b)\n")
	harnessString.WriteString("\tif err != nil {\n")
	harnessString.WriteString("\t\treturn err\n")
	harnessString.WriteString("\t}\n")
	harnessString.WriteString("\tnewBytes, err := f.GetBytes()\n")
	harnessString.WriteString("\tif err != nil {\n")
	harnessString.WriteString("\t\treturn err\n")
	harnessString.WriteString("\t}\n")
	harnessString.WriteString(fmt.Sprintf("\ts3 := &%s.%s{}\n", shortName, structName))
	harnessString.WriteString("\terr = s3.UnmarshalVT(newBytes)\n")
	harnessString.WriteString("\treturn err\n")
	harnessString.WriteString("}\n")
	return harnessString.String()
}

// createMainFuzzer creates a file and writes all
// the accumulated sub-elements (imports, main harness,
// internal harnesses) of the fuzzer to it.
func createMainFuzzer(functionList, harnesses []string) {
	var mainFuzzer strings.Builder

	// package and imports
	mainFuzzer.WriteString("package fuzzing\n\n")
	mainFuzzer.WriteString("import (\n")
	for k, v := range importPathShort {
		mainFuzzer.WriteString(fmt.Sprintf("\t%s \"%s\"\n", v, k))
	}
	mainFuzzer.WriteString("\tfuzz \"github.com/AdaLogics/go-fuzz-headers\"\n")
	mainFuzzer.WriteString(")\n\n")

	// main entrypoint of the fuzzer
	mainFuzzer.WriteString("func FuzzAPIMarshal(data []byte) int {\n")
	mainFuzzer.WriteString("\tif len(data)<10{ return 0 }\n")

	// calls to each internal harness
	maxOps := len(functionList)
	mainFuzzer.WriteString(fmt.Sprintf("\tfuncOp := int(data[0])%%%d\n", maxOps))
	mainFuzzer.WriteString("\tdata2 := data[1:]\n")
	mainFuzzer.WriteString("\tswitch funcOp {\n")
	for i := 0; i < len(functionList); i++ {
		mainFuzzer.WriteString(fmt.Sprintf("\tcase %d:\n", i))
		mainFuzzer.WriteString(fmt.Sprintf("\t%s\n", functionList[i]))
	}
	mainFuzzer.WriteString("\t}\n")
	mainFuzzer.WriteString("\treturn 1\n")
	mainFuzzer.WriteString("}")

	// add all the internal harnesses
	for i := 0; i < len(harnesses); i++ {
		mainFuzzer.WriteString(harnesses[i])
	}

	// write the file
	fuzzFile, err := os.Create("api_marshal_fuzzer.go")
	if err != nil {
		panic(err)
	}
	defer fuzzFile.Close()
	fuzzFile.WriteString(mainFuzzer.String())
}

// getGrepData extracts the required data from the grep result
func getGrepData() ([]string, []string) {
	// paths contains all import paths
	paths := make([]string, 0)

	// functionList contains the calls to all the harnesses
	functionList := make([]string, 0)

	// harnesses contains all the harnesses
	harnesses := make([]string, 0)

	// extract data from the grep results:
	f, _ := os.Open("/tmp/marshal_targets.txt")
	defer f.Close()
	scanner := bufio.NewScanner(f)
	for scanner.Scan() {
		line := scanner.Text()
		if strings.Contains(line, "func (") {
			thePath := strings.Split(line, ":")[0]
			if !contains(paths, thePath) {
				paths = append(paths, thePath)
			}
			targetStructTemp := strings.Split(line, "func (m *")
			if len(targetStructTemp) == 1 {
				continue
			}

			// get name of target struct
			structName := strings.Split(targetStructTemp[1], ") ")[0]

			// get import path and short name
			importPath := pathToImportPath[thePath]
			shortName := importPathShort[importPath]
			//fmt.Println(createHarness(shortName, structName))

			// create harness
			harness := createHarness(shortName, structName)
			harnesses = append(harnesses, harness)

			// add harness to list of function calls
			functionList = append(functionList, createFunctionCall(shortName, structName))
		}
	}
	return functionList, harnesses
}

func main() {
	functionList, harnesses := getGrepData()
	createMainFuzzer(functionList, harnesses)
}