go: support embed configuration in stdlib packages (pantsbuild#18271)

Go standard library packages are permitted to use `//go:embed` directives. Add support for such embed configurations.

Fixes pantsbuild#18268.

src/python/pants/backend/go/util_rules/build_pkg_target.py

@@ -4,10 +4,12 @@
 from __future__ import annotations
 
 import dataclasses
+import json
 from dataclasses import dataclass
 from typing import ClassVar, Type, cast
 
 from pants.backend.go.dependency_inference import GoModuleImportPathsMapping
+from pants.backend.go.go_sources.load_go_binary import LoadedGoBinary, LoadedGoBinaryRequest
 from pants.backend.go.target_type_rules import GoImportPathMappingRequest
 from pants.backend.go.target_types import (
     GoAssemblerFlagsField,
@@ -40,7 +42,11 @@
     OwningGoModRequest,
 )
 from pants.backend.go.util_rules.goroot import GoRoot
-from pants.backend.go.util_rules.import_analysis import GoStdLibPackages, GoStdLibPackagesRequest
+from pants.backend.go.util_rules.import_analysis import (
+    GoStdLibPackage,
+    GoStdLibPackages,
+    GoStdLibPackagesRequest,
+)
 from pants.backend.go.util_rules.pkg_pattern import match_simple_pattern
 from pants.backend.go.util_rules.third_party_pkg import (
     ThirdPartyPkgAnalysis,
@@ -49,9 +55,11 @@
 from pants.build_graph.address import Address
 from pants.engine.engine_aware import EngineAwareParameter
 from pants.engine.environment import EnvironmentName
+from pants.engine.fs import CreateDigest, FileContent
 from pants.engine.internals.graph import AmbiguousCodegenImplementationsException
-from pants.engine.internals.native_engine import EMPTY_DIGEST
+from pants.engine.internals.native_engine import EMPTY_DIGEST, Digest, MergeDigests
 from pants.engine.internals.selectors import Get, MultiGet
+from pants.engine.process import FallibleProcessResult, Process
 from pants.engine.rules import collect_rules, rule
 from pants.engine.target import (
     Dependencies,
@@ -533,6 +541,59 @@ def _is_coverage_enabled_for_stdlib_package(import_path: str, build_opts: GoBuil
     return False
 
 
+@dataclass(frozen=True)
+class _ResolveStdlibEmbedConfigRequest:
+    package: GoStdLibPackage
+
+
+@dataclass(frozen=True)
+class _ResolveStdlibEmbedConfigResult:
+    embed_config: EmbedConfig | None
+    stderr: str | None
+
+
+@rule
+async def resolve_go_stdlib_embed_config(
+    request: _ResolveStdlibEmbedConfigRequest,
+) -> _ResolveStdlibEmbedConfigResult:
+    patterns_json = json.dumps(
+        {
+            "EmbedPatterns": request.package.embed_patterns,
+            "TestEmbedPatterns": [],
+            "XTestEmbedPatterns": [],
+        }
+    ).encode("utf-8")
+
+    embedder, patterns_json_digest = await MultiGet(
+        Get(LoadedGoBinary, LoadedGoBinaryRequest("embedcfg", ("main.go",), "./embedder")),
+        Get(Digest, CreateDigest([FileContent("patterns.json", patterns_json)])),
+    )
+    input_digest = await Get(
+        Digest,
+        MergeDigests((patterns_json_digest, embedder.digest)),
+    )
+    embed_result = await Get(
+        FallibleProcessResult,
+        Process(
+            ("./embedder", "patterns.json", request.package.pkg_source_path),
+            input_digest=input_digest,
+            description=f"Create embed mapping for {request.package.import_path}",
+            level=LogLevel.DEBUG,
+        ),
+    )
+    if embed_result.exit_code != 0:
+        return _ResolveStdlibEmbedConfigResult(
+            embed_config=None,
+            stderr=embed_result.stderr.decode(),
+        )
+    metadata = json.loads(embed_result.stdout)
+    embed_config = EmbedConfig.from_json_dict(metadata.get("EmbedConfig", {}))
+    return _ResolveStdlibEmbedConfigResult(
+        embed_config=embed_config,
+        stderr=None,
+    )
+
+
 @rule
 async def setup_build_go_package_target_request_for_stdlib(
     request: BuildGoPackageRequestForStdlibRequest,
@@ -580,6 +641,21 @@ async def setup_build_go_package_target_request_for_stdlib(
 
     with_coverage = _is_coverage_enabled_for_stdlib_package(request.import_path, request.build_opts)
 
+    embed_config: EmbedConfig | None = None
+    if pkg_info.embed_patterns and pkg_info.embed_files:
+        embed_config_result = await Get(
+            _ResolveStdlibEmbedConfigResult, _ResolveStdlibEmbedConfigRequest(pkg_info)
+        )
+        if not embed_config_result.embed_config:
+            assert embed_config_result.stderr is not None
+            return FallibleBuildGoPackageRequest(
+                request=None,
+                import_path=request.import_path,
+                exit_code=1,
+                stderr=embed_config_result.stderr,
+            )
+        embed_config = embed_config_result.embed_config
+
     return FallibleBuildGoPackageRequest(
         request=BuildGoPackageRequest(
             import_path=pkg_info.import_path,
@@ -602,6 +678,7 @@ async def setup_build_go_package_target_request_for_stdlib(
             cgo_flags=pkg_info.cgo_flags,
             with_coverage=with_coverage,
             is_stdlib=True,
+            embed_config=embed_config,
         ),
         import_path=request.import_path,
     )

src/python/pants/backend/go/util_rules/build_pkg_target_test.py

@@ -36,10 +36,12 @@
     FallibleBuiltGoPackage,
 )
 from pants.backend.go.util_rules.build_pkg_target import (
+    BuildGoPackageRequestForStdlibRequest,
     BuildGoPackageTargetRequest,
     GoCodegenBuildRequest,
 )
 from pants.backend.go.util_rules.go_mod import OwningGoMod, OwningGoModRequest
+from pants.backend.go.util_rules.import_analysis import GoStdLibPackages, GoStdLibPackagesRequest
 from pants.core.target_types import FilesGeneratorTarget, FileSourceField, FileTarget
 from pants.engine.addresses import Address, Addresses
 from pants.engine.fs import CreateDigest, Digest, FileContent, Snapshot
@@ -170,6 +172,8 @@ def rule_runner() -> RuleRunner:
             QueryRule(FallibleBuiltGoPackage, [BuildGoPackageRequest]),
             QueryRule(BuildGoPackageRequest, [BuildGoPackageTargetRequest]),
             QueryRule(FallibleBuildGoPackageRequest, [BuildGoPackageTargetRequest]),
+            QueryRule(GoStdLibPackages, (GoStdLibPackagesRequest,)),
+            QueryRule(BuildGoPackageRequest, (BuildGoPackageRequestForStdlibRequest,)),
             UnionRule(GoCodegenBuildRequest, GoCodegenBuildFilesRequest),
             UnionRule(GoModuleImportPathsMappingsHook, GenerateFromFileImportPathsMappingHook),
             FileTarget.register_plugin_field(GoOwningGoModAddressField),
@@ -640,3 +644,33 @@ def test_xtest_deps(rule_runner: RuleRunner) -> None:
         expected_direct_dependency_import_paths=[],
         expected_transitive_dependency_import_paths=[],
     )
+
+
+def test_stdlib_embed_config(rule_runner: RuleRunner) -> None:
+    import_path = "crypto/internal/nistec"
+    stdlib_packages = rule_runner.request(
+        GoStdLibPackages, [GoStdLibPackagesRequest(with_race_detector=False, cgo_enabled=False)]
+    )
+    pkg_info = stdlib_packages.get(import_path)
+    if not pkg_info:
+        pytest.skip(
+            f"Skipping test since `{import_path}` import path not available in Go standard library."
+        )
+
+    assert "embed" in pkg_info.imports
+    assert pkg_info.embed_patterns
+    assert pkg_info.embed_files
+
+    build_request = rule_runner.request(
+        BuildGoPackageRequest,
+        [
+            BuildGoPackageRequestForStdlibRequest(
+                import_path=import_path, build_opts=GoBuildOptions(cgo_enabled=False)
+            )
+        ],
+    )
+
+    embed_config = build_request.embed_config
+    assert embed_config is not None
+    assert embed_config.patterns
+    assert embed_config.files

src/python/pants/backend/go/util_rules/import_analysis.py

@@ -40,6 +40,13 @@ class GoStdLibPackage:
     syso_files: tuple[str, ...]
     cgo_flags: CGoCompilerFlags
 
+    # Embed configuration.
+    #
+    # Note: `EmbedConfig` is not resolved here to avoid issues with trying to build the the embed analyzer.
+    # The `EmbedConfig` will be resolved in `build_pkg_target.py` rules.
+    embed_patterns: tuple[str, ...]
+    embed_files: tuple[str, ...]
+
 
 class GoStdLibPackages(FrozenDict[str, GoStdLibPackage]):
     """A mapping of standard library import paths to an analysis of the package at that import
@@ -95,6 +102,8 @@ async def analyze_go_stdlib_packages(request: GoStdLibPackagesRequest) -> GoStdL
                 ldflags=tuple(pkg_json.get("CgoLDFLAGS", [])),
                 pkg_config=tuple(pkg_json.get("CgoPkgConfig", [])),
             ),
+            embed_patterns=tuple(pkg_json.get("EmbedPatterns", [])),
+            embed_files=tuple(pkg_json.get("EmbedFiles", [])),
         )
 
     return GoStdLibPackages(stdlib_packages)