Implemented file-containing-extension request for Reflection Service (#1677)

Motivation:

The file-containing-extension request enables users to get the file descriptor protos of
the proto file containing the extension they are looking for and its transitive dependencies.

Modifications:

- Created a new struct (ExtensionDescriptor) to represent the type that is extended and the field number of each extension that
exists inside the protos passed to the Reflection Service.
- Added a <ExtensionDescriptor, FileName> dictionary inside the ReflectionServiceData registry to store the extensions, avoid
duplicates and retrieve nicely the file name of the proto that contains the requested extension. The file name is then used to
get the serialized file descriptor protos of the proto containing the extension and its transitive dependencies.
- Added integration and unit tests.

Result:

The Reflection Service can now be used to get the serialized file descriptor protos of the proto containing the provided
extension and its transitive dependencies.

Sources/GRPCReflectionService/Server/ReflectionService.swift

@@ -43,15 +43,21 @@ internal struct ReflectionServiceData: Sendable {
     internal var serializedFileDescriptorProto: Data
     internal var dependencyFileNames: [String]
   }
+  private struct ExtensionDescriptor: Sendable, Hashable {
+    internal let extendeeTypeName: String
+    internal let fieldNumber: Int32
+  }
 
   internal var fileDescriptorDataByFilename: [String: FileDescriptorProtoData]
   internal var serviceNames: [String]
   internal var fileNameBySymbol: [String: String]
+  private var fileNameByExtensionDescriptor: [ExtensionDescriptor: String]
 
   internal init(fileDescriptors: [Google_Protobuf_FileDescriptorProto]) throws {
     self.serviceNames = []
     self.fileDescriptorDataByFilename = [:]
     self.fileNameBySymbol = [:]
+    self.fileNameByExtensionDescriptor = [:]
 
     for fileDescriptorProto in fileDescriptors {
       let serializedFileDescriptorProto: Data
@@ -70,6 +76,8 @@ internal struct ReflectionServiceData: Sendable {
       )
       self.fileDescriptorDataByFilename[fileDescriptorProto.name] = protoData
       self.serviceNames.append(contentsOf: fileDescriptorProto.service.map { $0.name })
+
+      // Populating the <symbol, file name> dictionary.
       for qualifiedSybolName in fileDescriptorProto.qualifiedSymbolNames {
         let oldValue = self.fileNameBySymbol.updateValue(
           fileDescriptorProto.name,
@@ -83,6 +91,28 @@ internal struct ReflectionServiceData: Sendable {
           )
         }
       }
+
+      // Populating the <extension descriptor, file name> dictionary.
+      for `extension` in fileDescriptorProto.extension {
+        let extensionDescriptor = ExtensionDescriptor(
+          extendeeTypeName: `extension`.extendee,
+          fieldNumber: `extension`.number
+        )
+        let oldFileName = self.fileNameByExtensionDescriptor.updateValue(
+          fileDescriptorProto.name,
+          forKey: extensionDescriptor
+        )
+        if let oldFileName = oldFileName {
+          throw GRPCStatus(
+            code: .alreadyExists,
+            message:
+              """
+              The extension of the \(extensionDescriptor.extendeeTypeName) type with the field number equal to \
+              \(extensionDescriptor.fieldNumber) from \(fileDescriptorProto.name) already exists in \(oldFileName).
+              """
+          )
+        }
+      }
     }
   }
 
@@ -119,6 +149,14 @@ internal struct ReflectionServiceData: Sendable {
   internal func nameOfFileContainingSymbol(named symbolName: String) -> String? {
     return self.fileNameBySymbol[symbolName]
   }
+
+  internal func nameOfFileContainingExtension(
+    named extendeeName: String,
+    fieldNumber number: Int32
+  ) -> String? {
+    let key = ExtensionDescriptor(extendeeTypeName: extendeeName, fieldNumber: number)
+    return self.fileNameByExtensionDescriptor[key]
+  }
 }
 
 @available(macOS 10.15, iOS 13, tvOS 13, watchOS 6, *)
@@ -172,6 +210,24 @@ internal final class ReflectionServiceProvider: Reflection_ServerReflectionAsync
     return try self.findFileByFileName(fileName, request: request)
   }
 
