Add support for extensions in CRuby, JRuby, and FFI Ruby (#14703) (#14756)

Follow up to #14594, which added support for custom options, this PR implements extensions support, which should fully resolve #1198.

Closes #14703

COPYBARA_INTEGRATE_REVIEW=https://github.com/protocolbuffers/protobuf/pull/14703 from protocolbuffers:add-support-for-extensions-in-ruby 601aca4121212c1633e9c7cd8abc65aebe9da9df
PiperOrigin-RevId: 582460674

Co-authored-by: Jason Lunn <jason.lunn@gmail.com>

13 files changed, 257 insertions, 55 deletions

diff --git a/ruby/ext/google/protobuf_c/defs.c b/ruby/ext/google/protobuf_c/defs.c
index 4e8054717..ea9185321 100644
--- a/ruby/ext/google/protobuf_c/defs.c
+++ b/ruby/ext/google/protobuf_c/defs.c
@@ -144,20 +144,26 @@ VALUE DescriptorPool_add_serialized_file(VALUE _self,
  * call-seq:
  *     DescriptorPool.lookup(name) => descriptor
  *
- * Finds a Descriptor or EnumDescriptor by name and returns it, or nil if none
- * exists with the given name.
+ * Finds a Descriptor, EnumDescriptor or FieldDescriptor by name and returns it,
+ * or nil if none exists with the given name.
  */
 static VALUE DescriptorPool_lookup(VALUE _self, VALUE name) {
   DescriptorPool* self = ruby_to_DescriptorPool(_self);
   const char* name_str = get_str(name);
   const upb_MessageDef* msgdef;
   const upb_EnumDef* enumdef;
+  const upb_FieldDef* fielddef;
 
   msgdef = upb_DefPool_FindMessageByName(self->symtab, name_str);
   if (msgdef) {
     return get_msgdef_obj(_self, msgdef);
   }
 
+  fielddef = upb_DefPool_FindExtensionByName(self->symtab, name_str);
+  if (fielddef) {
+    return get_fielddef_obj(_self, fielddef);
+  }
+
   enumdef = upb_DefPool_FindEnumByName(self->symtab, name_str);
   if (enumdef) {
     return get_enumdef_obj(_self, enumdef);
diff --git a/ruby/ext/google/protobuf_c/message.c b/ruby/ext/google/protobuf_c/message.c
index 2dec31a89..d64d411f8 100644
--- a/ruby/ext/google/protobuf_c/message.c
+++ b/ruby/ext/google/protobuf_c/message.c
@@ -977,9 +977,12 @@ VALUE Message_decode_bytes(int size, const char* bytes, int options,
   VALUE msg_rb = initialize_rb_class_with_no_args(klass);
   Message* msg = ruby_to_Message(msg_rb);
 
+  const upb_FileDef* file = upb_MessageDef_File(msg->msgdef);
+  const upb_ExtensionRegistry* extreg =
+      upb_DefPool_ExtensionRegistry(upb_FileDef_Pool(file));
   upb_DecodeStatus status = upb_Decode(bytes, size, (upb_Message*)msg->msg,
                                        upb_MessageDef_MiniTable(msg->msgdef),
-                                       NULL, options, Arena_get(msg->arena));
+                                       extreg, options, Arena_get(msg->arena));
   if (status != kUpb_DecodeStatus_Ok) {
     rb_raise(cParseError, "Error occurred during parsing");
   }
@@ -1303,9 +1306,12 @@ upb_Message* Message_deep_copy(const upb_Message* msg, const upb_MessageDef* m,
   upb_Message* new_msg = upb_Message_New(layout, arena);
   char* data;
 
+  const upb_FileDef* file = upb_MessageDef_File(m);
+  const upb_ExtensionRegistry* extreg =
+      upb_DefPool_ExtensionRegistry(upb_FileDef_Pool(file));
   if (upb_Encode(msg, layout, 0, tmp_arena, &data, &size) !=
           kUpb_EncodeStatus_Ok ||
-      upb_Decode(data, size, new_msg, layout, NULL, 0, arena) !=
+      upb_Decode(data, size, new_msg, layout, extreg, 0, arena) !=
           kUpb_DecodeStatus_Ok) {
     upb_Arena_Free(tmp_arena);
     rb_raise(cParseError, "Error occurred copying proto");
diff --git a/ruby/lib/google/protobuf/ffi/descriptor_pool.rb b/ruby/lib/google/protobuf/ffi/descriptor_pool.rb
index f0543adbb..96c7d0941 100644
--- a/ruby/lib/google/protobuf/ffi/descriptor_pool.rb
+++ b/ruby/lib/google/protobuf/ffi/descriptor_pool.rb
@@ -9,13 +9,16 @@ module Google
   module Protobuf
     class FFI
       # DefPool
-      attach_function :add_serialized_file,   :upb_DefPool_AddFile,           [:DefPool, :FileDescriptorProto, Status.by_ref], :FileDef
-      attach_function :free_descriptor_pool,  :upb_DefPool_Free,              [:DefPool], :void
-      attach_function :create_descriptor_pool,:upb_DefPool_New,               [], :DefPool
-      attach_function :lookup_enum,           :upb_DefPool_FindEnumByName,    [:DefPool, :string], EnumDescriptor
-      attach_function :lookup_msg,            :upb_DefPool_FindMessageByName, [:DefPool, :string], Descriptor
-      # FileDescriptorProto
-      attach_function :parse,                 :FileDescriptorProto_parse,     [:binary_string, :size_t, Internal::Arena], :FileDescriptorProto
+      attach_function :add_serialized_file,   :upb_DefPool_AddFile,            [:DefPool, :FileDescriptorProto, Status.by_ref], :FileDef
+      attach_function :free_descriptor_pool,  :upb_DefPool_Free,               [:DefPool], :void
+      attach_function :create_descriptor_pool,:upb_DefPool_New,                [], :DefPool
+      attach_function :get_extension_registry,:upb_DefPool_ExtensionRegistry,  [:DefPool],  :ExtensionRegistry
+      attach_function :lookup_enum,           :upb_DefPool_FindEnumByName,     [:DefPool, :string], EnumDescriptor
+      attach_function :lookup_extension,      :upb_DefPool_FindExtensionByName,[:DefPool, :string], FieldDescriptor
+      attach_function :lookup_msg,            :upb_DefPool_FindMessageByName,  [:DefPool, :string], Descriptor
+
+        # FileDescriptorProto
+      attach_function :parse,                 :FileDescriptorProto_parse,      [:binary_string, :size_t, Internal::Arena], :FileDescriptorProto
     end
     class DescriptorPool
       attr :descriptor_pool
@@ -50,7 +53,8 @@ module Google
 
       def lookup name
         Google::Protobuf::FFI.lookup_msg(@descriptor_pool, name) ||
-          Google::Protobuf::FFI.lookup_enum(@descriptor_pool, name)
+          Google::Protobuf::FFI.lookup_enum(@descriptor_pool, name) ||
+          Google::Protobuf::FFI.lookup_extension(@descriptor_pool, name)
       end
 
       def self.generated_pool
diff --git a/ruby/lib/google/protobuf/ffi/message.rb b/ruby/lib/google/protobuf/ffi/message.rb
index 045f67f90..39eb40385 100644
--- a/ruby/lib/google/protobuf/ffi/message.rb
+++ b/ruby/lib/google/protobuf/ffi/message.rb
@@ -170,7 +170,15 @@ module Google
 
             message = new
             mini_table_ptr = Google::Protobuf::FFI.get_mini_table(message.class.descriptor)
-            status = Google::Protobuf::FFI.decode_message(data, data.bytesize, message.instance_variable_get(:@msg), mini_table_ptr, nil, decoding_options, message.instance_variable_get(:@arena))
+            status = Google::Protobuf::FFI.decode_message(
+              data,
+              data.bytesize,
+              message.instance_variable_get(:@msg),
+              mini_table_ptr,
+              Google::Protobuf::FFI.get_extension_registry(message.class.descriptor.send(:pool).descriptor_pool),
+              decoding_options,
+              message.instance_variable_get(:@arena)
+            )
             raise ParseError.new "Error occurred during parsing" unless status == :Ok
             message
           end
diff --git a/ruby/src/main/java/com/google/protobuf/jruby/RubyDescriptorPool.java b/ruby/src/main/java/com/google/protobuf/jruby/RubyDescriptorPool.java
index d65b412a0..7c01eb927 100644
--- a/ruby/src/main/java/com/google/protobuf/jruby/RubyDescriptorPool.java
+++ b/ruby/src/main/java/com/google/protobuf/jruby/RubyDescriptorPool.java
@@ -36,7 +36,9 @@ import com.google.protobuf.DescriptorProtos.FileDescriptorProto;
 import com.google.protobuf.Descriptors.Descriptor;
 import com.google.protobuf.Descriptors.DescriptorValidationException;
 import com.google.protobuf.Descriptors.EnumDescriptor;
+import com.google.protobuf.Descriptors.FieldDescriptor;
 import com.google.protobuf.Descriptors.FileDescriptor;
+import com.google.protobuf.ExtensionRegistry;
 import com.google.protobuf.InvalidProtocolBufferException;
 import java.util.ArrayList;
 import java.util.HashMap;
@@ -70,6 +72,7 @@ public class RubyDescriptorPool extends RubyObject {
             cDescriptorPool.newInstance(runtime.getCurrentContext(), Block.NULL_BLOCK);
     cDescriptor = (RubyClass) runtime.getClassFromPath("Google::Protobuf::Descriptor");
     cEnumDescriptor = (RubyClass) runtime.getClassFromPath("Google::Protobuf::EnumDescriptor");
+    cFieldDescriptor = (RubyClass) runtime.getClassFromPath("Google::Protobuf::FieldDescriptor");
   }
 
   public RubyDescriptorPool(Ruby runtime, RubyClass klazz) {
@@ -92,7 +95,7 @@ public class RubyDescriptorPool extends RubyObject {
    * call-seq:
    *     DescriptorPool.lookup(name) => descriptor
    *
-   * Finds a Descriptor or EnumDescriptor by name and returns it, or nil if none
+   * Finds a Descriptor, EnumDescriptor or FieldDescriptor by name and returns it, or nil if none
    * exists with the given name.
    *
    * This currently lazy loads the ruby descriptor objects as they are requested.
@@ -121,7 +124,8 @@ public class RubyDescriptorPool extends RubyObject {
   public IRubyObject add_serialized_file(ThreadContext context, IRubyObject data) {
     byte[] bin = data.convertToString().getBytes();
     try {
-      FileDescriptorProto.Builder builder = FileDescriptorProto.newBuilder().mergeFrom(bin);
+      FileDescriptorProto.Builder builder =
+          FileDescriptorProto.newBuilder().mergeFrom(bin, registry);
       registerFileDescriptor(context, builder);
     } catch (InvalidProtocolBufferException e) {
       throw RaiseException.from(
@@ -150,6 +154,8 @@ public class RubyDescriptorPool extends RubyObject {
     for (EnumDescriptor ed : fd.getEnumTypes()) registerEnumDescriptor(context, ed, packageName);
     for (Descriptor message : fd.getMessageTypes())
       registerDescriptor(context, message, packageName);
+    for (FieldDescriptor fieldDescriptor : fd.getExtensions())
+      registerExtension(context, fieldDescriptor, packageName);
 
     // Mark this as a loaded file
     fileDescriptors.add(fd);
@@ -170,6 +176,24 @@ public class RubyDescriptorPool extends RubyObject {
       registerEnumDescriptor(context, ed, fullPath);
     for (Descriptor message : descriptor.getNestedTypes())
       registerDescriptor(context, message, fullPath);
+    for (FieldDescriptor fieldDescriptor : descriptor.getExtensions())
+      registerExtension(context, fieldDescriptor, fullPath);
+  }
+
+  private void registerExtension(
+      ThreadContext context, FieldDescriptor descriptor, String parentPath) {
+    if (descriptor.getJavaType() == FieldDescriptor.JavaType.MESSAGE) {
+      registry.add(descriptor, descriptor.toProto());
+    } else {
+      registry.add(descriptor);
+    }
+    RubyString name = context.runtime.newString(parentPath + descriptor.getName());
+    RubyFieldDescriptor des =
+        (RubyFieldDescriptor) cFieldDescriptor.newInstance(context, Block.NULL_BLOCK);
+    des.setName(name);
+    des.setDescriptor(context, descriptor, this);
+    // For MessageSet extensions, there is the possibility of a name conflict. Prefer the Message.
+    symtab.putIfAbsent(name, des);
   }
 
   private void registerEnumDescriptor(
@@ -188,8 +212,10 @@ public class RubyDescriptorPool extends RubyObject {
 
   private static RubyClass cDescriptor;
   private static RubyClass cEnumDescriptor;
+  private static RubyClass cFieldDescriptor;
   private static RubyDescriptorPool descriptorPool;
 
   private List<FileDescriptor> fileDescriptors;
   private Map<IRubyObject, IRubyObject> symtab;
+  protected static final ExtensionRegistry registry = ExtensionRegistry.newInstance();
 }
diff --git a/ruby/src/main/java/com/google/protobuf/jruby/RubyFieldDescriptor.java b/ruby/src/main/java/com/google/protobuf/jruby/RubyFieldDescriptor.java
index 5ba86ef2a..1d647ea8d 100644
--- a/ruby/src/main/java/com/google/protobuf/jruby/RubyFieldDescriptor.java
+++ b/ruby/src/main/java/com/google/protobuf/jruby/RubyFieldDescriptor.java
@@ -103,6 +103,10 @@ public class RubyFieldDescriptor extends RubyObject {
     return this.name;
   }
 
+  protected void setName(IRubyObject name) {
+    this.name = name;
+  }
+
   /*
    * call-seq:
    *     FieldDescriptor.subtype => message_or_enum_descriptor
@@ -229,7 +233,7 @@ public class RubyFieldDescriptor extends RubyObject {
    */
   @JRubyMethod(name = "set")
   public IRubyObject setValue(ThreadContext context, IRubyObject message, IRubyObject value) {
-    ((RubyMessage) message).setField(context, descriptor, value);
+    ((RubyMessage) message).setField(context, this, value);
     return context.nil;
   }
 
@@ -263,6 +267,10 @@ public class RubyFieldDescriptor extends RubyObject {
     this.pool = pool;
   }
 
+  protected FieldDescriptor getDescriptor() {
+    return descriptor;
+  }
+
   private void calculateLabel(ThreadContext context) {
     if (descriptor.isRepeated()) {
       this.label = context.runtime.newSymbol("repeated");
diff --git a/ruby/src/main/java/com/google/protobuf/jruby/RubyMessage.java b/ruby/src/main/java/com/google/protobuf/jruby/RubyMessage.java
index 25f9dca58..12e609191 100644
--- a/ruby/src/main/java/com/google/protobuf/jruby/RubyMessage.java
+++ b/ruby/src/main/java/com/google/protobuf/jruby/RubyMessage.java
@@ -634,7 +634,7 @@ public class RubyMessage extends RubyObject {
   public static IRubyObject decodeBytes(
       ThreadContext context, RubyMessage ret, CodedInputStream input, boolean freeze) {
     try {
-      ret.builder.mergeFrom(input);
+      ret.builder.mergeFrom(input, RubyDescriptorPool.registry);
     } catch (Exception e) {
       throw RaiseException.from(
           context.runtime,
@@ -965,6 +965,12 @@ public class RubyMessage extends RubyObject {
     return setFieldInternal(context, fieldDescriptor, value);
   }
 
+  protected IRubyObject setField(
+      ThreadContext context, RubyFieldDescriptor fieldDescriptor, IRubyObject value) {
+    validateMessageType(context, fieldDescriptor.getDescriptor(), "set");
+    return setFieldInternal(context, fieldDescriptor.getDescriptor(), fieldDescriptor, value);
+  }
+
   private RubyRepeatedField getRepeatedField(
       ThreadContext context, FieldDescriptor fieldDescriptor) {
     if (fields.containsKey(fieldDescriptor)) {
@@ -1275,6 +1281,14 @@ public class RubyMessage extends RubyObject {
 
   private IRubyObject setFieldInternal(
       ThreadContext context, FieldDescriptor fieldDescriptor, IRubyObject value) {
+    return setFieldInternal(context, fieldDescriptor, null, value);
+  }
+
+  private IRubyObject setFieldInternal(
+      ThreadContext context,
+      FieldDescriptor fieldDescriptor,
+      RubyFieldDescriptor rubyFieldDescriptor,
+      IRubyObject value) {
     testFrozen("can't modify frozen " + getMetaClass());
 
     if (fieldDescriptor.isMapField()) {
@@ -1299,8 +1313,12 @@ public class RubyMessage extends RubyObject {
       // Determine the typeclass, if any
       IRubyObject typeClass = context.runtime.getObject();
       if (fieldType == FieldDescriptor.Type.MESSAGE) {
-        typeClass =
-            ((RubyDescriptor) getDescriptorForField(context, fieldDescriptor)).msgclass(context);
+        if (rubyFieldDescriptor != null) {
+          typeClass = ((RubyDescriptor) rubyFieldDescriptor.getSubtype(context)).msgclass(context);
+        } else {
+          typeClass =
+              ((RubyDescriptor) getDescriptorForField(context, fieldDescriptor)).msgclass(context);
+        }
         if (value.isNil()) {
           addValue = false;
         }
diff --git a/ruby/tests/basic.rb b/ruby/tests/basic.rb
index 9cd2d705e..fe2119e42 100755
--- a/ruby/tests/basic.rb
+++ b/ruby/tests/basic.rb
@@ -729,6 +729,19 @@ module BasicTest
       oneof_descriptor = descriptor.lookup_oneof("test_deprecated_message_oneof")
 
       assert_instance_of Google::Protobuf::OneofOptions, oneof_descriptor.options
+      test_top_level_option = Google::Protobuf::DescriptorPool.generated_pool.lookup 'basic_test.test_top_level_option'
+      assert_instance_of Google::Protobuf::FieldDescriptor, test_top_level_option
+      assert_equal "Custom option value", test_top_level_option.get(oneof_descriptor.options)
+    end
+
+    def test_nested_extension
+      descriptor = TestDeprecatedMessage.descriptor
+      oneof_descriptor = descriptor.lookup_oneof("test_deprecated_message_oneof")
+
+      assert_instance_of Google::Protobuf::OneofOptions, oneof_descriptor.options
+      test_nested_option = Google::Protobuf::DescriptorPool.generated_pool.lookup 'basic_test.TestDeprecatedMessage.test_nested_option'
+      assert_instance_of Google::Protobuf::FieldDescriptor, test_nested_option
+      assert_equal "Another custom option value", test_nested_option.get(oneof_descriptor.options)
     end
 
     def test_options_deep_freeze
@@ -739,6 +752,25 @@ module BasicTest
           Google::Protobuf::UninterpretedOption.new
       end
     end
+
+    def test_message_deep_freeze
+      message = TestDeprecatedMessage.new
+      omit(":internal_deep_freeze only exists under FFI") unless message.respond_to? :internal_deep_freeze, true
+      nested_message_2 = TestMessage2.new
+
+      message.map_string_msg["message"] = TestMessage2.new
+      message.repeated_msg.push(TestMessage2.new)
+
+      message.send(:internal_deep_freeze)
+
+      assert_raise FrozenError do
+        message.map_string_msg["message"].foo = "bar"
+      end
+
+      assert_raise FrozenError do
+        message.repeated_msg[0].foo = "bar"
+      end
+    end
   end
 
   def test_oneof_fields_respond_to? # regression test for issue 9202
diff --git a/ruby/tests/basic_proto2.rb b/ruby/tests/basic_proto2.rb
index e9e99e545..1eefc6ff8 100755
--- a/ruby/tests/basic_proto2.rb
+++ b/ruby/tests/basic_proto2.rb
@@ -269,5 +269,57 @@ module BasicTestProto2
       assert msg.respond_to? :has_d?
       refute msg.has_d?
     end
+
+    def test_extension
+      message = TestExtensions.new
+      extension = Google::Protobuf::DescriptorPool.generated_pool.lookup 'basic_test_proto2.optional_int32_extension'
+      assert_instance_of Google::Protobuf::FieldDescriptor, extension
+      assert_equal 0, extension.get(message)
+      extension.set message, 42
+      assert_equal 42, extension.get(message)
+    end
+
+    def test_nested_extension
+      message = TestExtensions.new
+      extension = Google::Protobuf::DescriptorPool.generated_pool.lookup 'basic_test_proto2.TestNestedExtension.test'
+      assert_instance_of Google::Protobuf::FieldDescriptor, extension
+      assert_equal 'test', extension.get(message)
+      extension.set message, 'another test'
+      assert_equal 'another test', extension.get(message)
+    end
+
+    def test_message_set_extension_json_roundtrip
+      omit "Java Protobuf JsonFormat does not handle Proto2 extensions" if defined? JRUBY_VERSION and :NATIVE == Google::Protobuf::IMPLEMENTATION
+      message = TestMessageSet.new
+      ext1 = Google::Protobuf::DescriptorPool.generated_pool.lookup 'basic_test_proto2.TestMessageSetExtension1.message_set_extension'
+      assert_instance_of Google::Protobuf::FieldDescriptor, ext1
+      ext2 = Google::Protobuf::DescriptorPool.generated_pool.lookup 'basic_test_proto2.TestMessageSetExtension2.message_set_extension'
+      assert_instance_of Google::Protobuf::FieldDescriptor, ext2
+      ext3 = Google::Protobuf::DescriptorPool.generated_pool.lookup 'basic_test_proto2.message_set_extension3'
+      assert_instance_of Google::Protobuf::FieldDescriptor, ext3
+      ext1.set(message, ext1.subtype.msgclass.new(i: 42))
+      ext2.set(message, ext2.subtype.msgclass.new(str: 'foo'))
+      ext3.set(message, ext3.subtype.msgclass.new(text: 'bar'))
+      message_text = message.to_json
+      parsed_message = TestMessageSet.decode_json message_text
+      assert_equal message, parsed_message
+    end
+
+
+    def test_message_set_extension_roundtrip
+      message = TestMessageSet.new
+      ext1 = Google::Protobuf::DescriptorPool.generated_pool.lookup 'basic_test_proto2.TestMessageSetExtension1.message_set_extension'
+      assert_instance_of Google::Protobuf::FieldDescriptor, ext1
+      ext2 = Google::Protobuf::DescriptorPool.generated_pool.lookup 'basic_test_proto2.TestMessageSetExtension2.message_set_extension'
+      assert_instance_of Google::Protobuf::FieldDescriptor, ext2
+      ext3 = Google::Protobuf::DescriptorPool.generated_pool.lookup 'basic_test_proto2.message_set_extension3'
+      assert_instance_of Google::Protobuf::FieldDescriptor, ext3
+      ext1.set(message, ext1.subtype.msgclass.new(i: 42))
+      ext2.set(message, ext2.subtype.msgclass.new(str: 'foo'))
+      ext3.set(message, ext3.subtype.msgclass.new(text: 'bar'))
+      encoded_message = TestMessageSet.encode message
+      decoded_message = TestMessageSet.decode encoded_message
+      assert_equal message, decoded_message
+    end
   end
 end
diff --git a/ruby/tests/basic_test.proto b/ruby/tests/basic_test.proto
index 89eb9cc3c..8feab6ed6 100644
--- a/ruby/tests/basic_test.proto
+++ b/ruby/tests/basic_test.proto
@@ -2,6 +2,7 @@ syntax = "proto3";
 
 package basic_test;
 
+import "google/protobuf/descriptor.proto";
 import "google/protobuf/duration.proto";
 import "google/protobuf/struct.proto";
 import "google/protobuf/timestamp.proto";
@@ -70,12 +71,23 @@ message TestMessage2 {
   optional int32 foo = 1;
 }
 
+extend google.protobuf.OneofOptions {
+  optional string test_top_level_option = 1000;
+}
+
 message TestDeprecatedMessage {
   option deprecated = true;
 
+  extend google.protobuf.OneofOptions {
+    optional string test_nested_option = 1001;
+  }
+
   optional int32 foo = 1 [deprecated = true];
 
   oneof test_deprecated_message_oneof {
+    option (test_top_level_option) = "Custom option value";
+    option (test_nested_option) = "Another custom option value";
+
     string a = 2;
     int32 b = 3;
   }
diff --git a/ruby/tests/basic_test_proto2.proto b/ruby/tests/basic_test_proto2.proto
index ac705ed63..777b4dd77 100644
--- a/ruby/tests/basic_test_proto2.proto
+++ b/ruby/tests/basic_test_proto2.proto
@@ -2,10 +2,10 @@ syntax = "proto2";
 
 package basic_test_proto2;
 
-import "google/protobuf/wrappers.proto";
-import "google/protobuf/timestamp.proto";
 import "google/protobuf/duration.proto";
 import "google/protobuf/struct.proto";
+import "google/protobuf/timestamp.proto";
+import "google/protobuf/wrappers.proto";
 
 message Foo {
   optional Bar bar = 1;
@@ -188,3 +188,46 @@ message MyStruct {
   optional string string = 1;
   optional google.protobuf.Struct struct = 2;
 }
+
+message TestExtensions {
+  extensions 1 to max;
+}
+
+message TestNestedExtension {
+  extend TestExtensions {
+    optional string test = 1002 [default = "test"];
+  }
+}
+
+extend TestExtensions {
+  optional int32 optional_int32_extension = 1;
+}
+
+// A message with message_set_wire_format.
+message TestMessageSet {
+  option message_set_wire_format = true;
+
+  extensions 4 to max;
+}
+
+message TestMessageSetExtension1 {
+  extend TestMessageSet {
+    optional TestMessageSetExtension1 message_set_extension = 98418603;
+  }
+  optional int32 i = 15;
+}
+
+message TestMessageSetExtension2 {
+  extend TestMessageSet {
+    optional TestMessageSetExtension2 message_set_extension = 98418634;
+  }
+  optional string str = 25;
+}
+
+message TestMessageSetExtension3 {
+  optional string text = 35;
+}
+
+extend TestMessageSet {
+  optional TestMessageSetExtension3 message_set_extension3 = 98418655;
+}
diff --git a/src/google/protobuf/compiler/ruby/ruby_generator.cc b/src/google/protobuf/compiler/ruby/ruby_generator.cc
index 7bc828337..12a5c12e7 100644
--- a/src/google/protobuf/compiler/ruby/ruby_generator.cc
+++ b/src/google/protobuf/compiler/ruby/ruby_generator.cc
@@ -40,7 +40,7 @@ void GenerateEnumAssignment(absl::string_view prefix, const EnumDescriptor* en,
                             io::Printer* printer);
 std::string DefaultValueForField(const FieldDescriptor* field);
 
-template<class numeric_type>
+template <class numeric_type>
 std::string NumberToString(numeric_type value) {
   std::ostringstream os;
   os << value;
@@ -65,7 +65,6 @@ bool IsAlpha(char ch) { return IsLower(ch) || IsUpper(ch); }
 
 char UpperChar(char ch) { return IsLower(ch) ? (ch - 'a' + 'A') : ch; }
 
-
 // Package names in protobuf are snake_case by convention, but Ruby module
 // names must be PascalCased.
 //
@@ -123,14 +122,12 @@ void GenerateMessageAssignment(absl::string_view prefix,
     return;
   }
 
+  printer->Print("$prefix$$name$ = ", "prefix", prefix, "name",
+                 RubifyConstant(message->name()));
   printer->Print(
-    "$prefix$$name$ = ",
-    "prefix", prefix,
-    "name", RubifyConstant(message->name()));
-  printer->Print(
-    "::Google::Protobuf::DescriptorPool.generated_pool."
-    "lookup(\"$full_name$\").msgclass\n",
-    "full_name", message->full_name());
+      "::Google::Protobuf::DescriptorPool.generated_pool."
+      "lookup(\"$full_name$\").msgclass\n",
+      "full_name", message->full_name());
 
   std::string nested_prefix =
       absl::StrCat(prefix, RubifyConstant(message->name()), "::");
@@ -144,14 +141,12 @@ void GenerateMessageAssignment(absl::string_view prefix,
 
 void GenerateEnumAssignment(absl::string_view prefix, const EnumDescriptor* en,
                             io::Printer* printer) {
+  printer->Print("$prefix$$name$ = ", "prefix", prefix, "name",
+                 RubifyConstant(en->name()));
   printer->Print(
-    "$prefix$$name$ = ",
-    "prefix", prefix,
-    "name", RubifyConstant(en->name()));
-  printer->Print(
-    "::Google::Protobuf::DescriptorPool.generated_pool."
-    "lookup(\"$full_name$\").enummodule\n",
-    "full_name", en->full_name());
+      "::Google::Protobuf::DescriptorPool.generated_pool."
+      "lookup(\"$full_name$\").enummodule\n",
+      "full_name", en->full_name());
 }
 
 int GeneratePackageModules(const FileDescriptor* file, io::Printer* printer) {
@@ -197,9 +192,7 @@ int GeneratePackageModules(const FileDescriptor* file, io::Printer* printer) {
     if (need_change_to_module) {
       component = PackageToModule(component);
     }
-    printer->Print(
-      "module $name$\n",
-      "name", component);
+    printer->Print("module $name$\n", "name", component);
     printer->Indent();
     levels++;
   }
@@ -210,8 +203,7 @@ void EndPackageModules(int levels, io::Printer* printer) {
   while (levels > 0) {
     levels--;
     printer->Outdent();
-    printer->Print(
-      "end\n");
+    printer->Print("end\n");
   }
 }
 
@@ -304,16 +296,12 @@ bool GenerateFile(const FileDescriptor* file, io::Printer* printer,
 
   if (file->dependency_count() != 0) {
     for (int i = 0; i < file->dependency_count(); i++) {
-      printer->Print("require '$name$'\n", "name", GetRequireName(file->dependency(i)->name()));
+      printer->Print("require '$name$'\n", "name",
+                     GetRequireName(file->dependency(i)->name()));
     }
     printer->Print("\n");
   }
 
-  // TODO: Remove this when ruby supports extensions.
-  if (file->extension_count() > 0) {
-    ABSL_LOG(WARNING) << "Extensions are not yet supported in Ruby.";
-  }
-
   GenerateBinaryDescriptor(file, printer, error);
 
   int levels = GeneratePackageModules(file, printer);
@@ -328,11 +316,10 @@ bool GenerateFile(const FileDescriptor* file, io::Printer* printer,
   return true;
 }
 
-bool Generator::Generate(
-    const FileDescriptor* file,
-    const std::string& parameter,
-    GeneratorContext* generator_context,
-    std::string* error) const {
+bool Generator::Generate(const FileDescriptor* file,
+                         const std::string& parameter,
+                         GeneratorContext* generator_context,
+                         std::string* error) const {
   if (FileDescriptorLegacy(file).syntax() ==
       FileDescriptorLegacy::Syntax::SYNTAX_UNKNOWN) {
     *error = "Invalid or unsupported proto syntax";
diff --git a/upb/reflection/def_pool.h b/upb/reflection/def_pool.h
index f21a0dbd2..513e29de8 100644
--- a/upb/reflection/def_pool.h
+++ b/upb/reflection/def_pool.h
@@ -48,7 +48,7 @@ const upb_FileDef* upb_DefPool_FindFileByNameWithSize(const upb_DefPool* s,
 const upb_FieldDef* upb_DefPool_FindExtensionByMiniTable(
     const upb_DefPool* s, const upb_MiniTableExtension* ext);
 
-const upb_FieldDef* upb_DefPool_FindExtensionByName(const upb_DefPool* s,
+UPB_API const upb_FieldDef* upb_DefPool_FindExtensionByName(const upb_DefPool* s,
                                                     const char* sym);
 
 const upb_FieldDef* upb_DefPool_FindExtensionByNameWithSize(
@@ -71,7 +71,7 @@ UPB_API const upb_FileDef* upb_DefPool_AddFile(
     upb_DefPool* s, const UPB_DESC(FileDescriptorProto) * file_proto,
     upb_Status* status);
 
-const upb_ExtensionRegistry* upb_DefPool_ExtensionRegistry(
+UPB_API const upb_ExtensionRegistry* upb_DefPool_ExtensionRegistry(
     const upb_DefPool* s);
 
 const upb_FieldDef** upb_DefPool_GetAllExtensions(const upb_DefPool* s,