/
transformer.go
136 lines (117 loc) · 3.83 KB
/
transformer.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
package hook
import (
"go/ast"
"go/token"
"go/types"
"path/filepath"
"strconv"
"golang.org/x/tools/go/ast/astutil"
"github.com/CodeIntelligenceTesting/gofuzz/internal/pkg/log"
)
type Transformer struct {
file *ast.File
fileSet *token.FileSet
filePath string
typeInfo *types.Info
transformedImports map[string]string
hookIdFn IdFunction
numAddedHooks int
}
func NewTransformer(file *ast.File, fileSet *token.FileSet, filePath string, typeInfo *types.Info, hookIdFn IdFunction) *Transformer {
return &Transformer{
file: file,
fileSet: fileSet,
filePath: filePath,
typeInfo: typeInfo,
transformedImports: make(map[string]string),
hookIdFn: hookIdFn,
numAddedHooks: 0,
}
}
func (t *Transformer) TransformFile() int {
astutil.Apply(t.file, nil, func(cursor *astutil.Cursor) bool {
callExpr, ok := cursor.Node().(*ast.CallExpr)
if !ok {
return true
}
selectorExpr, ok := callExpr.Fun.(*ast.SelectorExpr)
if !ok {
return true
}
receiver := selectorExpr.X
if pkgName, pkgPath, ok := importedPackage(receiver, t.file); ok {
if h := MatchingFunctionHook(selectorExpr.Sel.Name, pkgPath); h != nil {
log.Debugf("Hooked function %s.%s", pkgPath, selectorExpr.Sel.Name)
selectorExpr.Sel.Name = h.HookName
selectorExpr.X = &ast.Ident{
Name: sanitizersPackageName,
}
callExpr.Args = append([]ast.Expr{t.hookIdFn(callExpr, t.filePath, t.fileSet)}, callExpr.Args...)
astutil.AddNamedImport(t.fileSet, t.file, sanitizersPackageName, sanitizersPackagePath)
t.transformedImports[pkgPath] = pkgName
t.numAddedHooks += 1
}
} else if receiverType := t.typeInfo.Types[receiver]; receiverType.IsValue() {
if h := MatchingMethodHook(selectorExpr.Sel.Name, receiverType.Type.String()); h != nil {
log.Debugf("Hooked method %s.%s", receiverType.Type.String(), selectorExpr.Sel.Name)
selectorExpr.Sel.Name = h.HookName
selectorExpr.X = &ast.Ident{
Name: sanitizersPackageName,
}
callExpr.Args = append([]ast.Expr{t.hookIdFn(callExpr, t.filePath, t.fileSet), receiver}, callExpr.Args...)
astutil.AddNamedImport(t.fileSet, t.file, sanitizersPackageName, sanitizersPackagePath)
t.numAddedHooks += 1
}
}
return true
})
if t.numAddedHooks > 0 {
t.removeUnusedImport()
}
return t.numAddedHooks
}
func (t *Transformer) removeUnusedImport() {
for path, name := range t.transformedImports {
if !astutil.UsesImport(t.file, path) {
astutil.DeleteNamedImport(t.fileSet, t.file, name, path)
}
}
}
func importedPackage(expr ast.Expr, file *ast.File) (name string, path string, found bool) {
// the package part of function call expressions (e.g., os.Open()) is an identifier
ident, ok := expr.(*ast.Ident)
if !ok {
return "", "", false
}
// package identifiers should top-level unresolved identifiers
if ident.Obj != nil {
return "", "", false
}
for _, importSpec := range file.Imports {
// remove the quotations from the AST string value of the import path
importPath, err := strconv.Unquote(importSpec.Path.Value)
if err != nil {
return "", "", false
}
var importName string
if importSpec.Name != nil {
// handle named imports: import newMyPkg "github.com/myMod/myPkg".
// In this case importSpec.Name = "newMyPkg" and newMyPkg will be the identifier
// used to call functions exported by the package: newMyPkg.Foo()
importName = importSpec.Name.Name
name = importName
} else {
// handle unnamed imports: import "github.com/myMod/myPkg"
// In this case myPkg will be the identifier used to call functions exported
// by the package: myPkg.Foo()
importName = filepath.Base(importPath)
name = ""
}
if ident.Name == importName {
found = true
path = importPath
return
}
}
return "", "", false
}