From fcfefc15b17dd70ab249e3d8d09d1ccc5da7d347 Mon Sep 17 00:00:00 2001
From: salma-samy <salmasamy@google.com>
Date: Thu, 6 Jul 2023 14:26:10 -0700
Subject: [PATCH] handle exception instead of crashing Fixes:
 https://github.com/bazelbuild/bazel/issues/18641

PiperOrigin-RevId: 546092727
Change-Id: I66365813a40390ec61a1a0565b06655b3ca50fcd
---
 .../starlark/StarlarkRuleClassFunctions.java  | 10 ++++--
 src/test/py/bazel/bzlmod/bazel_module_test.py | 31 +++++++++++++++++++
 2 files changed, 39 insertions(+), 2 deletions(-)

diff --git a/src/main/java/com/google/devtools/build/lib/analysis/starlark/StarlarkRuleClassFunctions.java b/src/main/java/com/google/devtools/build/lib/analysis/starlark/StarlarkRuleClassFunctions.java
index 747ab3b4fe9e46..78991c3e776b0d 100644
--- a/src/main/java/com/google/devtools/build/lib/analysis/starlark/StarlarkRuleClassFunctions.java
+++ b/src/main/java/com/google/devtools/build/lib/analysis/starlark/StarlarkRuleClassFunctions.java
@@ -758,9 +758,15 @@ public Label getExtensionLabel() {
     public Object call(StarlarkThread thread, Tuple args, Dict<String, Object> kwargs)
         throws EvalException, InterruptedException {
       if (!args.isEmpty()) {
-        throw new EvalException("unexpected positional arguments");
+        throw new EvalException("Unexpected positional arguments");
+      }
+      try {
+        BazelStarlarkContext.from(thread).checkLoadingPhase(getName());
+      } catch (IllegalStateException e) {
+        throw new EvalException(
+            "A rule can only be instantiated in a BUILD file, or a macro "
+                + "invoked from a BUILD file");
       }
-      BazelStarlarkContext.from(thread).checkLoadingPhase(getName());
       if (ruleClass == null) {
         throw new EvalException("Invalid rule class hasn't been exported by a bzl file");
       }
diff --git a/src/test/py/bazel/bzlmod/bazel_module_test.py b/src/test/py/bazel/bzlmod/bazel_module_test.py
index 99ef7ff277485e..ed9ec4f1bf067a 100644
--- a/src/test/py/bazel/bzlmod/bazel_module_test.py
+++ b/src/test/py/bazel/bzlmod/bazel_module_test.py
@@ -666,6 +666,37 @@ def testWorkspaceToolchainRegistrationWithPlatformsConstraint(self):
     with open(self.Path('bazel-bin/my_consumer'), 'r') as f:
       self.assertEqual(f.read().strip(), 'my_value = Hello, Bzlmod!')
 
+  def testModuleExtensionWithRuleError(self):
+    self.ScratchFile(
+        'MODULE.bazel',
+        [
+            'ext = use_extension("extensions.bzl", "ext")',
+            'use_repo(ext, "ext")',
+        ],
+    )
+    self.ScratchFile('BUILD')
+    self.ScratchFile(
+        'extensions.bzl',
+        [
+            'def _rule_impl(ctx):',
+            '  print("RULE CALLED")',
+            'init_rule = rule(_rule_impl)',
+            'def ext_impl(module_ctx):',
+            '  init_rule()',
+            'ext = module_extension(implementation = ext_impl,)',
+        ],
+    )
+    exit_code, _, stderr = self.RunBazel(
+        ['build', '--nobuild', '@ext//:all'],
+        allow_failure=True,
+    )
+    self.AssertExitCode(exit_code, 48, stderr)
+    self.assertIn(
+        'Error in init_rule: A rule can only be instantiated in a BUILD file, '
+        'or a macro invoked from a BUILD file',
+        stderr,
+    )
+
 
 if __name__ == '__main__':
   unittest.main()