Add nogo support to go_binary and go_test targets.

Updates #3374

PiperOrigin-RevId: 328378700

6 files changed, 101 insertions, 14 deletions

diff --git a/tools/bazeldefs/defs.bzl b/tools/bazeldefs/defs.bzl
index db7f379b8..4bbcda054 100644
--- a/tools/bazeldefs/defs.bzl
+++ b/tools/bazeldefs/defs.bzl
@@ -87,13 +87,14 @@ def cc_binary(name, static = False, **kwargs):
         **kwargs
     )
 
-def go_binary(name, static = False, pure = False, **kwargs):
+def go_binary(name, static = False, pure = False, x_defs = None, **kwargs):
     """Build a go binary.
 
     Args:
         name: name of the target.
         static: build a static binary.
         pure: build without cgo.
+        x_defs: additional definitions.
         **kwargs: rest of the arguments are passed to _go_binary.
     """
     if static:
@@ -102,6 +103,7 @@ def go_binary(name, static = False, pure = False, **kwargs):
         kwargs["pure"] = "on"
     _go_binary(
         name = name,
+        x_defs = x_defs,
         **kwargs
     )
 
@@ -151,6 +153,11 @@ def go_rule(rule, implementation, **kwargs):
     toolchains = kwargs.get("toolchains", []) + ["@io_bazel_rules_go//go:toolchain"]
     return rule(implementation, attrs = attrs, toolchains = toolchains, **kwargs)
 
