From cd63ea84aabb6414f3471cd4cb99480c245a148f Mon Sep 17 00:00:00 2001
From: Sauyon Lee <sauyon@github.com>
Date: Thu, 10 Sep 2020 00:46:36 -0700
Subject: [PATCH] extractor: revamp argument parsing

---
 extractor/cli/go-extractor/go-extractor.go | 68 ++++++++++++++++++----
 1 file changed, 58 insertions(+), 10 deletions(-)

diff --git a/extractor/cli/go-extractor/go-extractor.go b/extractor/cli/go-extractor/go-extractor.go
index 8ab34255e77..4b2dbe2297e 100644
--- a/extractor/cli/go-extractor/go-extractor.go
+++ b/extractor/cli/go-extractor/go-extractor.go
@@ -20,23 +20,70 @@ func usage() {
 	fmt.Fprintf(os.Stderr, "--help                Print this help.\n")
 }
 
-func parseFlags(args []string) ([]string, []string) {
+func parseFlags(args []string, mimic bool) ([]string, []string) {
 	i := 0
 	buildFlags := []string{}
-	for i < len(args) && strings.HasPrefix(args[i], "-") {
+	for ; i < len(args) && strings.HasPrefix(args[i], "-"); i++ {
 		if args[i] == "--" {
 			i++
 			break
 		}
 
-		if args[i] == "--help" {
-			usage()
-			os.Exit(0)
-		} else {
-			buildFlags = append(buildFlags, args[i])
+		if !mimic {
+			// we're not in mimic mode, try to parse our arguments
+			switch args[i] {
+			case "--help":
+				usage()
+				os.Exit(0)
+			case "--mimic":
+				if i+1 < len(args) {
+					i++
+					compiler := args[i]
+					log.Printf("Compiler: %s", compiler)
+					if i+1 < len(args) {
+						i++
+						command := args[i]
+						if command == "build" || command == "install" || command == "run" {
+							log.Printf("Intercepting build")
+							return parseFlags(args[i+1:], true)
+						} else {
+							log.Printf("Non-build command '%s'; skipping", strings.Join(args[1:], " "))
+							os.Exit(0)
+						}
+					} else {
+						log.Printf("Non-build command '%s'; skipping", strings.Join(args[1:], " "))
+						os.Exit(0)
+					}
+				} else {
+					log.Fatalf("Invalid --mimic: no compiler specified")
+				}
+			}
 		}
 
-		i++
+		// parse go build flags
+		switch args[i] {
+		// skip `-o output` and `-i`, if applicable
+		case "-o":
+			if i+1 < len(args) {
+				i++
+			}
+		case "-i":
+		case "-p", "-asmflags", "-buildmode", "-compiler", "-gccgoflags", "-gcflags", "-installsuffix",
+			"-ldflags", "-mod", "-modfile", "-pkgdir", "-tags", "-toolexec":
+			if i+1 < len(args) {
+				buildFlags = append(buildFlags, args[i], args[i+1])
+				i++
+			} else {
+				buildFlags = append(buildFlags, args[i])
+			}
+		default:
+			if strings.HasPrefix(args[i], "-") {
+				buildFlags = append(buildFlags, args[i])
+			} else {
+				// stop parsing if the argument is not a flag (and so is positional)
+				break
+			}
+		}
 	}
 
 	cpuprofile = os.Getenv("CODEQL_EXTRACTOR_GO_CPU_PROFILE")
@@ -46,7 +93,7 @@ func parseFlags(args []string) ([]string, []string) {
 }
 
 func main() {
-	buildFlags, patterns := parseFlags(os.Args[1:])
+	buildFlags, patterns := parseFlags(os.Args[1:], false)
 
 	if cpuprofile != "" {
 		f, err := os.Create(cpuprofile)
@@ -63,9 +110,10 @@ func main() {
 	if len(patterns) == 0 {
 		log.Println("Nothing to extract.")
 	} else {
+		log.Printf("Build flags: '%s'; patterns: '%s'\n", strings.Join(buildFlags, " "), strings.Join(patterns, " "))
 		err := extractor.ExtractWithFlags(buildFlags, patterns)
 		if err != nil {
-			log.Fatal(err)
+			log.Fatalf("Error running go tooling: %s\n", err.Error())
 		}
 	}