+  internal func findFileByExtension(
+    extensionRequest: Reflection_ExtensionRequest,
+    request: Reflection_ServerReflectionRequest
+  ) throws -> Reflection_ServerReflectionResponse {
+    guard
+      let fileName = self.protoRegistry.nameOfFileContainingExtension(
+        named: extensionRequest.containingType,
+        fieldNumber: extensionRequest.extensionNumber
+      )
+    else {
+      throw GRPCStatus(
+        code: .notFound,
+        message: "The provided extension could not be found."
+      )
+    }
+    return try self.findFileByFileName(fileName, request: request)
+  }
+
   internal func serverReflectionInfo(
     requestStream: GRPCAsyncRequestStream<Reflection_ServerReflectionRequest>,
     responseStream: GRPCAsyncResponseStreamWriter<Reflection_ServerReflectionResponse>,
@@ -197,6 +253,13 @@ internal final class ReflectionServiceProvider: Reflection_ServerReflectionAsync
         )
         try await responseStream.send(response)
 
+      case let .fileContainingExtension(extensionRequest):
+        let response = try self.findFileByExtension(
+          extensionRequest: extensionRequest,
+          request: request
+        )
+        try await responseStream.send(response)
+
       default:
         throw GRPCStatus(code: .unimplemented)
       }

Tests/GRPCTests/GRPCReflectionServiceTests/ReflectionServiceIntegrationTests.swift

@@ -29,7 +29,7 @@ final class ReflectionServiceIntegrationTests: GRPCTestCase {
   private let protos: [Google_Protobuf_FileDescriptorProto] = makeProtosWithDependencies()
   private let independentProto: Google_Protobuf_FileDescriptorProto = generateFileDescriptorProto(
     fileName: "independentBar",
-    suffix: 5
+    suffix: "5"
   )
 
   private func setUpServerAndChannel() throws {
@@ -174,4 +174,66 @@ final class ReflectionServiceIntegrationTests: GRPCTestCase {
       }
     }
   }
+
+  func testFileByExtension() async throws {
+    try self.setUpServerAndChannel()
+    let client = Reflection_ServerReflectionAsyncClient(channel: self.channel!)
+    let serviceReflectionInfo = client.makeServerReflectionInfoCall()
+
+    try await serviceReflectionInfo.requestStream.send(
+      .with {
+        $0.host = "127.0.0.1"
+        $0.fileContainingExtension = .with {
+          $0.containingType = "inputMessage1"
+          $0.extensionNumber = 2
+        }
+      }
+    )
+
+    serviceReflectionInfo.requestStream.finish()
+    var iterator = serviceReflectionInfo.responseStream.makeAsyncIterator()
+    guard let message = try await iterator.next() else {
+      return XCTFail("Could not get a response message.")
+    }
+    let receivedData: [Google_Protobuf_FileDescriptorProto]
+    do {
+      receivedData = try message.fileDescriptorResponse.fileDescriptorProto.map {
+        try Google_Protobuf_FileDescriptorProto(serializedData: $0)
+      }
+    } catch {
+      return XCTFail("Could not serialize data received as a message.")
+    }
+
+    let fileToFind = self.protos[0]
+    let dependentProtos = self.protos[1...]
+    var receivedProtoContainingExtension = 0
+    var dependenciesCount = 0
+    for fileDescriptorProto in receivedData {
+      if fileDescriptorProto == fileToFind {
+        receivedProtoContainingExtension += 1
+        XCTAssert(
+          fileDescriptorProto.extension.map { $0.name }.contains("extensionInputMessage1"),
+          """
+          The response doesn't contain the serialized file descriptor proto \
+          containing the \"extensionInputMessage1\" extension.
+          """
+        )
+      } else {
+        dependenciesCount += 1
+        XCTAssert(
+          dependentProtos.contains(fileDescriptorProto),
+          """
+          The \(fileDescriptorProto.name) is not a dependency of the \
+          proto file containing the \"extensionInputMessage1\" symbol.
+          """
+        )
+      }
+    }
+    XCTAssertEqual(
+      receivedProtoContainingExtension,
+      1,
+      "The file descriptor proto of the proto containing the extension was not received."
+    )
+    XCTAssertEqual(dependenciesCount, 3)
+  }
 }

Tests/GRPCTests/GRPCReflectionServiceTests/ReflectionServiceUnitTests.swift

@@ -62,7 +62,7 @@ final class ReflectionServiceUnitTests: GRPCTestCase {
     XCTAssertEqual(registryServices, servicesNames)
   }
 
