From f01d9e5601a82acbda9a2927fa9828fff5c0c6fb Mon Sep 17 00:00:00 2001
From: Jon Skeet <jonskeet@google.com>
Date: Wed, 20 Jan 2021 08:11:27 +0000
Subject: [PATCH] Add an extension registry option when parsing file
 descriptors

This is important when parsing descriptor sets that contain extensions.
(The alternative is to get the descriptor bytes again for individual
elements, e.g. message descriptors, and reparse them with the
appropriate extensions. It's really ugly.)
---
 .../Reflection/DescriptorsTest.cs             | 20 +++++++++++++++++++
 .../Reflection/FileDescriptor.cs              | 19 ++++++++++++++++--
 2 files changed, 37 insertions(+), 2 deletions(-)

diff --git a/csharp/src/Google.Protobuf.Test/Reflection/DescriptorsTest.cs b/csharp/src/Google.Protobuf.Test/Reflection/DescriptorsTest.cs
index ebb8394d2355..fab983d46351 100644
--- a/csharp/src/Google.Protobuf.Test/Reflection/DescriptorsTest.cs
+++ b/csharp/src/Google.Protobuf.Test/Reflection/DescriptorsTest.cs
@@ -35,7 +35,9 @@
 using ProtobufUnittest;
 using System;
 using System.Collections.Generic;
+using System.IO;
 using System.Linq;
+using UnitTest.Issues.TestProtos;
 
 namespace Google.Protobuf.Reflection
 {
@@ -70,6 +72,24 @@ public void FileDescriptor_BuildFromByteStrings()
             TestFileDescriptor(converted[2], converted[1], converted[0]);
         }
 
+        [Test]
+        public void FileDescriptor_BuildFromByteStrings_WithExtensionRegistry()
+        {
+            var extension = UnittestCustomOptionsProto3Extensions.MessageOpt1;
+
+            var byteStrings = new[]
+            {
+                DescriptorReflection.Descriptor.Proto.ToByteString(),
+                UnittestCustomOptionsProto3Reflection.Descriptor.Proto.ToByteString()
+            };
+            var registry = new ExtensionRegistry { extension };
+
+            var descriptor = FileDescriptor.BuildFromByteStrings(byteStrings, registry).Last();
+            var message = descriptor.MessageTypes.Single(t => t.Name == nameof(TestMessageWithCustomOptions));
+            var extensionValue = message.GetOptions().GetExtension(extension);
+            Assert.AreEqual(-56, extensionValue);
+        }
+
         private void TestFileDescriptor(FileDescriptor file, FileDescriptor importedFile, FileDescriptor importedPublicFile)
         {
             Assert.AreEqual("unittest_proto3.proto", file.Name);
diff --git a/csharp/src/Google.Protobuf/Reflection/FileDescriptor.cs b/csharp/src/Google.Protobuf/Reflection/FileDescriptor.cs
index 88e4a9de9621..724bb3a04dad 100644
--- a/csharp/src/Google.Protobuf/Reflection/FileDescriptor.cs
+++ b/csharp/src/Google.Protobuf/Reflection/FileDescriptor.cs
@@ -481,18 +481,21 @@ private static IEnumerable<Extension> GetAllDependedExtensionsFromMessage(Messag
         /// dependencies must come before the descriptor which depends on them. (If A depends on B, and B
         /// depends on C, then the descriptors must be presented in the order C, B, A.) This is compatible
         /// with the order in which protoc provides descriptors to plugins.</param>
+        /// <param name="registry">The extension registry to use when parsing, or null if no extensions are required.</param>
         /// <returns>The file descriptors corresponding to <paramref name="descriptorData"/>.</returns>
-        public static IReadOnlyList<FileDescriptor> BuildFromByteStrings(IEnumerable<ByteString> descriptorData)
+        public static IReadOnlyList<FileDescriptor> BuildFromByteStrings(IEnumerable<ByteString> descriptorData, ExtensionRegistry registry)
         {
             ProtoPreconditions.CheckNotNull(descriptorData, nameof(descriptorData));
 
+            var parser = FileDescriptorProto.Parser.WithExtensionRegistry(registry);
+
             // TODO: See if we can build a single DescriptorPool instead of building lots of them.
             // This will all behave correctly, but it's less efficient than we'd like.
             var descriptors = new List<FileDescriptor>();
             var descriptorsByName = new Dictionary<string, FileDescriptor>();
             foreach (var data in descriptorData)
             {
-                var proto = FileDescriptorProto.Parser.ParseFrom(data);
+                var proto = parser.ParseFrom(data);
                 var dependencies = new List<FileDescriptor>();
                 foreach (var dependencyName in proto.Dependency)
                 {
@@ -518,6 +521,18 @@ public static IReadOnlyList<FileDescriptor> BuildFromByteStrings(IEnumerable<Byt
             return new ReadOnlyCollection<FileDescriptor>(descriptors);
         }
 
+        /// <summary>
+        /// Converts the given descriptor binary data into FileDescriptor objects.
+        /// Note: reflection using the returned FileDescriptors is not currently supported.
+        /// </summary>
+        /// <param name="descriptorData">The binary file descriptor proto data. Must not be null, and any
+        /// dependencies must come before the descriptor which depends on them. (If A depends on B, and B
+        /// depends on C, then the descriptors must be presented in the order C, B, A.) This is compatible
+        /// with the order in which protoc provides descriptors to plugins.</param>
+        /// <returns>The file descriptors corresponding to <paramref name="descriptorData"/>.</returns>
+        public static IReadOnlyList<FileDescriptor> BuildFromByteStrings(IEnumerable<ByteString> descriptorData) =>
+            BuildFromByteStrings(descriptorData, null);
+
         /// <summary>
         /// Returns a <see cref="System.String" /> that represents this instance.
         /// </summary>