From 7e312c557fd10df8b19a534f1feb4168fbc328f5 Mon Sep 17 00:00:00 2001
From: Pagran <67878280+pagran@users.noreply.github.com>
Date: Tue, 29 Sep 2020 22:45:09 +0300
Subject: [PATCH] Fix short name collision Reduce short name length

---
 main.go                     | 38 ++++++++++++++++++++++++++++++++-----
 testdata/scripts/syntax.txt | 20 +++++++++++++++++++
 2 files changed, 53 insertions(+), 5 deletions(-)

diff --git a/main.go b/main.go
index eb27339..5d7c698 100644
--- a/main.go
+++ b/main.go
@@ -589,6 +589,8 @@ func transformCompile(args []string) ([]string, error) {
 	}
 
 	privateNameMap := make(map[string]string)
+	existsNames := collectNames(files)
+	packageCounter := 0
 
 	// TODO: randomize the order and names of the files
 	newPaths := make([]string, 0, len(files))
@@ -623,7 +625,7 @@ func transformCompile(args []string) ([]string, error) {
 			if !envGarbleTiny {
 				extraComments, file = transformLineInfo(file)
 			}
-			file = transformGo(file, info, blacklist, privateNameMap, pkgPath)
+			file = transformGo(file, info, blacklist, privateNameMap, pkgPath, existsNames, &packageCounter)
 
 			// Uncomment for some quick debugging. Do not delete.
 			// fmt.Fprintf(os.Stderr, "\n-- %s/%s --\n", pkgPath, origName)
@@ -830,11 +832,14 @@ 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 builder.Len() == 0 && !unicode.IsLetter(c) {
+			builder.WriteByte('_')
+		}
+		builder.WriteRune(c)
 	}
 	return builder.String()
 }
@@ -903,8 +908,24 @@ func buildBlacklist(files []*ast.File, info *types.Info, pkg *types.Package) map
 	return blacklist
 }
 
+// collectNames collects all names, including the names of local variables,
+// functions, global fields, etc.
+func collectNames(files []*ast.File) map[string]struct{} {
+	blacklist := make(map[string]struct{})
+	visit := func(node ast.Node) bool {
+		if ident, ok := node.(*ast.Ident); ok {
+			blacklist[ident.Name] = struct{}{}
+		}
+		return true
+	}
+	for _, file := range files {
+		ast.Inspect(file, visit)
+	}
+	return blacklist
+}
+
 // 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, pkgPath string, existsNames map[string]struct{}, packageCounter *int) *ast.File {
 	// Shuffle top level declarations
 	mathrand.Shuffle(len(file.Decls), func(i, j int) {
 		decl1 := file.Decls[i]
@@ -1026,7 +1047,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
diff --git a/testdata/scripts/syntax.txt b/testdata/scripts/syntax.txt
index 20c5c48..8f8e1e3 100644
--- a/testdata/scripts/syntax.txt
+++ b/testdata/scripts/syntax.txt
@@ -26,6 +26,7 @@ import (
 	"go/ast"
 
 	"rsc.io/quote"
+	"test/main/sub"
 )
 
 // This comment contains valuable information. Ensure it's not in the final binary.
@@ -71,6 +72,7 @@ func main() {
 	println(string(enc))
 	scopesTest()
 	println(quote.Go())
+	sub.Test()
 }
 
 -- scopes.go --
@@ -103,6 +105,24 @@ func scopesTest() {
 
 func input(localNameParam string) (localNameReturn string) { return localNameParam }
 
+-- sub\names.go --
+package sub
+
+// Do not move these fields! Moving can break a test
+var someGlobalVar0 = "0"
+var someGlobalVar1 = "1"
+var someGlobalVar2 = "2"
+
+func Test() {
+    var A, B, C, D, E string
+    noop(A, B, C, D, E)
+    if someGlobalVar0 != "0" || someGlobalVar1 != "1" || someGlobalVar2 != "2"{
+        panic("name collision detected")
+    }
+}
+
+func noop(...interface{}) {}
+
 -- main.stderr --
 nil case
 {"Foo":3}