make detection of reflect more robust

It now works with variables and composite type expressions too.

main.go

@@ -301,7 +301,8 @@ func transformCompile(args []string) ([]string, error) {
 		Defs: make(map[*ast.Ident]types.Object),
 		Uses: make(map[*ast.Ident]types.Object),
 	}
-	if _, err := origTypesConfig.Check(pkgPath, fset, files, info); err != nil {
+	pkg, err := origTypesConfig.Check(pkgPath, fset, files, info)
+	if err != nil {
 		return nil, fmt.Errorf("typecheck error: %v", err)
 	}
 
@@ -320,7 +321,7 @@ func transformCompile(args []string) ([]string, error) {
 	// log.Println(flags)
 	args = flags
 
-	blacklist := buildBlacklist(files, info)
+	blacklist := buildBlacklist(files, info, pkg)
 
 	// TODO: randomize the order and names of the files
 	for i, file := range files {
@@ -447,38 +448,49 @@ func hashWith(salt, value string) string {
 	return "z" + sum[:length]
 }
 
-func buildBlacklist(files []*ast.File, info *types.Info) (blacklist []types.Object) {
-	pre := func(node ast.Node) bool {
-		nodeExpr, ok := node.(*ast.CallExpr)
+// buildBlacklist collects all the objects in a package which are known to be
+// used with reflect.TypeOf or reflect.ValueOf. Since we obfuscate one package
+// at a time, we only detect those if the type definition and the reflect usage
+// are both in the same package.
+func buildBlacklist(files []*ast.File, info *types.Info, pkg *types.Package) (blacklist []types.Object) {
+	// Keep track of the current syntax tree level. If reflectCallLevel is
+	// non-negative, we are under a reflect call.
+	level := 0
+	reflectCallLevel := -1
+
+	visit := func(node ast.Node) bool {
+		if node == nil {
+			if level == reflectCallLevel {
+				reflectCallLevel = -1
+			}
+			level--
+			return true
+		}
+		if reflectCallLevel >= 0 && level >= reflectCallLevel {
+			expr, _ := node.(ast.Expr)
+			if obj := objOf(info.TypeOf(expr)); obj != nil && obj.Pkg() == pkg {
+				blacklist = append(blacklist, obj)
+			}
+		}
+		level++
+		call, ok := node.(*ast.CallExpr)
 		if !ok {
 			return true
 		}
-
-		exprFunc, ok := nodeExpr.Fun.(*ast.SelectorExpr)
+		sel, ok := call.Fun.(*ast.SelectorExpr)
 		if !ok {
 			return true
 		}
+		fnType := info.ObjectOf(sel.Sel)
 
-		exprFuncType := info.ObjectOf(exprFunc.Sel)
-
-		if exprFuncType.Pkg().Path() == "reflect" && (exprFuncType.Name() == "TypeOf" || exprFuncType.Name() == "ValueOf") {
-			for _, arg := range nodeExpr.Args {
-				if expr, ok := arg.(*ast.CallExpr); ok {
-					if f, ok := expr.Fun.(*ast.Ident); ok {
-						blacklistItem := info.ObjectOf(f)
-						blacklist = append(blacklist, blacklistItem)
-					}
-				}
-			}
+		if fnType.Pkg().Path() == "reflect" && (fnType.Name() == "TypeOf" || fnType.Name() == "ValueOf") {
+			reflectCallLevel = level
 		}
-
 		return true
 	}
-
 	for _, file := range files {
-		ast.Inspect(file, pre)
+		ast.Inspect(file, visit)
 	}
-
 	return blacklist
 }
 
@@ -536,8 +548,9 @@ func transformGo(file *ast.File, info *types.Info, blacklist []types.Object) *as
 			return true // could be a Go plugin API
 		}
 
+		// TODO: also do this for method receivers
 		for _, item := range blacklist {
-			if item.Pkg().Path() == obj.Pkg().Path() && item.Name() == obj.Name() {
+			if obj == item {
 				return true
 			}
 		}
@@ -611,6 +624,7 @@ func implementedOutsideGo(obj *types.Func) bool {
 // pointer type. This is useful to obtain "testing.T" from "*testing.T", or to
 // obtain the type declaration object from an embedded field.
 func objOf(t types.Type) types.Object {
+	fmt.Printf("t: %T\n", t)
 	switch t := t.(type) {
 	case *types.Named:
 		return t.Obj()

testdata/scripts/imports.txt

@@ -3,7 +3,6 @@ exec ./main
 cmp stdout main.stdout
 
 ! binsubstr main$exe 'ImportedVar' 'ImportedConst' 'ImportedFunc' 'ImportedType' 'main.go' 'imported.go'
-binsubstr main$exe 'ImportedAPI'
 
 [short] stop # checking that the build is reproducible is slow
 
@@ -24,7 +23,6 @@ package main
 
 import (
 	"fmt"
-	"reflect"
 	_ "unsafe"
 
 	"test/main/imported"
@@ -41,18 +39,13 @@ func main() {
 	imported.ImportedFunc('x')
 	fmt.Println(imported.ImportedType(3))
 
-	api := new(imported.ImportedAPI)
-	fmt.Println(reflect.TypeOf(api))
+	fmt.Printf("%T\n", imported.ReflectTypeOf(2))
+	fmt.Printf("%T\n", imported.ReflectTypeOfIndirect(4))
+	fmt.Printf("%#v\n", imported.ReflectValueOfVar)
+
 	linkedPrintln(nil)
 	fmt.Println(quote.Go())
 }
--- main.stdout --
-imported var value
-imported const value
-3
-*imported.ImportedAPI
-<nil>
-Don't communicate by sharing memory, share memory by communicating.
 -- imported/imported.go --
 package imported
 
@@ -66,8 +59,31 @@ func ImportedFunc(param rune) string {
 	return string(param)
 }
 
-type ImportedType int
+type ReflectTypeOf int
+
+var _ = reflect.TypeOf(ReflectTypeOf(0))
+
+type ReflectTypeOfIndirect int
+
+var _ = reflect.TypeOf(new([]*ReflectTypeOfIndirect))
 
-type ImportedAPI int
+type ReflectValueOf struct {
+	Foo int `bar:"baz"`
+}
+
+var ReflectValueOfVar = ReflectValueOf{Foo: 3}
+
+var _ = reflect.TypeOf(ReflectValueOfVar)
 
-var _ = reflect.TypeOf(ImportedAPI(0))
+// ImportedType comes after the calls to reflect, to ensure no false positives.
+type ImportedType int
+
+-- main.stdout --
+imported var value
+imported const value
+3
+imported.ReflectTypeOf
+imported.ReflectTypeOfIndirect
+imported.ReflectValueOf{Foo:3}
+<nil>
+Don't communicate by sharing memory, share memory by communicating.