From 97b95337b64aa25cc749c959bc33b3eda2e002dd Mon Sep 17 00:00:00 2001
From: "Michael B. Gale" <mbg@github.com>
Date: Wed, 8 May 2024 15:06:20 +0100
Subject: [PATCH] Go: Make toolchain functions resilient to different version
 formats

---
 go/extractor/toolchain/BUILD.bazel       |  1 +
 go/extractor/toolchain/toolchain.go      | 10 +++---
 go/extractor/toolchain/toolchain_test.go | 42 ++++++++++++++++++++++--
 3 files changed, 45 insertions(+), 8 deletions(-)

diff --git a/go/extractor/toolchain/BUILD.bazel b/go/extractor/toolchain/BUILD.bazel
index fde8d327e9e2..f3c0fc3d217c 100644
--- a/go/extractor/toolchain/BUILD.bazel
+++ b/go/extractor/toolchain/BUILD.bazel
@@ -17,4 +17,5 @@ go_test(
     name = "toolchain_test",
     srcs = ["toolchain_test.go"],
     embed = [":toolchain"],
+    deps = ["//go/extractor/util"],
 )
diff --git a/go/extractor/toolchain/toolchain.go b/go/extractor/toolchain/toolchain.go
index 03f42b5a851f..4f45c36f461f 100644
--- a/go/extractor/toolchain/toolchain.go
+++ b/go/extractor/toolchain/toolchain.go
@@ -25,7 +25,7 @@ var goVersions = map[string]struct{}{}
 
 // Adds an entry to the set of installed Go versions for the normalised `version` number.
 func addGoVersion(version string) {
-	goVersions[semver.Canonical("v"+version)] = struct{}{}
+	goVersions[semver.Canonical(util.FormatSemVer(version))] = struct{}{}
 }
 
 // Returns the current Go version as returned by 'go version', e.g. go1.14.4
@@ -58,7 +58,7 @@ func GetEnvGoVersion() string {
 
 // Determines whether, to our knowledge, `version` is available on the current system.
 func HasGoVersion(version string) bool {
-	_, found := goVersions[semver.Canonical("v"+version)]
+	_, found := goVersions[semver.Canonical(util.FormatSemVer(version))]
 	return found
 }
 
@@ -72,7 +72,7 @@ func InstallVersion(workingDir string, version string) bool {
 	// Construct a command to invoke `go version` with `GOTOOLCHAIN=go1.N.0` to give
 	// Go a valid toolchain version to download the toolchain we need; subsequent commands
 	// should then work even with an invalid version that's still in `go.mod`
-	toolchainArg := "GOTOOLCHAIN=go" + semver.Canonical("v" + version)[1:]
+	toolchainArg := "GOTOOLCHAIN=go" + semver.Canonical(util.FormatSemVer(version))[1:]
 	versionCmd := Version()
 	versionCmd.Dir = workingDir
 	versionCmd.Env = append(os.Environ(), toolchainArg)
@@ -112,9 +112,9 @@ func GoVersionToSemVer(goVersion string) string {
 	// which is compatible with the SemVer specification
 	rcIndex := strings.Index(goVersion, "rc")
 	if rcIndex != -1 {
-		return semver.Canonical("v"+goVersion[:rcIndex]) + "-" + goVersion[rcIndex:]
+		return semver.Canonical(util.FormatSemVer(goVersion[:rcIndex])) + "-" + goVersion[rcIndex:]
 	} else {
-		return semver.Canonical("v" + goVersion)
+		return semver.Canonical(util.FormatSemVer(goVersion))
 	}
 }
 
diff --git a/go/extractor/toolchain/toolchain_test.go b/go/extractor/toolchain/toolchain_test.go
index 8f60aa9a665f..574370117825 100644
--- a/go/extractor/toolchain/toolchain_test.go
+++ b/go/extractor/toolchain/toolchain_test.go
@@ -1,6 +1,10 @@
 package toolchain
 
-import "testing"
+import (
+	"testing"
+
+	"github.com/github/codeql-go/extractor/util"
+)
 
 func TestParseGoVersion(t *testing.T) {
 	tests := map[string]string{
@@ -16,9 +20,41 @@ func TestParseGoVersion(t *testing.T) {
 }
 
 func TestHasGoVersion(t *testing.T) {
-	if HasGoVersion("1.21") {
-		t.Error("Expected HasGoVersion(\"1.21\") to be false, but got true")
+	versions := []string{"1.21", "v1.22", "1.22.3", "v1.21rc4"}
+
+	// All versions should be unknown.
+	for _, version := range versions {
+		if HasGoVersion(version) {
+			t.Errorf("Expected HasGoVersion(\"%s\") to be false, but got true", version)
+		}
+
+		if HasGoVersion(util.FormatSemVer(version)) {
+			t.Errorf("Expected HasGoVersion(\"%s\") to be false, but got true", util.FormatSemVer(version))
+		}
+
+		if HasGoVersion(util.UnformatSemVer(version)) {
+			t.Errorf("Expected HasGoVersion(\"%s\") to be false, but got true", util.UnformatSemVer(version))
+		}
+
+		// Add the version in preparation for the next part of the test.
+		addGoVersion(version)
 	}
+
+	// Now we should have all of the versions.
+	for _, version := range versions {
+		if !HasGoVersion(version) {
+			t.Errorf("Expected HasGoVersion(\"%s\") to be true, but got false", version)
+		}
+
+		if !HasGoVersion(util.FormatSemVer(version)) {
+			t.Errorf("Expected HasGoVersion(\"%s\") to be true, but got false", util.FormatSemVer(version))
+		}
+
+		if !HasGoVersion(util.UnformatSemVer(version)) {
+			t.Errorf("Expected HasGoVersion(\"%s\") to be true, but got false", util.UnformatSemVer(version))
+		}
+	}
+
 }
 
 func testGoVersionToSemVer(t *testing.T, goVersion string, expectedSemVer string) {