Fix attribute lookup for subrules used in aspects

Fixes https://github.com/bazelbuild/bazel/issues/23282

PiperOrigin-RevId: 667946340
Change-Id: Icac6c4e4c695e5b580c13e9aa9cd7cee608b8e3b
diff --git a/src/main/java/com/google/devtools/build/lib/analysis/starlark/StarlarkSubrule.java b/src/main/java/com/google/devtools/build/lib/analysis/starlark/StarlarkSubrule.java
index 5f3f0de..fc393c4 100644
--- a/src/main/java/com/google/devtools/build/lib/analysis/starlark/StarlarkSubrule.java
+++ b/src/main/java/com/google/devtools/build/lib/analysis/starlark/StarlarkSubrule.java
@@ -135,12 +135,7 @@
             "got invalid named argument: '%s' is an implicit dependency and cannot be overridden",
             attr.attrName);
       }
-      Attribute attribute =
-          ruleContext
-              .getRuleContext()
-              .getRule()
-              .getRuleClassObject()
-              .getAttributeByName(attr.ruleAttrName);
+      Attribute attribute = getAttributeByName(ruleContext, attr.ruleAttrName);
       // We need to use the underlying RuleContext because the subrule attributes are hidden from
       // the rule ctx.attr
       Object value;
@@ -179,6 +174,14 @@
     }
   }
 
+  private static Attribute getAttributeByName(StarlarkRuleContext ruleContext, String attr) {
+    if (ruleContext.isForAspect()) {
+      return ruleContext.getRuleContext().getMainAspect().getDefinition().getAttributes().get(attr);
+    } else {
+      return ruleContext.getRuleContext().getRule().getRuleClassObject().getAttributeByName(attr);
+    }
+  }
+
   private ImmutableSet<StarlarkSubrule> getDeclaredSubrules() {
     return subrules;
   }
diff --git a/src/test/java/com/google/devtools/build/lib/analysis/starlark/StarlarkSubruleTest.java b/src/test/java/com/google/devtools/build/lib/analysis/starlark/StarlarkSubruleTest.java
index ab715d5..88dc566 100644
--- a/src/test/java/com/google/devtools/build/lib/analysis/starlark/StarlarkSubruleTest.java
+++ b/src/test/java/com/google/devtools/build/lib/analysis/starlark/StarlarkSubruleTest.java
@@ -890,7 +890,7 @@
   }
 
   @Test
-  public void testSubruleAttrs_implicitLabelDepsAreResolvedToTargets() throws Exception {
+  public void testSubruleAttrs_implicitLabelDepsAreResolvedToTargets_inRule() throws Exception {
     scratch.file(
         "some/pkg/BUILD",
         //
@@ -932,6 +932,53 @@
   }
 
   @Test
+  public void testSubruleAttrs_implicitLabelDepsAreResolvedToTargets_inAspect() throws Exception {
+    scratch.file(
+        "some/pkg/BUILD",
+        //
+        "genrule(name = 'tool', cmd = '', outs = ['tool.exe'])");
+    scratch.file(
+        "subrule_testing/myrule.bzl",
+        """
+        def _subrule_impl(ctx, _tool):
+            return _tool
+
+        _my_subrule = subrule(
+            implementation = _subrule_impl,
+            attrs = {"_tool": attr.label(default = "//some/pkg:tool")},
+        )
+
+        MyInfo = provider()
+
+        def _aspect_impl(ctx, target):
+            res = _my_subrule()
+            return [MyInfo(result = res)]
+
+        _my_aspect = aspect(implementation = _aspect_impl, subrules = [_my_subrule])
+
+        my_rule = rule(
+            implementation = lambda ctx: [ctx.attr.dep[MyInfo]],
+            attrs = {"dep": attr.label(mandatory = True, aspects = [_my_aspect])},
+        )
+        """);
+    scratch.file(
+        "subrule_testing/BUILD",
+        """
+        load("myrule.bzl", "my_rule")
+        filegroup(name = 'bar')
+        my_rule(name = "foo", dep = ":bar")
+        """);
+
+    StructImpl provider =
+        getProvider("//subrule_testing:foo", "//subrule_testing:myrule.bzl", "MyInfo");
+
+    assertThat(provider).isNotNull();
+    Object value = provider.getValue("result");
+    assertThat(value).isInstanceOf(ConfiguredTarget.class);
+    assertThat(((ConfiguredTarget) value).getLabel().toString()).isEqualTo("//some/pkg:tool");
+  }
+
+  @Test
   public void testSubruleAttrs_singleFileLabelAttributesAreResolvedToFile() throws Exception {
     scratch.file(
         "some/pkg/BUILD",