From f1bf6f91eeea8b6058cf0dc9659e2b880c148c69 Mon Sep 17 00:00:00 2001
From: lu4p <lu4p@pm.me>
Date: Mon, 6 Jul 2020 09:49:21 +0200
Subject: [PATCH] skip literals used in constant expressions

Fixes #39.
---
 .gitignore                   |  1 +
 main.go                      | 10 +++++--
 strings.go                   | 57 +++++++++++++++++++++++++++++++++++-
 testdata/scripts/strings.txt | 38 ++++++++++++++++++++++--
 4 files changed, 100 insertions(+), 6 deletions(-)

diff --git a/.gitignore b/.gitignore
index e14c665..3a3e340 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1 +1,2 @@
 /garble
+/test
\ No newline at end of file
diff --git a/main.go b/main.go
index 10b4e9f..c5ae0eb 100644
--- a/main.go
+++ b/main.go
@@ -386,8 +386,10 @@ func transformCompile(args []string) ([]string, error) {
 		return nil, fmt.Errorf("typecheck error: %v", err)
 	}
 
+	blacklist := buildBlacklist(files, info, pkg)
+
 	if envGarbleLiterals {
-		files = obfuscateLiterals(files, info)
+		files = obfuscateLiterals(files, info, blacklist)
 		// ast changed so we need to typecheck again
 		pkg, err = origTypesConfig.Check(pkgPath, fset, files, info)
 		if err != nil {
@@ -410,8 +412,6 @@ func transformCompile(args []string) ([]string, error) {
 	// log.Println(flags)
 	args = flags
 
-	blacklist := buildBlacklist(files, info, pkg)
-
 	pkgDebugDir := ""
 	if envGarbleDebugDir != "" {
 		osPkgPath := filepath.FromSlash(pkgPath)
@@ -665,6 +665,10 @@ func buildBlacklist(files []*ast.File, info *types.Info, pkg *types.Package) map
 		}
 	}
 	visit := func(node ast.Node) bool {
+		if envGarbleLiterals {
+			constBlacklist(node, info, blacklist)
+		}
+
 		if node == nil {
 			if level == reflectCallLevel {
 				reflectCallLevel = -1
diff --git a/strings.go b/strings.go
index 9c33fc4..e95a02e 100644
--- a/strings.go
+++ b/strings.go
@@ -43,7 +43,7 @@ func containsTypeDefStr(expr ast.Expr, info *types.Info) bool {
 	return false
 }
 
-func obfuscateLiterals(files []*ast.File, info *types.Info) []*ast.File {
+func obfuscateLiterals(files []*ast.File, info *types.Info, blacklist map[types.Object]struct{}) []*ast.File {
 	pre := func(cursor *astutil.Cursor) bool {
 		switch x := cursor.Node().(type) {
 		case *ast.ValueSpec:
@@ -91,6 +91,15 @@ func obfuscateLiterals(files []*ast.File, info *types.Info) []*ast.File {
 					return false
 				}
 
+				for _, name := range spec.Names {
+					obj := info.ObjectOf(name)
+
+					// The object itself is blacklisted, e.g. a value that needs to be constant
+					if _, ok := blacklist[obj]; ok {
+						return false
+					}
+				}
+
 				for _, val := range spec.Values {
 					if v, ok := val.(*ast.BasicLit); !ok || v.Kind != token.STRING {
 						return false // skip the block if it contains non basic literals
@@ -306,3 +315,49 @@ func keyStmt(key []byte) *ast.GenDecl {
 		}},
 	}
 }
+
+func constBlacklist(node ast.Node, info *types.Info, blacklist map[types.Object]struct{}) {
+
+	blacklistObjects := func(node ast.Node) bool {
+		ident, ok := node.(*ast.Ident)
+		if !ok {
+			return true
+		}
+
+		obj := info.ObjectOf(ident)
+		blacklist[obj] = struct{}{}
+
+		return true
+	}
+
+	switch x := node.(type) {
+	// in a slice or array composite literal all explicit keys must be constant representable
+	case *ast.CompositeLit:
+		if _, ok := x.Type.(*ast.ArrayType); !ok {
+			break
+		}
+		for _, elt := range x.Elts {
+			if kv, ok := elt.(*ast.KeyValueExpr); ok {
+				ast.Inspect(kv.Key, blacklistObjects)
+			}
+		}
+	// in an array type the length must be a constant representable
+	case *ast.ArrayType:
+		if x.Len != nil {
+			ast.Inspect(x.Len, blacklistObjects)
+		}
+
+	// in a const declaration all values must be constant representable
+	case *ast.GenDecl:
+		if x.Tok != token.CONST {
+			break
+		}
+		for _, spec := range x.Specs {
+			spec := spec.(*ast.ValueSpec)
+
+			for _, val := range spec.Values {
+				ast.Inspect(val, blacklistObjects)
+			}
+		}
+	}
+}
diff --git a/testdata/scripts/strings.txt b/testdata/scripts/strings.txt
index db2f051..a0de26d 100644
--- a/testdata/scripts/strings.txt
+++ b/testdata/scripts/strings.txt
@@ -78,6 +78,7 @@ func main() {
 	println(skip1, skip2)
 	println(i, foo, bar)
 	typedTest()
+	constantTest()
 }
 
 type stringType string
@@ -98,7 +99,7 @@ func typedTest() {
 	println(skipTypedConst, skipTypedVar, skipTypedVarAssign)
 
 	y := stringTypeStruct{
-		str:     "stringTypeField String",     // obfuscate
+		str:     "stringTypeField String",  // obfuscate
 		strType: "stringTypeField strType", // skip
 	}
 	println(y.str, y.strType)
@@ -120,6 +121,37 @@ func typedTest() {
 	println(stringTypeFunc("stringType func param")) // skip
 }
 
+// constantTest tests that string constants which need to be constant are skipped
+func constantTest() {
+	const a = "foo" // skip
+	const length = len(a)
+
+	const b = "bar" // skip
+	type T [len(b)]byte
+
+	const c = "foo" // skip
+	var _ [len(c)]byte
+
+	const d = "foo" // skip
+	var arr = [5]string{len(d): "foo"}
+	for _, elm := range arr {
+		if elm != "" {
+			println(elm)
+		}
+	}
+
+	const e = "foo" // skip
+	var slice = []string{len(e): "foo"}
+	for _, elm := range slice {
+		if elm != "" {
+			println(elm)
+		}
+	}
+
+	const f = "foo" // skip
+	const i = length + len(f)
+}
+
 func stringTypeFunc(s stringType) stringType {
 	println(s)
 	return "stringType return" // skip
@@ -141,4 +173,6 @@ skip typed const skip typed var skip typed var assign
 stringTypeField String stringTypeField strType
 stringType lambda func return
 stringType func param
-stringType return
\ No newline at end of file
+stringType return
+foo
+foo