Fix obfuscating built-in packages

Fix short name collision
Optimize encodeIntToName (remove underscore)

line_obfuscator.go

@@ -16,6 +16,50 @@ import (
 // Source: https://go.googlesource.com/go/+/refs/heads/master/src/cmd/compile/internal/syntax/parser_test.go#229
 const PosMin = 1
 
+const buildTagPrefix = "// +build"
+
+var nameSpecialDirectives = []string{
+	"//go:linkname",
+
+	"//go:cgo_export_static",
+	"//go:cgo_export_dynamic",
+	"//go:cgo_import_static",
+	"//go:cgo_import_dynamic",
+}
+
+var specialDirectives = append([]string{
+	"//go:cgo_ldflag",
+	"//go:cgo_dynamic_linker",
+	// Not necessarily, but it is desirable to prevent unexpected consequences in cases where "//go:generate" is linked to "node.Doc"
+	"//go:generate",
+}, nameSpecialDirectives...)
+
+func isOneOfDirective(text string, directives []string) bool {
+	for _, prefix := range directives {
+		if strings.HasPrefix(text, prefix) {
+			return true
+		}
+	}
+	return false
+}
+
+func getLocalName(text string) (string, bool) {
+	if !isOneOfDirective(text, nameSpecialDirectives) {
+		return "", false
+	}
+	parts := strings.SplitN(text, " ", 3)
+	if len(parts) < 2 {
+		return "", false
+	}
+
+	name := strings.TrimSpace(parts[1])
+	if len(name) == 0 {
+		return "", false
+	}
+
+	return name, false
+}
+
 func prependComment(group *ast.CommentGroup, comment *ast.Comment) *ast.CommentGroup {
 	if group == nil {
 		return &ast.CommentGroup{List: []*ast.Comment{comment}}
@@ -32,9 +76,8 @@ func clearCommentGroup(group *ast.CommentGroup) *ast.CommentGroup {
 	}
 
 	var comments []*ast.Comment
-
 	for _, comment := range group.List {
-		if strings.HasPrefix(comment.Text, "//go:") {
+		if strings.HasPrefix(comment.Text, "//go:") && !isOneOfDirective(comment.Text, specialDirectives) {
 			comments = append(comments, &ast.Comment{Text: comment.Text})
 		}
 	}
@@ -49,12 +92,16 @@ func clearNodeComments(node ast.Node) {
 	switch n := node.(type) {
 	case *ast.Field:
 		n.Doc = clearCommentGroup(n.Doc)
+		n.Comment = nil
 	case *ast.ImportSpec:
 		n.Doc = clearCommentGroup(n.Doc)
+		n.Comment = nil
 	case *ast.ValueSpec:
 		n.Doc = clearCommentGroup(n.Doc)
+		n.Comment = nil
 	case *ast.TypeSpec:
 		n.Doc = clearCommentGroup(n.Doc)
+		n.Comment = nil
 	case *ast.GenDecl:
 		n.Doc = clearCommentGroup(n.Doc)
 	case *ast.FuncDecl:
@@ -64,21 +111,39 @@ func clearNodeComments(node ast.Node) {
 	}
 }
 
-func findBuildTags(commentGroups []*ast.CommentGroup) (buildTags []string) {
-	for _, group := range commentGroups {
-		for _, comment := range group.List {
-			if !strings.Contains(comment.Text, "+build") {
+func processSpecialComments(commentGroups []*ast.CommentGroup) (extraComments, localnameBlacklist []string) {
+	var buildTags []string
+	var specialComments []string
+	for _, commentGroup := range commentGroups {
+		for _, comment := range commentGroup.List {
+			if strings.HasPrefix(comment.Text, buildTagPrefix) {
+				buildTags = append(buildTags, comment.Text)
 				continue
 			}
-			buildTags = append(buildTags, comment.Text)
+
+			if !isOneOfDirective(comment.Text, specialDirectives) {
+				continue
+			}
+
+			specialComments = append(specialComments, comment.Text)
+			localName, ok := getLocalName(comment.Text)
+			if ok {
+				localnameBlacklist = append(localnameBlacklist, localName)
+			}
 		}
 	}
-	return buildTags
+
+	extraComments = append(extraComments, buildTags...)
+	extraComments = append(extraComments, "")
+	extraComments = append(extraComments, specialComments...)
+	extraComments = append(extraComments, "")
+	return extraComments, localnameBlacklist
 }
 
-func transformLineInfo(file *ast.File) ([]string, *ast.File) {
+func transformLineInfo(file *ast.File) ([]string, []string, *ast.File) {
 	// Save build tags and add file name leak protection
-	extraComments := append(findBuildTags(file.Comments), "", "//line :1")
+	extraComments, localNameBlacklist := processSpecialComments(file.Comments)
+	extraComments = append(extraComments, "", "//line :1")
 	file.Comments = nil
 
 	newLines := mathrand.Perm(len(file.Decls))
@@ -88,6 +153,9 @@ func transformLineInfo(file *ast.File) ([]string, *ast.File) {
 		node := cursor.Node()
 		clearNodeComments(node)
 
+		if envGarbleTiny {
+			return true
+		}
 		funcDecl, ok := node.(*ast.FuncDecl)
 		if !ok {
 			return true
@@ -99,5 +167,5 @@ func transformLineInfo(file *ast.File) ([]string, *ast.File) {
 		return true
 	}
 
-	return extraComments, astutil.Apply(file, pre, nil).(*ast.File)
+	return extraComments, localNameBlacklist, astutil.Apply(file, pre, nil).(*ast.File)
 }

main.go

@@ -494,15 +494,6 @@ func transformCompile(args []string) ([]string, error) {
 	flags = append(flags, "-dwarf=false")
 
 	pkgPath := flagValue(flags, "-p")
-	if (pkgPath == "runtime" && envGarbleTiny) || pkgPath == "runtime/internal/sys" {
-		// Even though these packages aren't private, we will still process
-		// them later to remove build information and strip code from the
-		// runtime. However, we only want flags to work on private packages.
-		envGarbleLiterals = false
-		envGarbleDebugDir = ""
-	} else if !isPrivate(pkgPath) {
-		return append(flags, paths...), nil
-	}
 	for i, path := range paths {
 		if filepath.Base(path) == "_gomod_.go" {
 			// never include module info
@@ -555,16 +546,12 @@ func transformCompile(args []string) ([]string, error) {
 		return nil, fmt.Errorf("typecheck error: %v", err)
 	}
 
-	blacklist := buildBlacklist(files, info, pkg)
+	blacklist, existsNames := buildBlacklist(files, info, pkg)
 
 	// unsafe.Pointer is a special type that doesn't exist as a plain Go
 	// type definition, so we can't change its name.
 	blacklist[types.Unsafe.Scope().Lookup("Pointer")] = struct{}{}
 
-	if envGarbleLiterals {
-		files = literals.Obfuscate(files, info, fset, blacklist)
-	}
-
 	tempDir, err := ioutil.TempDir("", "garble-build")
 	if err != nil {
 		return nil, err
@@ -589,11 +576,31 @@ func transformCompile(args []string) ([]string, error) {
 	}
 
 	privateNameMap := make(map[string]string)
+	packageCounter := 0
+
+	filesExtraComments := make([][]string, len(files))
+
+	for i, file := range files {
+		println(paths[i])
+		extraComments, localNameBlacklist, file := transformLineInfo(file)
+		for _, name := range localNameBlacklist {
+			localFunc := pkg.Scope().Lookup(name)
+			if localFunc == nil {
+				continue
+			}
+			blacklist[localFunc] = struct{}{}
+		}
+		filesExtraComments[i] = extraComments
+		files[i] = file
+	}
+
+	if envGarbleLiterals && isPrivate(pkgPath) {
+		files = literals.Obfuscate(files, info, fset, blacklist)
+	}
 
 	// TODO: randomize the order and names of the files
 	newPaths := make([]string, 0, len(files))
 	for i, file := range files {
-		var extraComments []string
 		origName := filepath.Base(filepath.Clean(paths[i]))
 		name := origName
 		switch {
@@ -614,23 +621,16 @@ func transformCompile(args []string) ([]string, error) {
 			spec := file.Decls[0].(*ast.GenDecl).Specs[0].(*ast.ValueSpec)
 			lit := spec.Values[0].(*ast.BasicLit)
 			lit.Value = "`unknown`"
-		case strings.HasPrefix(origName, "_cgo_"):
-			// Cgo generated code requires a prefix. Also, don't
-			// garble it, since it's just generated code and it gets
-			// messy.
-			name = "_cgo_" + name
-		default:
-			if !envGarbleTiny {
-				extraComments, file = transformLineInfo(file)
-			}
-			file = transformGo(file, info, blacklist, privateNameMap, pkgPath)
-
-			// Uncomment for some quick debugging. Do not delete.
-			// fmt.Fprintf(os.Stderr, "\n-- %s/%s --\n", pkgPath, origName)
-			// if err := printConfig.Fprint(os.Stderr, fset, file); err != nil {
-			// 	return nil, err
-			// }
 		}
+
+		file = transformGo(file, info, blacklist, privateNameMap, existsNames, pkgPath, &packageCounter)
+
+		// Uncomment for some quick debugging. Do not delete.
+		// fmt.Fprintf(os.Stderr, "\n-- %s/%s --\n", pkgPath, origName)
+		// if err := printConfig.Fprint(os.Stderr, fset, file); err != nil {
+		// 	return nil, err
+		// }
+
 		tempFile, err := os.Create(filepath.Join(tempDir, name))
 		if err != nil {
 			return nil, err
@@ -649,6 +649,7 @@ func transformCompile(args []string) ([]string, error) {
 			printWriter = io.MultiWriter(tempFile, debugFile)
 		}
 
+		extraComments := filesExtraComments[i]
 		if len(extraComments) > 0 {
 			for _, comment := range extraComments {
 				if _, err = printWriter.Write([]byte(comment + "\n")); err != nil {
@@ -830,12 +831,16 @@ var privateNameCharset = buildNameCharset()
 
 func encodeIntToName(i int) string {
 	builder := strings.Builder{}
-	builder.WriteByte('_')
 	for i > 0 {
 		charIdx := i % len(privateNameCharset)
 		i -= charIdx + 1
-		builder.WriteRune(privateNameCharset[charIdx])
+		c := privateNameCharset[charIdx]
+		if unicode.IsDigit(c) && builder.Len() == 0 {
+			builder.WriteByte('_')
+		}
+		builder.WriteRune(c)
 	}
+
 	return builder.String()
 }
 
@@ -845,8 +850,9 @@ func encodeIntToName(i int) string {
 // are both in the same package.
 //
 // The blacklist mainly contains named types and their field declarations.
-func buildBlacklist(files []*ast.File, info *types.Info, pkg *types.Package) map[types.Object]struct{} {
+func buildBlacklist(files []*ast.File, info *types.Info, pkg *types.Package) (map[types.Object]struct{}, map[string]struct{}) {
 	blacklist := make(map[types.Object]struct{})
+	existsNames := make(map[string]struct{})
 
 	reflectBlacklist := func(node ast.Node) bool {
 		expr, _ := node.(ast.Expr)
@@ -872,6 +878,10 @@ func buildBlacklist(files []*ast.File, info *types.Info, pkg *types.Package) map
 	}
 
 	visit := func(node ast.Node) bool {
+		if ident, ok := node.(*ast.Ident); ok {
+			existsNames[ident.Name] = struct{}{}
+		}
+
 		if envGarbleLiterals {
 			literals.ConstBlacklist(node, info, blacklist)
 		}
@@ -900,11 +910,11 @@ func buildBlacklist(files []*ast.File, info *types.Info, pkg *types.Package) map
 	for _, file := range files {
 		ast.Inspect(file, visit)
 	}
-	return blacklist
+	return blacklist, existsNames
 }
 
 // transformGo garbles the provided Go syntax node.
-func transformGo(file *ast.File, info *types.Info, blacklist map[types.Object]struct{}, privateNameMap map[string]string, pkgPath string) *ast.File {
+func transformGo(file *ast.File, info *types.Info, blacklist map[types.Object]struct{}, privateNameMap map[string]string, existsNames map[string]struct{}, pkgPath string, packageCounter *int) *ast.File {
 	// Shuffle top level declarations
 	mathrand.Shuffle(len(file.Decls), func(i, j int) {
 		decl1 := file.Decls[i]
@@ -1026,7 +1036,14 @@ func transformGo(file *ast.File, info *types.Info, blacklist map[types.Object]st
 			return true
 		}
 
-		name := encodeIntToName(len(privateNameMap) + 1)
+		var name string
+		for {
+			*packageCounter++
+			name = encodeIntToName(*packageCounter)
+			if _, ok := existsNames[name]; !ok {
+				break
+			}
+		}
 
 		// orig := node.Name
 		privateNameMap[fullName] = name