file_walker.go

Documentation: github.com/AdamKorcz/go-118-fuzz-build

     1  package main
     2  
     3  import (
     4  	"bytes"
     5  	//"fmt"
     6  	"go/ast"
     7  	"go/parser"
     8  	"go/printer"
     9  	"go/token"
    10  	"os"
    11  
    12  	"golang.org/x/tools/go/ast/astutil"
    13  	"golang.org/x/tools/go/packages"
    14  )
    15  
    16  // rewriteTestingImports rewrites imports for:
    17  // - all package files
    18  // - the fuzzer
    19  // - dependencies
    20  //
    21  // it rewrites "testing" => "github.com/AdamKorcz/go-118-fuzz-build/testing"
    22  func rewriteTestingImports(pkgs []*packages.Package, fuzzName string) (string, []byte, error) {
    23  	var fuzzFilepath string
    24  	var originalFuzzContents []byte
    25  	originalFuzzContents = []byte("NONE")
    26  
    27  	// First find file with fuzz harness
    28  	for _, pkg := range pkgs {
    29  		for _, file := range pkg.GoFiles {
    30  			//fmt.Println(file)
    31  			err := rewriteTestingImport(file)
    32  			if err != nil {
    33  				panic(err)
    34  			}
    35  		}
    36  	}
    37  
    38  	// rewrite testing in imported packages
    39  	packages.Visit(pkgs, rewriteImportTesting, nil)
    40  
    41  	for _, pkg := range pkgs {
    42  		for _, file := range pkg.GoFiles {
    43  			fuzzFile, b, err := rewriteFuzzer(file, fuzzName)
    44  			if err != nil {
    45  				panic(err)
    46  			}
    47  			if fuzzFile != "" {
    48  				fuzzFilepath = fuzzFile
    49  				originalFuzzContents = b
    50  			}
    51  		}
    52  	}
    53  	return fuzzFilepath, originalFuzzContents, nil
    54  }
    55  
    56  func rewriteFuzzer(path, fuzzerName string) (originalPath string, originalFile []byte, err error) {
    57  	var fileHasOurHarness bool // to determine whether we should rewrite filename
    58  	fileHasOurHarness = false
    59  
    60  	var originalFuzzContents []byte
    61  	originalFuzzContents = []byte("NONE")
    62  
    63  	fset := token.NewFileSet()
    64  	f, err := parser.ParseFile(fset, path, nil, 0)
    65  	if err != nil {
    66  		return "", originalFuzzContents, err
    67  	}
    68  	for _, decl := range f.Decls {
    69  		if _, ok := decl.(*ast.FuncDecl); ok {
    70  			if decl.(*ast.FuncDecl).Name.Name == fuzzerName {
    71  				fileHasOurHarness = true
    72  
    73  			}
    74  		}
    75  	}
    76  
    77  	if fileHasOurHarness {
    78  		originalFuzzContents, err = os.ReadFile(path)
    79  		if err != nil {
    80  			panic(err)
    81  		}
    82  
    83  		// Replace import path
    84  		astutil.DeleteImport(fset, f, "testing")
    85  		astutil.AddImport(fset, f, "github.com/AdamKorcz/go-118-fuzz-build/testing")
    86  	}
    87  
    88  	// Rewrite filename
    89  	if fileHasOurHarness {
    90  		var buf bytes.Buffer
    91  		printer.Fprint(&buf, fset, f)
    92  
    93  		newFile, err := os.Create(path + "_fuzz.go")
    94  		if err != nil {
    95  			panic(err)
    96  		}
    97  		defer newFile.Close()
    98  		newFile.Write(buf.Bytes())
    99  		return path, originalFuzzContents, nil
   100  	}
   101  	return "", originalFuzzContents, nil
   102  }
   103  
   104  // Rewrites testing import of a single path
   105  func rewriteTestingImport(path string) error {
   106  	//fmt.Println("Rewriting ", path)
   107  	fsetCheck := token.NewFileSet()
   108  	fCheck, err := parser.ParseFile(fsetCheck, path, nil, parser.ImportsOnly)
   109  	if err != nil {
   110  		return err
   111  	}
   112  
   113  	// First check if the import already exists
   114  	// Return if it does.
   115  	for _, imp := range fCheck.Imports {
   116  		if imp.Path.Value == "github.com/AdamKorcz/go-118-fuzz-build/testing" {
   117  			return nil
   118  		}
   119  	}
   120  
   121  	// Replace import path
   122  	for _, imp := range fCheck.Imports {
   123  		if imp.Path.Value == "testing" {
   124  			imp.Path.Value = "github.com/AdamKorcz/go-118-fuzz-build/testing"
   125  		}
   126  	}
   127  	return nil
   128  }
   129  
   130  // Rewrites testing import of a package
   131  func rewriteImportTesting(pkg *packages.Package) bool {
   132  	for _, file := range pkg.GoFiles {
   133  		err := rewriteTestingImport(file)
   134  		if err != nil {
   135  			panic(err)
   136  		}
   137  	}
   138  	return true
   139  }
   140  
   141  // Checks whether a fuzz test exists in a given file
   142  func rewriteFuzzerImports(path, fuzzName string) error {
   143  	fset := token.NewFileSet()
   144  	f, err := parser.ParseFile(fset, path, nil, 0)
   145  	if err != nil {
   146  		return err
   147  	}
   148  	for _, decl := range f.Decls {
   149  		if _, ok := decl.(*ast.FuncDecl); ok {
   150  			if decl.(*ast.FuncDecl).Name.Name == fuzzName {
   151  				// First rewrite testing.F
   152  
   153  			}
   154  		}
   155  	}
   156  	return nil
   157  }
   158
View as plain text