+def go_test_library(target):
+    if hasattr(target.attr, "embed") and len(target.attr.embed) > 0:
+        return target.attr.embed[0]
+    return None
+
 def go_context(ctx):
     go_ctx = _go_context(ctx)
     return struct(
diff --git a/tools/defs.bzl b/tools/defs.bzl
index e71a26cf4..290d564f2 100644
--- a/tools/defs.bzl
+++ b/tools/defs.bzl
@@ -27,7 +27,6 @@ gbenchmark = _gbenchmark
 gazelle = _gazelle
 go_embed_data = _go_embed_data
 go_path = _go_path
-go_test = _go_test
 gtest = _gtest
 grpcpp = _grpcpp
 loopback = _loopback
@@ -45,17 +44,35 @@ vdso_linker_option = _vdso_linker_option
 default_platform = _default_platform
 platforms = _platforms
 
-def go_binary(name, **kwargs):
+def go_binary(name, nogo = True, pure = False, static = False, x_defs = None, **kwargs):
     """Wraps the standard go_binary.
 
     Args:
       name: the rule name.
+      nogo: enable nogo analysis.
+      pure: build a pure Go (no CGo) binary.
+      static: build a static binary.
+      x_defs: additional linker definitions.
       **kwargs: standard go_binary arguments.
     """
     _go_binary(
         name = name,
+        pure = pure,
+        static = static,
+        x_defs = x_defs,
         **kwargs
     )
+    if nogo:
+        # Note that the nogo rule applies only for go_library and go_test
+        # targets, therefore we construct a library from the binary sources.
+        _go_library(
+            name = name + "_nogo_library",
+            **kwargs
+        )
+        nogo_test(
+            name = name + "_nogo",
+            deps = [":" + name + "_nogo_library"],
+        )
 
 def calculate_sets(srcs):
     """Calculates special Go sets for templates.
@@ -119,6 +136,7 @@ def go_library(name, srcs, deps = [], imports = [], stateify = True, marshal = F
       stateify: whether statify is enabled (default: true).
       marshal: whether marshal is enabled (default: false).
       marshal_debug: whether the gomarshal tools emits debugging output (default: false).
+      nogo: enable nogo analysis.
       **kwargs: standard go_library arguments.
     """
     all_srcs = srcs
@@ -202,6 +220,24 @@ def go_library(name, srcs, deps = [], imports = [], stateify = True, marshal = F
                 **kwargs
             )
 
+def go_test(name, nogo = True, **kwargs):
+    """Wraps the standard go_test.
+
+    Args:
+      name: the rule name.
+      nogo: enable nogo analysis.
+      **kwargs: standard go_test arguments.
+    """
+    _go_test(
+        name = name,
+        **kwargs
+    )
+    if nogo:
+        nogo_test(
+            name = name + "_nogo",
+            deps = [":" + name],
+        )
+
 def proto_library(name, srcs, deps = None, has_services = 0, **kwargs):
     """Wraps the standard proto_library.
 
diff --git a/tools/issue_reviver/BUILD b/tools/issue_reviver/BUILD
index 4ef1a3124..35b0111ca 100644
--- a/tools/issue_reviver/BUILD
+++ b/tools/issue_reviver/BUILD
@@ -5,6 +5,7 @@ package(licenses = ["notice"])
 go_binary(
     name = "issue_reviver",
     srcs = ["main.go"],
+    nogo = False,
     deps = [
         "//tools/issue_reviver/github",
         "//tools/issue_reviver/reviver",
diff --git a/tools/issue_reviver/github/BUILD b/tools/issue_reviver/github/BUILD
index 0eabc2835..555abd296 100644
--- a/tools/issue_reviver/github/BUILD
+++ b/tools/issue_reviver/github/BUILD
@@ -21,4 +21,5 @@ go_test(
     size = "small",
     srcs = ["github_test.go"],
     library = ":github",
+    nogo = False,
 )
diff --git a/tools/nogo/check/BUILD b/tools/nogo/check/BUILD
index e2d76cd5c..21ba2c306 100644
--- a/tools/nogo/check/BUILD
+++ b/tools/nogo/check/BUILD
@@ -7,6 +7,7 @@ package(licenses = ["notice"])
 go_binary(
     name = "check",
     srcs = ["main.go"],
+    nogo = False,
     visibility = ["//visibility:public"],
     deps = ["//tools/nogo"],
 )
diff --git a/tools/nogo/defs.bzl b/tools/nogo/defs.bzl
index d399079c5..5377620b0 100644
--- a/tools/nogo/defs.bzl
+++ b/tools/nogo/defs.bzl
@@ -1,6 +1,6 @@
 """Nogo rules."""
 
-load("//tools/bazeldefs:defs.bzl", "go_context", "go_importpath", "go_rule")
+load("//tools/bazeldefs:defs.bzl", "go_context", "go_importpath", "go_rule", "go_test_library")
 
 # NogoInfo is the serialized set of package facts for a nogo analysis.
 #
@@ -8,10 +8,13 @@ load("//tools/bazeldefs:defs.bzl", "go_context", "go_importpath", "go_rule")
 # with the source files as input. Note however, that the individual nogo rules
 # are simply stubs that enter into the shadow dependency tree (the "aspect").
 NogoInfo = provider(
+    "information for nogo analysis",
     fields = {
         "facts": "serialized package facts",
         "importpath": "package import path",
         "binaries": "package binary files",
+        "srcs": "original source files (for go_test support)",
+        "deps": "original deps (for go_test support)",
     },
 )
 
@@ -21,16 +24,29 @@ def _nogo_aspect_impl(target, ctx):
     # All work is done in the shadow properties for go rules. For a proto
     # library, we simply skip the analysis portion but still need to return a
     # valid NogoInfo to reference the generated binary.
-    if ctx.rule.kind == "go_library":
+    if ctx.rule.kind in ("go_library", "go_binary", "go_test", "go_tool_library"):
         srcs = ctx.rule.files.srcs
-    elif ctx.rule.kind == "go_proto_library" or ctx.rule.kind == "go_wrap_cc":
+        deps = ctx.rule.attr.deps
+    elif ctx.rule.kind in ("go_proto_library", "go_wrap_cc"):
         srcs = []
+        deps = ctx.rule.attr.deps
     else:
         return [NogoInfo()]
 
-    go_ctx = go_context(ctx)
+    # If we're using the "library" attribute, then we need to aggregate the
+    # original library sources and dependencies into this target to perform
+    # proper type analysis.
+    if ctx.rule.kind == "go_test":
+        library = go_test_library(ctx.rule)
+        if library != None:
+            info = library[NogoInfo]
+            if hasattr(info, "srcs"):
+                srcs = srcs + info.srcs
+            if hasattr(info, "deps"):
+                deps = deps + info.deps
 
     # Construct the Go environment from the go_ctx.env dictionary.
+    go_ctx = go_context(ctx)
     env_prefix = " ".join(["%s=%s" % (key, value) for (key, value) in go_ctx.env.items()])
 
     # Start with all target files and srcs as input.
@@ -41,6 +57,13 @@ def _nogo_aspect_impl(target, ctx):
     # to cleanly allow us redirect stdout to the actual output file. Perhaps
     # I'm missing something here, but the intermediate script does work.
     binaries = target.files.to_list()
+    objfiles = [f for f in binaries if f.path.endswith(".a")]
+    if len(objfiles) > 0:
+        # Prefer the .a files for go_library targets.
+        target_objfile = objfiles[0]
+    else:
+        # Use the raw binary for go_binary and go_test targets.
+        target_objfile = binaries[0]
     disasm_file = ctx.actions.declare_file(target.label.name + ".out")
     dumper = ctx.actions.declare_file("%s-dumper" % ctx.label.name)
     ctx.actions.write(dumper, "\n".join([
@@ -48,12 +71,12 @@ def _nogo_aspect_impl(target, ctx):
         "%s %s tool objdump %s > %s\n" % (
             env_prefix,
             go_ctx.go.path,
-            [f.path for f in binaries if f.path.endswith(".a")][0],
+            target_objfile.path,
             disasm_file.path,
         ),
     ]), is_executable = True)
     ctx.actions.run(
-        inputs = binaries,
+        inputs = [target_objfile],
         outputs = [disasm_file],
         tools = go_ctx.runfiles,
         mnemonic = "GoObjdump",
@@ -63,7 +86,15 @@ def _nogo_aspect_impl(target, ctx):
     inputs.append(disasm_file)
 
     # Extract the importpath for this package.
-    importpath = go_importpath(target)
+    if ctx.rule.kind == "go_test":
+        # If this is a test, then it will not be imported by anything else.
+        # We can safely set the importapth to just "test". Note that this
+        # is necessary if the library also imports the core library (in
+        # addition to including the sources directly), which happens in
+        # some complex cases (seccomp_victim).
+        importpath = "test"
+    else:
+        importpath = go_importpath(target)
 
     # The nogo tool requires a configfile serialized in JSON format to do its
     # work. This must line up with the nogo.Config fields.
@@ -84,7 +115,7 @@ def _nogo_aspect_impl(target, ctx):
     )
 
     # Collect all info from shadow dependencies.
-    for dep in ctx.rule.attr.deps:
+    for dep in deps:
         # There will be no file attribute set for all transitive dependencies
         # that are not go_library or go_binary rules, such as a proto rules.
         # This is handled by the ctx.rule.kind check above.
@@ -126,12 +157,18 @@ def _nogo_aspect_impl(target, ctx):
         facts = facts,
         importpath = importpath,
         binaries = binaries,
+        srcs = srcs,
+        deps = deps,
     )]
 
 nogo_aspect = go_rule(
     aspect,
     implementation = _nogo_aspect_impl,
-    attr_aspects = ["deps"],
+    attr_aspects = [
+        "deps",
+        "library",
+        "embed",
+    ],
     attrs = {
         "_nogo": attr.label(
             default = "//tools/nogo/check:check",
@@ -171,6 +208,10 @@ _nogo_test = rule(
     test = True,
 )
 
-def nogo_test(**kwargs):
+def nogo_test(name, **kwargs):
     tags = kwargs.pop("tags", []) + ["nogo"]
-    _nogo_test(tags = tags, **kwargs)
+    _nogo_test(
+        name = name,
+        tags = tags,
+        **kwargs
+    )