From f6d46fb89f4b6569d672ae7676f68a284a89b21e Mon Sep 17 00:00:00 2001
From: Patrick Ziegler <49288965+pziggo@users.noreply.github.com>
Date: Mon, 18 Sep 2023 15:03:17 +0200
Subject: [PATCH] go/tools/gopackagesdriver: fix package walk panic if bzlmod
 enabled

With bzlmod enabled, also canonical repository names are enabled. Thus
labels from the local repository have to be prefixed with two '@' (in
order to refer to the canonical repository name) to match the IDs in
the package registry.
Also, the result of a bazel query does not return the canonical name of
the rules_go repository, but the apparent repository name. Thus, the
query for go targets it will not match the stdlib label as constructed
with the native.repository_name() rule as injected via x_defs.

This fix must be considered as workaround as it leads to a mix of
canonical and apparant repository names in the package lists. Once
bazel query also supports returning the canonincal repository names, a
more correct fix can be applied.
---
 go/tools/gopackagesdriver/main.go            | 23 ++++++++++----------
 go/tools/gopackagesdriver/packageregistry.go |  8 +++++--
 go/tools/gopackagesdriver/utils.go           | 12 ++++++++++
 3 files changed, 30 insertions(+), 13 deletions(-)

diff --git a/go/tools/gopackagesdriver/main.go b/go/tools/gopackagesdriver/main.go
index 8cc904d3ce..3d0f708153 100644
--- a/go/tools/gopackagesdriver/main.go
+++ b/go/tools/gopackagesdriver/main.go
@@ -53,17 +53,18 @@ type driverResponse struct {
 var (
 	// Injected via x_defs.
 
-	rulesGoRepositoryName string
-	goDefaultAspect       = rulesGoRepositoryName + "//go/tools/gopackagesdriver:aspect.bzl%go_pkg_info_aspect"
-	bazelBin              = getenvDefault("GOPACKAGESDRIVER_BAZEL", "bazel")
-	bazelStartupFlags     = strings.Fields(os.Getenv("GOPACKAGESDRIVER_BAZEL_FLAGS"))
-	bazelQueryFlags       = strings.Fields(os.Getenv("GOPACKAGESDRIVER_BAZEL_QUERY_FLAGS"))
-	bazelQueryScope       = getenvDefault("GOPACKAGESDRIVER_BAZEL_QUERY_SCOPE", "")
-	bazelBuildFlags       = strings.Fields(os.Getenv("GOPACKAGESDRIVER_BAZEL_BUILD_FLAGS"))
-	workspaceRoot         = os.Getenv("BUILD_WORKSPACE_DIRECTORY")
-	additionalAspects     = strings.Fields(os.Getenv("GOPACKAGESDRIVER_BAZEL_ADDTL_ASPECTS"))
-	additionalKinds       = strings.Fields(os.Getenv("GOPACKAGESDRIVER_BAZEL_KINDS"))
-	emptyResponse         = &driverResponse{
+	rulesGoRepositoryName               string
+	goDefaultAspect                     = rulesGoRepositoryName + "//go/tools/gopackagesdriver:aspect.bzl%go_pkg_info_aspect"
+	bazelSupportsCanonicalLabelLiterals = strings.HasPrefix(rulesGoRepositoryName, "@@")
+	bazelBin                            = getenvDefault("GOPACKAGESDRIVER_BAZEL", "bazel")
+	bazelStartupFlags                   = strings.Fields(os.Getenv("GOPACKAGESDRIVER_BAZEL_FLAGS"))
+	bazelQueryFlags                     = strings.Fields(os.Getenv("GOPACKAGESDRIVER_BAZEL_QUERY_FLAGS"))
+	bazelQueryScope                     = getenvDefault("GOPACKAGESDRIVER_BAZEL_QUERY_SCOPE", "")
+	bazelBuildFlags                     = strings.Fields(os.Getenv("GOPACKAGESDRIVER_BAZEL_BUILD_FLAGS"))
+	workspaceRoot                       = os.Getenv("BUILD_WORKSPACE_DIRECTORY")
+	additionalAspects                   = strings.Fields(os.Getenv("GOPACKAGESDRIVER_BAZEL_ADDTL_ASPECTS"))
+	additionalKinds                     = strings.Fields(os.Getenv("GOPACKAGESDRIVER_BAZEL_KINDS"))
+	emptyResponse                       = &driverResponse{
 		NotHandled: true,
 		Compiler:   "gc",
 		Arch:       runtime.GOARCH,
diff --git a/go/tools/gopackagesdriver/packageregistry.go b/go/tools/gopackagesdriver/packageregistry.go
index 92f9eac707..32d1dd31c8 100644
--- a/go/tools/gopackagesdriver/packageregistry.go
+++ b/go/tools/gopackagesdriver/packageregistry.go
@@ -89,10 +89,14 @@ func (pr *PackageRegistry) Match(labels []string) ([]string, []*FlatPackage) {
 
 	for _, label := range labels {
 		if !strings.HasPrefix(label, "@") {
-			label = fmt.Sprintf("@%s", label)
+			if bazelSupportsCanonicalLabelLiterals {
+				label = fmt.Sprintf("@@%s", label)
+			} else {
+				label = fmt.Sprintf("@%s", label)
+			}
 		}
 
-		if label == RulesGoStdlibLabel {
+		if isRulesGoStdlibLabel(label) {
 			// For stdlib, we need to append all the subpackages as roots
 			// since RulesGoStdLibLabel doesn't actually show up in the stdlib pkg.json
 			for _, pkg := range pr.packagesByID {
diff --git a/go/tools/gopackagesdriver/utils.go b/go/tools/gopackagesdriver/utils.go
index d5524fddd1..4da80b5110 100644
--- a/go/tools/gopackagesdriver/utils.go
+++ b/go/tools/gopackagesdriver/utils.go
@@ -65,6 +65,18 @@ func isLocalPattern(pattern string) bool {
 	return build.IsLocalImport(pattern) || filepath.IsAbs(pattern)
 }
 
+func isRulesGoStdlibLabel(pattern string) bool {
+	if pattern == RulesGoStdlibLabel {
+		return true
+	} else if pattern == "@io_bazel_rules_go//:stdlib" {
+		// Check also for a static repository name since bazel query does
+		// not return the canonical repository name when bzlmod is enabled.
+		// The static name is also used while building the stdliblist.
+		return true
+	}
+	return false
+}
+
 func packageID(pattern string) string {
 	pattern = path.Clean(pattern)
 	if filepath.IsAbs(pattern) {