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 ebb8394d2..fab983d46 100644
--- a/csharp/src/Google.Protobuf.Test/Reflection/DescriptorsTest.cs
+++ b/csharp/src/Google.Protobuf.Test/Reflection/DescriptorsTest.cs
@@ -35,7 +35,9 @@ using NUnit.Framework;
 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 @@ namespace Google.Protobuf.Reflection
             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 88e4a9de9..724bb3a04 100644
--- a/csharp/src/Google.Protobuf/Reflection/FileDescriptor.cs
+++ b/csharp/src/Google.Protobuf/Reflection/FileDescriptor.cs
@@ -481,18 +481,21 @@ namespace Google.Protobuf.Reflection
         /// 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 @@ namespace Google.Protobuf.Reflection
             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>