Allow disabling lazy parsing in individual `ExtensionRegistryLite` which will signal the parsing whether should parse extension fields lazily.

PiperOrigin-RevId: 968674119

java/core/src/main/java/com/google/protobuf/ExtensionRegistryLite.java

diff --git a/java/core/src/main/java/com/google/protobuf/ExtensionRegistryLite.java b/java/core/src/main/java/com/google/protobuf/ExtensionRegistryLite.java
index 61d69e5..b876fe8 100644
--- a/java/core/src/main/java/com/google/protobuf/ExtensionRegistryLite.java
+++ b/java/core/src/main/java/com/google/protobuf/ExtensionRegistryLite.java
@@ -60,6 +60,12 @@
 
   private static volatile LazyExtensionMode lazyExtensionMode = LazyExtensionMode.EAGER;
 
+  // Override for the lazy extension mode for this specific registry.
+  // -1 means fallback to the static lazyExtensionMode
+  // 0 means eager
+  // 1 means LAZY_VERIFY_ON_ACCESS
+  private byte lazyExtensionModeOverride = -1;
+
   static void setLazyExtensionMode(LazyExtensionMode mode) {
     lazyExtensionMode = mode;
   }
@@ -68,10 +74,28 @@
     return lazyExtensionMode;
   }
 
-  static boolean lazyExtensionEnabled() {
+  boolean lazyExtensionEnabled() {
+    if (lazyExtensionModeOverride == 1) {
+      return true;
+    } else if (lazyExtensionModeOverride == 0) {
+      return false;
+    }
     return lazyExtensionMode == LazyExtensionMode.LAZY_VERIFY_ON_ACCESS;
   }
 
+  /**
+   * Returns a new {@link ExtensionRegistryLite} with the same contents as this registry but with
+   * the lazy extension mode overridden.
+   *
+   * @param lazy whether to enable lazy extensions
+   * @return a new {@link ExtensionRegistryLite} with the lazy extension mode overridden
+   */
+  ExtensionRegistryLite withLazyExtensionsOverride(boolean lazy) {
+    ExtensionRegistryLite ret = getUnmodifiable();
+    ret.lazyExtensionModeOverride = (byte) (lazy ? 1 : 0);
+    return ret;
+  }
+
   // Visible for testing.
   static final String EXTENSION_CLASS_NAME = "com.google.protobuf.Extension";
 
@@ -192,6 +216,7 @@
     } else {
       this.extensionsByNumber = Collections.unmodifiableMap(other.extensionsByNumber);
     }
+    this.lazyExtensionModeOverride = other.lazyExtensionModeOverride;
   }
 
   private final Map<ObjectIntPair, GeneratedMessageLite.GeneratedExtension<?, ?>>

java/core/src/main/java/com/google/protobuf/InternalLazyField.java

diff --git a/java/core/src/main/java/com/google/protobuf/InternalLazyField.java b/java/core/src/main/java/com/google/protobuf/InternalLazyField.java
index 81b106a..959d3c1 100644
--- a/java/core/src/main/java/com/google/protobuf/InternalLazyField.java
+++ b/java/core/src/main/java/com/google/protobuf/InternalLazyField.java
@@ -215,7 +215,7 @@
         // extension, and we should fall back to the old behavior of silently returning the default
         // instance on corrupted extensions i.e. full parse.
         value =
-            ExtensionRegistryLite.lazyExtensionEnabled()
+            extensionRegistry.lazyExtensionEnabled()
                 ? defaultInstance.getParserForType().parsePartialFrom(bytes, extensionRegistry)
                 : defaultInstance.getParserForType().parseFrom(bytes, extensionRegistry);
       } catch (InvalidProtocolBufferException e) {
@@ -237,7 +237,7 @@
       ensureInitialized();
       return value;
     } catch (InvalidProtocolBufferException e) {
-      if (ExtensionRegistryLite.lazyExtensionEnabled()) {
+      if (extensionRegistry.lazyExtensionEnabled()) {
         // New behavior: runtime exception on corrupted extensions.
         throw new InvalidProtobufRuntimeException(e);
       } else {

java/core/src/main/java/com/google/protobuf/MessageReflection.java

diff --git a/java/core/src/main/java/com/google/protobuf/MessageReflection.java b/java/core/src/main/java/com/google/protobuf/MessageReflection.java
index 064df1d..4359364 100644
--- a/java/core/src/main/java/com/google/protobuf/MessageReflection.java
+++ b/java/core/src/main/java/com/google/protobuf/MessageReflection.java
@@ -1071,7 +1071,7 @@
         throws IOException {
       if (!field.isRepeated()) {
         boolean isLazyField =
-            ExtensionRegistryLite.lazyExtensionEnabled()
+            extensionRegistry.lazyExtensionEnabled()
                 && field.isExtension()
                 && !field.getContainingType().isDescriptorProtoType();
         if (hasField(field)) {

java/core/src/test/java/com/google/protobuf/ParserTest.java

diff --git a/java/core/src/test/java/com/google/protobuf/ParserTest.java b/java/core/src/test/java/com/google/protobuf/ParserTest.java
index d4a4685..0279a6f 100644
--- a/java/core/src/test/java/com/google/protobuf/ParserTest.java
+++ b/java/core/src/test/java/com/google/protobuf/ParserTest.java
@@ -487,4 +487,41 @@
     assertThat(thrown).hasMessageThat().contains("invalid tag");
     ExtensionRegistryLite.setLazyExtensionMode(originalMode);
   }
+
+  private static final byte[] MISSING_REQUIRED_EXTENSION_BYTES =
+      createMissingRequiredExtensionBytes();
+
+  private static byte[] createMissingRequiredExtensionBytes() {
+    TestMergeException.Builder message = TestMergeException.newBuilder();
+    message
+        .getAllExtensionsBuilder()
+        .setExtension(TestRequired.single, TestRequired.newBuilder().setA(1).buildPartial());
+    ByteString byteString = message.buildPartial().toByteString();
+    return byteString.concat(byteString).toByteArray();
+  }
+
+  @Test
+  public void testLazyExtensionsOverride_true_doesNotThrow() throws Exception {
+    ExtensionRegistry registry = ExtensionRegistry.newInstance();
+    UnittestProto.registerAllExtensions(registry);
+
+    // Should pass without throwing exception when overridden to lazy
+    TestMergeException result =
+        TestMergeException.parseFrom(
+            MISSING_REQUIRED_EXTENSION_BYTES, registry.withLazyExtensionsOverride(true));
+    assertThat(result).isNotNull();
+  }
+
+  @Test
+  public void testLazyExtensionsOverride_false_throwsException() throws Exception {
+    ExtensionRegistry registry = ExtensionRegistry.newInstance();
+    UnittestProto.registerAllExtensions(registry);
+
+    // Should throw exception when overridden to eager
+    assertThrows(
+        InvalidProtocolBufferException.class,
+        () ->
+            TestMergeException.parseFrom(
+                MISSING_REQUIRED_EXTENSION_BYTES, registry.withLazyExtensionsOverride(false)));
+  }
 }