diff --git a/line_obfuscator.go b/line_obfuscator.go index 1c6cea0..6f582e6 100644 --- a/line_obfuscator.go +++ b/line_obfuscator.go @@ -67,7 +67,22 @@ func clearNodeComments(node ast.Node) { } } -func transformLineInfo(fileIndex int, file *ast.File) *ast.File { +func findBuildTags(commentGroups []*ast.CommentGroup) (buildTags []string) { + for _, group := range commentGroups { + for _, comment := range group.List { + if !strings.Contains(comment.Text, "+build") { + continue + } + buildTags = append(buildTags, comment.Text) + } + } + return buildTags +} + +func transformLineInfo(fileIndex int, file *ast.File) ([]string, *ast.File) { + // Save build tags and add file name leak protection + extraComments := append(findBuildTags(file.Comments), "", "//line :1") + file.Comments = nil pre := func(cursor *astutil.Cursor) bool { node := cursor.Node() @@ -89,5 +104,5 @@ func transformLineInfo(fileIndex int, file *ast.File) *ast.File { return true } - return astutil.Apply(file, pre, nil).(*ast.File) + return extraComments, astutil.Apply(file, pre, nil).(*ast.File) } diff --git a/main.go b/main.go index 90af7d3..a4c4042 100644 --- a/main.go +++ b/main.go @@ -483,7 +483,7 @@ func transformCompile(args []string) ([]string, error) { // TODO: randomize the order and names of the files for i, file := range files { - protectFileName := false + var extraComments []string origName := filepath.Base(filepath.Clean(paths[i])) name := origName switch { @@ -510,8 +510,7 @@ func transformCompile(args []string) ([]string, error) { // messy. name = "_cgo_" + name default: - protectFileName = true - file = transformLineInfo(i, file) + extraComments, file = transformLineInfo(i, file) file = transformGo(file, info, blacklist) // Uncomment for some quick debugging. Do not delete. @@ -538,11 +537,12 @@ func transformCompile(args []string) ([]string, error) { printWriter = io.MultiWriter(tempFile, debugFile) } - if protectFileName { - // File name leak protection - _, err = printWriter.Write([]byte("//line :1\n")) - if err != nil { - return nil, err + if len(extraComments) > 0 { + for _, comment := range extraComments { + _, err = printWriter.Write([]byte(comment + "\n")) + if err != nil { + return nil, err + } } } if err := printConfig.Fprint(printWriter, fset, file); err != nil {