Add nogo support to go_binary and go_test targets.

Updates #3374

PiperOrigin-RevId: 328378700

2 files changed, 53 insertions, 11 deletions

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
+    )