From 56a1fd0257864c54b80ff760c46f3463b6fd2f9c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Daniel=20Mart=C3=AD?= <mvdan@mvdan.cc>
Date: Mon, 25 May 2020 21:41:23 +0100
Subject: [PATCH] support -ldflags=-X=pkg.name=str with garbled names

Because the linker has access to all the build IDs, just like the
compiler, we can support this transparently. Add a test too.

Fixes #21.
---
 main.go                      | 64 ++++++++++++++++++++++++++++++++----
 testdata/scripts/ldflags.txt | 36 ++++++++++++++++++++
 2 files changed, 94 insertions(+), 6 deletions(-)
 create mode 100644 testdata/scripts/ldflags.txt

diff --git a/main.go b/main.go
index 53ae290..1d2ce03 100644
--- a/main.go
+++ b/main.go
@@ -112,8 +112,10 @@ func garbledImport(path string) (*types.Package, error) {
 }
 
 type packageInfo struct {
-	buildID string
-	imports map[string]importedPkg
+	buildID string // from -buildid
+
+	imports     map[string]importedPkg // from -importcfg
+	firstImport string                 // first from -importcfg; the main package when linking
 }
 
 type importedPkg struct {
@@ -310,6 +312,7 @@ func transformCompile(args []string) ([]string, error) {
 	deferred = append(deferred, func() error {
 		return os.RemoveAll(tempDir)
 	})
+
 	// Add our temporary dir to the beginning of -trimpath, so that we don't
 	// leak temporary dirs. Needs to be at the beginning, since there may be
 	// shorter prefixes later in the list, such as $PWD if TMPDIR=$PWD/tmp.
@@ -354,6 +357,8 @@ func transformCompile(args []string) ([]string, error) {
 func isPrivate(pkgPath string) bool {
 	if pkgPath == "main" {
 		// TODO: why don't we see the full package path for main packages?
+		// Hint: it seems like the real import path is at the top of
+		// -importcfg.
 		return true
 	}
 	return GlobsMatchPath(envGoPrivate, pkgPath)
@@ -396,6 +401,9 @@ func readBuildIDs(flags []string) error {
 		if err != nil {
 			return err
 		}
+		if len(buildInfo.imports) == 0 {
+			buildInfo.firstImport = importPath
+		}
 		buildInfo.imports[importPath] = importedPkg{
 			packagefile: objectPath,
 			buildID:     fileID,
@@ -587,6 +595,39 @@ func transformLink(args []string) ([]string, error) {
 		// Nothing to transform; probably just ["-V=full"].
 		return args, nil
 	}
+
+	// Make sure -X works with garbled identifiers. To cover both garbled
+	// and non-garbled names, duplicate each flag with a garbled version.
+	if err := readBuildIDs(flags); err != nil {
+		return nil, err
+	}
+	flagValueIter(flags, "-X", func(val string) {
+		// val is in the form of "pkg.name=str"
+		i := strings.IndexByte(val, '=')
+		if i <= 0 {
+			return
+		}
+		name := val[:i]
+		str := val[i+1:]
+		j := strings.IndexByte(name, '.')
+		if j <= 0 {
+			return
+		}
+		pkg := name[:j]
+		name = name[j+1:]
+
+		pkgPath := pkg
+		if pkgPath == "main" {
+			// The main package is known under its import path in
+			// the import config map.
+			pkgPath = buildInfo.firstImport
+		}
+		if id := buildInfo.imports[pkgPath].buildID; id != "" {
+			name = hashWith(id, name)
+			flags = append(flags, fmt.Sprintf("-X=%s.%s=%s", pkg, name, str))
+		}
+	})
+
 	flags = append(flags, "-w", "-s")
 	return append(flags, paths...), nil
 }
@@ -605,21 +646,32 @@ func splitFlagsFromFiles(args []string, ext string) (flags, paths []string) {
 }
 
 // flagValue retrieves the value of a flag such as "-foo", from strings in the
-// list of arguments like "-foo=bar" or "-foo" "bar".
+// list of arguments like "-foo=bar" or "-foo" "bar". If the flag is repeated,
+// the last value is returned.
 func flagValue(flags []string, name string) string {
+	lastVal := ""
+	flagValueIter(flags, name, func(val string) {
+		lastVal = val
+	})
+	return lastVal
+}
+
+// flagValueIter retrieves all the values for a flag such as "-foo", like
+// flagValue. The difference is that it allows handling complex flags, such as
+// those whose values compose a list.
+func flagValueIter(flags []string, name string, fn func(string)) {
 	for i, arg := range flags {
 		if val := strings.TrimPrefix(arg, name+"="); val != arg {
 			// -name=value
-			return val
+			fn(val)
 		}
 		if arg == name { // -name ...
 			if i+1 < len(flags) {
 				// -name value
-				return flags[i+1]
+				fn(flags[i+1])
 			}
 		}
 	}
-	return ""
 }
 
 func flagSetValue(flags []string, name, value string) []string {
diff --git a/testdata/scripts/ldflags.txt b/testdata/scripts/ldflags.txt
new file mode 100644
index 0000000..370ee2e
--- /dev/null
+++ b/testdata/scripts/ldflags.txt
@@ -0,0 +1,36 @@
+garble build -ldflags='-X=main.unexportedVersion=v1.0.0 -X=test/main/imported.ExportedVar=replaced'
+exec ./main
+cmp stdout main.stdout
+! binsubstr main$exe 'unexportedVersion'
+
+[short] stop # no need to verify this with -short
+
+exec go build -ldflags='-X=main.unexportedVersion=v1.0.0 -X=test/main/imported.ExportedVar=replaced'
+exec ./main
+cmp stdout main.stdout
+binsubstr main$exe 'unexportedVersion'
+
+-- go.mod --
+module test/main
+-- main.go --
+package main
+
+import (
+	"fmt"
+
+	"test/main/imported"
+)
+
+var unexportedVersion = "unknown"
+
+func main() {
+	fmt.Println("version:", unexportedVersion)
+	fmt.Println("var:", imported.ExportedVar)
+}
+-- imported/imported.go --
+package imported
+
+var ExportedVar = "original"
+-- main.stdout --
+version: v1.0.0
+var: replaced