Add support for extensions in CRuby, JRuby, and FFI Ruby

ruby/ext/google/protobuf_c/defs.c

@@ -144,14 +144,20 @@ 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;
+
+  fielddef = upb_DefPool_FindExtensionByName(self->symtab, name_str);
+  if (fielddef) {
+    return get_fielddef_obj(_self, fielddef);
+  }
 
   msgdef = upb_DefPool_FindMessageByName(self->symtab, name_str);
   if (msgdef) {
@@ -575,7 +581,7 @@ typedef struct {
   VALUE descriptor_pool;  // Owns the upb_FieldDef.
 } FieldDescriptor;
 
-static VALUE cFieldDescriptor = Qnil;
+VALUE cFieldDescriptor = Qnil;
 
 static void FieldDescriptor_mark(void* _self) {
   FieldDescriptor* self = _self;
@@ -833,7 +839,7 @@ static VALUE FieldDescriptor_subtype(VALUE _self) {
  * Returns the value set for this field on the given message. Raises an
  * exception if message is of the wrong type.
  */
-static VALUE FieldDescriptor_get(VALUE _self, VALUE msg_rb) {
+VALUE FieldDescriptor_get(VALUE _self, VALUE msg_rb) {
   FieldDescriptor* self = ruby_to_FieldDescriptor(_self);
   const upb_MessageDef* m;
 

ruby/ext/google/protobuf_c/defs.h

@@ -79,4 +79,8 @@ extern VALUE generated_pool;
 // Call at startup to register all types in this module.
 void Defs_register(VALUE module);
 
+// FieldDescriptor
+VALUE FieldDescriptor_get(VALUE _self, VALUE msg_rb);
+VALUE cFieldDescriptor;
+
 #endif  // RUBY_PROTOBUF_DEFS_H_

ruby/ext/google/protobuf_c/message.c

@@ -977,9 +977,14 @@ 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);
 
-  upb_DecodeStatus status = upb_Decode(bytes, size, (upb_Message*)msg->msg,
+  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 +1308,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");

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 @@ def add_serialized_file(file_contents)
 
       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

ruby/lib/google/protobuf/ffi/message.rb

@@ -170,7 +170,15 @@ def self.decode(data, options = {})
 
             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

ruby/src/main/java/com/google/protobuf/jruby/RubyDescriptorPool.java

@@ -36,7 +36,9 @@
 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 IRubyObject allocate(Ruby runtime, RubyClass klazz) {
             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 IRubyObject build(ThreadContext context, Block block) {
    * 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,7 @@ public static IRubyObject generatedPool(ThreadContext context, IRubyObject recv)
   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 +153,8 @@ protected void registerFileDescriptor(
     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 +175,19 @@ private void registerDescriptor(ThreadContext context, Descriptor descriptor, St
       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) {
+    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);
+    symtab.put(name, des);
   }
 
   private void registerEnumDescriptor(
@@ -188,8 +206,10 @@ private FileDescriptor[] existingFileDescriptors() {
 
   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();
 }

ruby/src/main/java/com/google/protobuf/jruby/RubyFieldDescriptor.java

@@ -102,7 +102,9 @@ public IRubyObject getLabel(ThreadContext context) {
   public IRubyObject getName(ThreadContext context) {
     return this.name;
   }
-
+  protected void setName(IRubyObject name) {
+    this.name = name;
+  }
   /*
    * call-seq:
    *     FieldDescriptor.subtype => message_or_enum_descriptor

ruby/src/main/java/com/google/protobuf/jruby/RubyMessage.java

@@ -634,7 +634,7 @@ public static IRubyObject decode(ThreadContext context, IRubyObject recv, IRubyO
   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,

ruby/tests/basic.rb

@@ -729,6 +729,7 @@ def test_oneof_descriptor_options
       oneof_descriptor = descriptor.lookup_oneof("test_deprecated_message_oneof")
 
       assert_instance_of Google::Protobuf::OneofOptions, oneof_descriptor.options
+      assert_equal "Custom option value", Test_option.get(oneof_descriptor.options)
     end
 
     def test_options_deep_freeze
@@ -739,6 +740,25 @@ def test_options_deep_freeze
           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

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,17 @@ message TestMessage2 {
   optional int32 foo = 1;
 }
 
+extend google.protobuf.OneofOptions {
+  optional string test_option = 1000;
+}
+
 message TestDeprecatedMessage {
   option deprecated = true;
 
   optional int32 foo = 1 [deprecated = true];
 
   oneof test_deprecated_message_oneof {
+    option (test_option) = "Custom option value";
     string a = 2;
     int32 b = 3;
   }

src/google/protobuf/compiler/ruby/ruby_generator.cc

@@ -154,6 +154,18 @@ void GenerateEnumAssignment(absl::string_view prefix, const EnumDescriptor* en,
     "full_name", en->full_name());
 }
 
+void GenerateExtensionAssignment(absl::string_view prefix, const FieldDescriptor* f,
+                            io::Printer* printer) {
+  printer->Print(
+    "$prefix$$name$ = ",
+    "prefix", prefix,
+    "name", RubifyConstant(f->name()));
+  printer->Print(
+    "::Google::Protobuf::DescriptorPool.generated_pool."
+    "lookup(\"$full_name$\")\n",
+    "full_name", f->full_name());
+}
+
 int GeneratePackageModules(const FileDescriptor* file, io::Printer* printer) {
   int levels = 0;
   bool need_change_to_module = true;
@@ -309,11 +321,6 @@ bool GenerateFile(const FileDescriptor* file, io::Printer* printer,
     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);
@@ -323,6 +330,9 @@ bool GenerateFile(const FileDescriptor* file, io::Printer* printer,
   for (int i = 0; i < file->enum_type_count(); i++) {
     GenerateEnumAssignment("", file->enum_type(i), printer);
   }
+  for (int i = 0; i < file->extension_count(); i++) {
+    GenerateExtensionAssignment("", file->extension(i), printer);
+  }
   EndPackageModules(levels, printer);
 
   return true;

upb/reflection/def_pool.h

@@ -51,7 +51,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(
@@ -74,7 +74,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,