-  /// Testing the fileNameBySymbol array of the ReflectionServiceData object.
+  /// Testing the fileNameBySymbol dictionary of the ReflectionServiceData object.
   func testFileNameBySymbol() throws {
     let protos = makeProtosWithDependencies()
     let registry = try ReflectionServiceData(fileDescriptors: protos)
@@ -85,7 +85,7 @@ final class ReflectionServiceUnitTests: GRPCTestCase {
     var protos = makeProtosWithDependencies()
     protos[1].messageType.append(
       Google_Protobuf_DescriptorProto.with {
-        $0.name = "inputMessage"
+        $0.name = "inputMessage2"
         $0.field = [
           Google_Protobuf_FieldDescriptorProto.with {
             $0.name = "inputField"
@@ -104,7 +104,7 @@ final class ReflectionServiceUnitTests: GRPCTestCase {
           code: .alreadyExists,
           message:
             """
-            The packagebar2.inputMessage symbol from bar2.proto \
+            The packagebar2.inputMessage2 symbol from bar2.proto \
             already exists in bar2.proto.
             """
         )
@@ -124,7 +124,7 @@ final class ReflectionServiceUnitTests: GRPCTestCase {
   func testNameOfFileContainingSymbolMessage() throws {
     let protos = makeProtosWithDependencies()
     let registry = try ReflectionServiceData(fileDescriptors: protos)
-    let fileName = registry.nameOfFileContainingSymbol(named: "packagebar1.inputMessage")
+    let fileName = registry.nameOfFileContainingSymbol(named: "packagebar1.inputMessage1")
     XCTAssertEqual(fileName, "bar1.proto")
   }
 
@@ -148,7 +148,7 @@ final class ReflectionServiceUnitTests: GRPCTestCase {
     let protos = makeProtosWithDependencies()
     let registry = try ReflectionServiceData(fileDescriptors: protos)
     let fileName = registry.nameOfFileContainingSymbol(named: "packagebar2.enumType3")
-    XCTAssertEqual(fileName, nil)
+    XCTAssertNil(fileName)
   }
 
   // Testing the serializedFileDescriptorProto method in different cases.
@@ -329,4 +329,86 @@ final class ReflectionServiceUnitTests: GRPCTestCase {
       )
     }
   }
+
+  // Testing the nameOfFileContainingExtension() method.
+
+  func testNameOfFileContainingExtensions() throws {
+    let protos = makeProtosWithDependencies()
+    let registry = try ReflectionServiceData(fileDescriptors: protos)
+    for proto in protos {
+      for `extension` in proto.extension {
+        let registryFileName = registry.nameOfFileContainingExtension(
+          named: `extension`.extendee,
+          fieldNumber: `extension`.number
+        )
+        XCTAssertEqual(registryFileName, proto.name)
+      }
+    }
+  }
+
+  func testNameOfFileContainingExtensionsSameTypeExtensionsDifferentNumbers() throws {
+    var protos = makeProtosWithDependencies()
+    protos[0].extension.append(
+      .with {
+        $0.extendee = "inputMessage1"
+        $0.number = 3
+      }
+    )
+    let registry = try ReflectionServiceData(fileDescriptors: protos)
+
+    for proto in protos {
+      for `extension` in proto.extension {
+        let registryFileName = registry.nameOfFileContainingExtension(
+          named: `extension`.extendee,
+          fieldNumber: `extension`.number
+        )
+        XCTAssertEqual(registryFileName, proto.name)
+      }
+    }
+  }
+
+  func testNameOfFileContainingExtensionsInvalidTypeName() throws {
+    let protos = makeProtosWithDependencies()
+    let registry = try ReflectionServiceData(fileDescriptors: protos)
+    let registryFileName = registry.nameOfFileContainingExtension(
+      named: "InvalidType",
+      fieldNumber: 2
+    )
+    XCTAssertNil(registryFileName)
+  }
+
+  func testNameOfFileContainingExtensionsInvalidFieldNumber() throws {
+    let protos = makeProtosWithDependencies()
+    let registry = try ReflectionServiceData(fileDescriptors: protos)
+    let registryFileName = registry.nameOfFileContainingExtension(
+      named: protos[0].extension[0].extendee,
+      fieldNumber: 4
+    )
+    XCTAssertNil(registryFileName)
+  }
+
+  func testNameOfFileContainingExtensionsDuplicatedExtensions() throws {
+    var protos = makeProtosWithDependencies()
+    protos[0].extension.append(
+      .with {
+        $0.extendee = "inputMessage1"
+        $0.number = 2
+      }
+    )
+    XCTAssertThrowsError(
+      try ReflectionServiceData(fileDescriptors: protos)
+    ) { error in
+      XCTAssertEqual(
+        error as? GRPCStatus,
+        GRPCStatus(
+          code: .alreadyExists,
+          message:
+            """
+            The extension of the inputMessage1 type with the field number equal to \
+            2 from \(protos[0].name) already exists in \(protos[0].name).
+            """
+        )
+      )
+    }
+  }
 }

Tests/GRPCTests/GRPCReflectionServiceTests/Utils.swift

@@ -20,10 +20,10 @@ import SwiftProtobuf
 
 internal func generateFileDescriptorProto(
   fileName name: String,
-  suffix id: Int
+  suffix: String
 ) -> Google_Protobuf_FileDescriptorProto {
   let inputMessage = Google_Protobuf_DescriptorProto.with {
-    $0.name = "inputMessage"
+    $0.name = "inputMessage" + suffix
     $0.field = [
       Google_Protobuf_FieldDescriptorProto.with {
         $0.name = "inputField"
@@ -32,8 +32,14 @@ internal func generateFileDescriptorProto(
     ]
   }
 
+  let inputMessageExtension = Google_Protobuf_FieldDescriptorProto.with {
+    $0.name = "extensionInputMessage" + suffix
+    $0.extendee = "inputMessage" + suffix
+    $0.number = 2
+  }
+
   let outputMessage = Google_Protobuf_DescriptorProto.with {
-    $0.name = "outputMessage"
+    $0.name = "outputMessage" + suffix
     $0.field = [
       Google_Protobuf_FieldDescriptorProto.with {
         $0.name = "outputField"
@@ -43,7 +49,7 @@ internal func generateFileDescriptorProto(
   }
 
   let enumType = Google_Protobuf_EnumDescriptorProto.with {
-    $0.name = "enumType" + String(id)
+    $0.name = "enumType" + suffix
     $0.value = [
       Google_Protobuf_EnumValueDescriptorProto.with {
         $0.name = "value1"
@@ -55,22 +61,23 @@ internal func generateFileDescriptorProto(
   }
 
   let method = Google_Protobuf_MethodDescriptorProto.with {
-    $0.name = "testMethod" + String(id)
+    $0.name = "testMethod" + suffix
     $0.inputType = inputMessage.name
     $0.outputType = outputMessage.name
   }
 
   let serviceDescriptor = Google_Protobuf_ServiceDescriptorProto.with {
     $0.method = [method]
-    $0.name = "service" + String(id)
+    $0.name = "service" + suffix
   }
 
   let fileDescriptorProto = Google_Protobuf_FileDescriptorProto.with {
     $0.service = [serviceDescriptor]
-    $0.name = name + String(id) + ".proto"
-    $0.package = "package" + name + String(id)
+    $0.name = name + suffix + ".proto"
+    $0.package = "package" + name + suffix
     $0.messageType = [inputMessage, outputMessage]
     $0.enumType = [enumType]
+    $0.extension = [inputMessageExtension]
   }
 
   return fileDescriptorProto
@@ -80,7 +87,7 @@ internal func generateFileDescriptorProto(
 internal func makeProtosWithDependencies() -> [Google_Protobuf_FileDescriptorProto] {
   var fileDependencies: [Google_Protobuf_FileDescriptorProto] = []
   for id in 1 ... 4 {
-    let fileDescriptorProto = generateFileDescriptorProto(fileName: "bar", suffix: id)
+    let fileDescriptorProto = generateFileDescriptorProto(fileName: "bar", suffix: String(id))
     if id != 1 {
       // Dependency of the first dependency.
       fileDependencies[0].dependency.append(fileDescriptorProto.name)
@@ -92,10 +99,16 @@ internal func makeProtosWithDependencies() -> [Google_Protobuf_FileDescriptorPro
 
 internal func makeProtosWithComplexDependencies() -> [Google_Protobuf_FileDescriptorProto] {
   var protos: [Google_Protobuf_FileDescriptorProto] = []
-  protos.append(generateFileDescriptorProto(fileName: "foo", suffix: 0))
+  protos.append(generateFileDescriptorProto(fileName: "foo", suffix: "0"))
   for id in 1 ... 10 {
-    let fileDescriptorProtoA = generateFileDescriptorProto(fileName: "fooA", suffix: id)
-    let fileDescriptorProtoB = generateFileDescriptorProto(fileName: "fooB", suffix: id)
+    let fileDescriptorProtoA = generateFileDescriptorProto(
+      fileName: "fooA",
+      suffix: String(id) + "A"
+    )
+    let fileDescriptorProtoB = generateFileDescriptorProto(
+      fileName: "fooB",
+      suffix: String(id) + "B"
+    )
     let parent = protos.count > 1 ? protos.count - Int.random(in: 1 ..< 3) : protos.count - 1
     protos[parent].dependency.append(fileDescriptorProtoA.name)
     protos[parent].dependency.append(fileDescriptorProtoB.name)