Fix a few extension field bugs

protobuf/CHANGELOG.md

@@ -5,7 +5,16 @@
 
   The new limits are consistent with the Java and C++ implementations. ([#1060])
 
+* Fix `GeneratedMessage.addExtension` returning non-frozen and
+  `GeneratedMessage.getExtension` allowing modifying an extension when the
+  message is frozen before initializing the extension field set. ([#1062])
+
+* Fix `GeneratedMessage.getExtension` returning differently typed lists when the
+  message extension field set is initialized and frozen and initialized but not
+  frozen. ([#1062])
+
 [#1060]: https://github.com/google/protobuf.dart/pull/1060
+[#1062]: https://github.com/google/protobuf.dart/pull/1062
 
 ## 5.0.0
 

protobuf/lib/src/protobuf/extension_field_set.dart

@@ -58,7 +58,7 @@ class ExtensionFieldSet {
     final value = _values[fi.tagNumber];
     if (value != null) return value;
     _checkNotInUnknown(fi);
-    if (_isReadOnly) return PbList<T>.unmodifiable();
+    if (_isReadOnly) return fi._createRepeatedField()..freeze();
     return _addInfoAndCreateList<T>(fi);
   }
 
@@ -121,7 +121,6 @@ class ExtensionFieldSet {
         ),
       );
     }
-    _ensureWritable();
     _validateInfo(fi);
     _parent._validateField(fi, value);
     _addInfoUnchecked(fi);

protobuf/lib/src/protobuf/field_set.dart

@@ -106,8 +106,18 @@ class FieldSet {
   /// The [FieldInfo] for each non-extension field in tag order.
   Iterable<FieldInfo> get _infosSortedByTag => _meta.sortedByTag;
 
-  ExtensionFieldSet _ensureExtensions() =>
-      _extensions ??= ExtensionFieldSet(this);
+  ExtensionFieldSet _ensureExtensions() {
+    var extensions = _extensions;
+    if (extensions != null) {
+      return extensions;
+    }
+    extensions = ExtensionFieldSet(this);
+    _extensions = extensions;
+    if (_isReadOnly) {
+      extensions._markReadOnly();
+    }
+    return extensions;
+  }
 
   UnknownFieldSet _ensureUnknownFields() {
     if (_unknownFields == null) {
@@ -990,6 +1000,30 @@ class FieldSet {
       _oneofCases!.addAll(originalOneofCases);
     }
   }
+
+  bool hasExtension(Extension extension) =>
+      _extensions?._getFieldOrNull(extension) != null;
+
+  dynamic getExtension(Extension extension) =>
+      _ensureExtensions()._getFieldOrDefault(extension);
+
+  void setExtension(Extension extension, Object value) =>
+      _ensureExtensions()._setFieldAndInfo(extension, value);
+
+  void addExtension(Extension extension, Object? value) {
+    _ensureWritable();
+    if (!extension.isRepeated) {
+      throw ArgumentError(
+        'Cannot add to a non-repeated field (use setExtension())',
+      );
+    }
+    _ensureExtensions()._ensureRepeatedField(extension).add(value);
+  }
+
+  void clearExtension(Extension extension) {
+    _ensureWritable();
+    _extensions?._clearFieldAndInfo(extension);
+  }
 }
 
 extension FieldSetInternalExtension on FieldSet {

protobuf/lib/src/protobuf/generated_message.dart

@@ -334,19 +334,12 @@ abstract class GeneratedMessage {
   ///
   /// The backing [List] will be created if necessary.
   /// If the list already exists, the old extension won't be overwritten.
-  void addExtension(Extension extension, Object? value) {
-    if (!extension.isRepeated) {
-      throw ArgumentError(
-        'Cannot add to a non-repeated field (use setExtension())',
-      );
-    }
-    _fieldSet._ensureExtensions()._ensureRepeatedField(extension).add(value);
-  }
+  void addExtension(Extension extension, Object? value) =>
+      _fieldSet.addExtension(extension, value);
 
   /// Clears an extension field and also removes the extension.
-  void clearExtension(Extension extension) {
-    _fieldSet._extensions?._clearFieldAndInfo(extension);
-  }
+  void clearExtension(Extension extension) =>
+      _fieldSet.clearExtension(extension);
 
   /// Clears the contents of a given field.
   ///
@@ -364,7 +357,7 @@ abstract class GeneratedMessage {
   ///
   /// If not set, returns the extension's default value.
   dynamic getExtension(Extension extension) =>
-      _fieldSet._ensureExtensions()._getFieldOrDefault(extension);
+      _fieldSet.getExtension(extension);
 
   /// Returns the value of the field associated with [tagNumber], or the
   /// default value if it is not set.
@@ -386,8 +379,7 @@ abstract class GeneratedMessage {
       _fieldSet._ensureInfo(tagNumber).readonlyDefault;
 
   /// Returns `true` if a value of [extension] is present.
-  bool hasExtension(Extension extension) =>
-      _fieldSet._extensions?._getFieldOrNull(extension) != null;
+  bool hasExtension(Extension extension) => _fieldSet.hasExtension(extension);
 
   /// Whether this message has a field associated with [tagNumber].
   bool hasField(int tagNumber) => _fieldSet._hasField(tagNumber);
@@ -406,18 +398,8 @@ abstract class GeneratedMessage {
       .mergeFromUnknownFieldSet(unknownFieldSet);
 
   /// Sets the value of a non-repeated extension field to [value].
-  void setExtension(Extension extension, Object value) {
-    if (PbFieldType.isRepeated(extension.type)) {
-      throw ArgumentError(
-        _fieldSet._setFieldFailedMessage(
-          extension,
-          value,
-          'repeating field (use get + .add())',
-        ),
-      );
-    }
-    _fieldSet._ensureExtensions()._setFieldAndInfo(extension, value);
-  }
+  void setExtension(Extension extension, Object value) =>
+      _fieldSet.setExtension(extension, value);
 
   /// Sets the value of a field by its [tagNumber].
   ///

protoc_plugin/test/extension_test.dart

@@ -696,4 +696,82 @@ void main() {
     final d = r.reparseMessage(c);
     expect(m.hashCode, d.hashCode);
   });
+
+  test('setExtension throws when the extension field is repeated', () {
+    final m = Outer();
+    m.addExtension(Extend_unittest.extensionRepeated, 'hi');
+    expect(() {
+      m.setExtension(Extend_unittest.extensionRepeated, 'bye');
+    }, throwsArgumentError);
+  });
+
+  test('addExtension throws when the extension field is not repeated', () {
+    final m = Outer();
+    m.setExtension(Extend_unittest.extensionInner, Inner());
+    expect(() {
+      m.addExtension(Extend_unittest.extensionInner, Inner());
+    }, throwsArgumentError);
+  });
+
+  test('addExtension throws when the message is frozen', () {
+    final m = Outer()..freeze();
+    expect(
+      () => m.addExtension(Extend_unittest.extensionRepeated, 'hi'),
+      throwsUnsupportedError,
+    );
+  });
+
+  test('setExtension throws when the message is frozen', () {
+    final m = Outer()..freeze();
+    expect(
+      () => m.setExtension(Extend_unittest.extensionInner, Inner()),
+      throwsUnsupportedError,
+    );
+  });
+
+  test(
+    'getExtension returns frozen non-repeated value when the parent message is frozen',
+    () {
+      final m = Outer()..freeze();
+      final Inner ext = m.getExtension(Extend_unittest.extensionInner);
+      expect(() => ext.value = 'hi', throwsUnsupportedError);
+    },
+  );
+
+  test(
+    'getExtension returns frozen repeated value when the parent message is frozen',
+    () {
+      final m = Outer()..freeze();
+      final List ext = m.getExtension(Extend_unittest.extensionRepeated);
+      expect(() => ext.add('hi'), throwsUnsupportedError);
+    },
+  );
+
+  test(
+    'getExtension repeated value turns frozen when the parent message is frozen',
+    () {
+      final m = Outer();
+      final List ext = m.getExtension(Extend_unittest.extensionRepeated);
+      ext.add('hi');
+      m.freeze();
+      expect(() => ext.add('bye'), throwsUnsupportedError);
+    },
+  );
+
+  test(
+    'getExtension repeated value type is right when the field set is frozen',
+    () {
+      {
+        final m = Outer();
+        final ext = m.getExtension(Extend_unittest.extensionRepeated);
+        expect(ext, isA<PbList<String>>());
+      }
+
+      {
+        final m = Outer().freeze();
+        final ext = m.getExtension(Extend_unittest.extensionRepeated);
+        expect(ext, isA<PbList<String>>());
+      }
+    },
+  );
 }

protoc_plugin/test/protos/extend_unittest.proto

@@ -42,3 +42,7 @@ extend InnerMost {
 extend protobuf_unittest.TestAllExtensions {
   optional Outer outer = 104;
 }
+
+extend Outer {
+    repeated string extension_repeated = 6;
+}