More correctly comments transformation

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, true
+}
+
 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})
 		}
 	}
@@ -68,21 +111,44 @@ 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
+			}
+
+			if !isOneOfDirective(comment.Text, specialDirectives) {
 				continue
 			}
-			buildTags = append(buildTags, comment.Text)
+
+			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, cgoFile bool) ([]string, []string, *ast.File) {
+	prefix := ""
+	if cgoFile {
+		prefix = "_cgo_"
+	}
+
 	// 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 "+prefix+":1")
 	file.Comments = nil
 
 	newLines := mathrand.Perm(len(file.Decls))
@@ -92,16 +158,19 @@ 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
 		}
 
-		comment := &ast.Comment{Text: fmt.Sprintf("//line %c.go:%d", nameCharset[mathrand.Intn(len(nameCharset))], PosMin+newLines[funcCounter])}
+		comment := &ast.Comment{Text: fmt.Sprintf("//line %s%c.go:%d", prefix, nameCharset[mathrand.Intn(len(nameCharset))], PosMin+newLines[funcCounter])}
 		funcDecl.Doc = prependComment(funcDecl.Doc, comment)
 		funcCounter++
 		return true
 	}
 
-	return extraComments, astutil.Apply(file, pre, nil).(*ast.File)
+	return extraComments, localNameBlacklist, astutil.Apply(file, pre, nil).(*ast.File)
 }

main.go

@@ -591,11 +591,27 @@ func transformCompile(args []string) ([]string, error) {
 	privateNameMap := make(map[string]string)
 	existingNames := collectNames(files)
 	packageCounter := 0
+	filesExtraComments := make([][]string, len(files))
+
+	for i, file := range files {
+		name := filepath.Base(filepath.Clean(paths[i]))
+		cgoFile := strings.HasPrefix(name, "_cgo_")
+		extraComments, localNameBlacklist, file := transformLineInfo(file, cgoFile)
+		for _, name := range localNameBlacklist {
+			obj := pkg.Scope().Lookup(name)
+			if obj == nil {
+				continue
+			}
+
+			blacklist[obj] = struct{}{}
+		}
+		filesExtraComments[i] = extraComments
+		files[i] = file
+	}
 
 	// 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 {
@@ -622,9 +638,6 @@ func transformCompile(args []string) ([]string, error) {
 			// messy.
 			name = "_cgo_" + name
 		default:
-			if !envGarbleTiny {
-				extraComments, file = transformLineInfo(file)
-			}
 			file = transformGo(file, info, blacklist, privateNameMap, pkgPath, existingNames, &packageCounter)
 
 			// Uncomment for some quick debugging. Do not delete.
@@ -651,6 +664,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 {

testdata/scripts/syntax.txt

@@ -1,14 +1,15 @@
 env GOPRIVATE='test/main,rsc.io/*'
 
-garble build
+garble build -tags directives
 exec ./main$exe
 cmp stderr main.stderr
 
-! binsubstr main$exe 'localName' 'globalConst' 'globalVar' 'globalType' 'valuable information' 'rsc.io'
+! binsubstr main$exe 'localName' 'globalConst' 'globalVar' 'globalType' 'valuable information' 'rsc.io' 'directiveLocalFunc'
+binsubstr main$exe 'newDirectiveFuncName'
 
 [short] stop # no need to verify this with -short
 
-go build
+go build -tags directives
 exec ./main$exe
 cmp stderr main.stderr
 
@@ -73,6 +74,7 @@ func main() {
 	scopesTest()
 	println(quote.Go())
 	sub.Test()
+	sub.TestDirectives()
 }
 
 -- scopes.go --
@@ -122,6 +124,22 @@ func Test() {
 
 func noop(...interface{}) {}
 
+-- sub/directives.go --
+// +build directives
+
+package sub
+
+import _ "unsafe"
+
+//go:linkname directiveLocalFunc newDirectiveFuncName
+
+//go:noinline
+func directiveLocalFunc() { }
+
+//go:noinline
+func TestDirectives() {
+    directiveLocalFunc()
+}
 -- main.stderr --
 nil case
 {"Foo":3}