Support all_extension_numbers_of_type reflection requests

Sources/GRPCReflectionService/Server/ReflectionService.swift

@@ -51,13 +51,21 @@ internal struct ReflectionServiceData: Sendable {
   internal var fileDescriptorDataByFilename: [String: FileDescriptorProtoData]
   internal var serviceNames: [String]
   internal var fileNameBySymbol: [String: String]
+
+  // Stores the file names for each extension identified by an ExtensionDescriptor object.
   private var fileNameByExtensionDescriptor: [ExtensionDescriptor: String]
+  // Stores the field numbers for each type that has extensions.
+  private var fieldNumbersByType: [String: [Int32]]
+  // Stores all the messageTypeNames from the protos.
+  private var messageTypeNames: [String]
 
   internal init(fileDescriptors: [Google_Protobuf_FileDescriptorProto]) throws {
     self.serviceNames = []
     self.fileDescriptorDataByFilename = [:]
     self.fileNameBySymbol = [:]
     self.fileNameByExtensionDescriptor = [:]
+    self.fieldNumbersByType = [:]
+    self.messageTypeNames = []
 
     for fileDescriptorProto in fileDescriptors {
       let serializedFileDescriptorProto: Data
@@ -92,10 +100,11 @@ internal struct ReflectionServiceData: Sendable {
         }
       }
 
-      // Populating the <extension descriptor, file name> dictionary.
+      // Populating the <extension descriptor, file name> dictionary and the <typeName, [FieldNumber]> one.
       for `extension` in fileDescriptorProto.extension {
+        let typeName = String(`extension`.extendee.drop(while: { $0 == "." }))
         let extensionDescriptor = ExtensionDescriptor(
-          extendeeTypeName: `extension`.extendee,
+          extendeeTypeName: typeName,
           fieldNumber: `extension`.number
         )
         let oldFileName = self.fileNameByExtensionDescriptor.updateValue(
@@ -112,7 +121,12 @@ internal struct ReflectionServiceData: Sendable {
               """
           )
         }
+        self.fieldNumbersByType[typeName, default: []].append(`extension`.number)
       }
+      // Populating messageTypeNames array.
+      self.messageTypeNames.append(
+        contentsOf: fileDescriptorProto.qualifiedMessageTypes
+      )
     }
   }
 
@@ -151,12 +165,21 @@ internal struct ReflectionServiceData: Sendable {
   }
 
   internal func nameOfFileContainingExtension(
-    named extendeeName: String,
+    extendeeName: String,
     fieldNumber number: Int32
   ) -> String? {
     let key = ExtensionDescriptor(extendeeTypeName: extendeeName, fieldNumber: number)
     return self.fileNameByExtensionDescriptor[key]
   }
+
+  // Returns nil if the type has no extensions or if it doesn't exist.
+  internal func extensionsFieldNumbersOfType(named typeName: String) -> [Int32]? {
+    return self.fieldNumbersByType[typeName]
+  }
+
+  internal func containsMessageType(name typeName: String) -> Bool {
+    return self.messageTypeNames.contains(typeName)
+  }
 }
 
 @available(macOS 10.15, iOS 13, tvOS 13, watchOS 6, *)
@@ -216,7 +239,7 @@ internal final class ReflectionServiceProvider: Reflection_ServerReflectionAsync
   ) throws -> Reflection_ServerReflectionResponse {
     guard
       let fileName = self.protoRegistry.nameOfFileContainingExtension(
-        named: extensionRequest.containingType,
+        extendeeName: extensionRequest.containingType,
         fieldNumber: extensionRequest.extensionNumber
       )
     else {
@@ -228,6 +251,31 @@ internal final class ReflectionServiceProvider: Reflection_ServerReflectionAsync
     return try self.findFileByFileName(fileName, request: request)
   }
 
+  internal func findExtensionsFieldNumbersOfType(
+    named typeName: String,
+    request: Reflection_ServerReflectionRequest
+  ) throws -> Reflection_ServerReflectionResponse {
+    var fieldNumbers = self.protoRegistry.extensionsFieldNumbersOfType(named: typeName)
+    if fieldNumbers == nil {
+      // Checks if the typeName is a valid Message type name, that has no extensions,
+      // or if it is invalid.
+      guard self.protoRegistry.containsMessageType(name: typeName) else {
+        throw GRPCStatus(
+          code: .notFound,
+          message: "The provided type doesn't have any extensions."
+        )
+      }
+      fieldNumbers = []
+    }
+    return Reflection_ServerReflectionResponse(
+      request: request,
+      extensionNumberResponse: .with {
+        $0.baseTypeName = typeName
+        $0.extensionNumber = fieldNumbers!
+      }
+    )
+  }
+
   internal func serverReflectionInfo(
     requestStream: GRPCAsyncRequestStream<Reflection_ServerReflectionRequest>,
     responseStream: GRPCAsyncResponseStreamWriter<Reflection_ServerReflectionResponse>,
@@ -260,6 +308,13 @@ internal final class ReflectionServiceProvider: Reflection_ServerReflectionAsync
         )
         try await responseStream.send(response)
 
+      case let .allExtensionNumbersOfType(typeName):
+        let response = try self.findExtensionsFieldNumbersOfType(
+          named: typeName,
+          request: request
+        )
+        try await responseStream.send(response)
+
       default:
         throw GRPCStatus(code: .unimplemented)
       }
@@ -289,6 +344,17 @@ extension Reflection_ServerReflectionResponse {
       $0.listServicesResponse = listServicesResponse
     }
   }
+
+  init(
+    request: Reflection_ServerReflectionRequest,
+    extensionNumberResponse: Reflection_ExtensionNumberResponse
+  ) {
+    self = .with {
+      $0.validHost = request.host
+      $0.originalRequest = request
+      $0.allExtensionNumbersResponse = extensionNumberResponse
+    }
+  }
 }
 
 extension Google_Protobuf_FileDescriptorProto {

Tests/GRPCTests/GRPCReflectionServiceTests/ReflectionServiceIntegrationTests.swift

@@ -184,7 +184,7 @@ final class ReflectionServiceIntegrationTests: GRPCTestCase {
       .with {
         $0.host = "127.0.0.1"
         $0.fileContainingExtension = .with {
-          $0.containingType = "inputMessage1"
+          $0.containingType = "packagebar1.inputMessage1"
           $0.extensionNumber = 2
         }
       }
@@ -212,10 +212,12 @@ final class ReflectionServiceIntegrationTests: GRPCTestCase {
       if fileDescriptorProto == fileToFind {
         receivedProtoContainingExtension += 1
         XCTAssert(
-          fileDescriptorProto.extension.map { $0.name }.contains("extensionInputMessage1"),
+          fileDescriptorProto.extension.map { $0.name }.contains(
+            "extension.packagebar1.inputMessage1-2"
+          ),
           """
           The response doesn't contain the serialized file descriptor proto \
-          containing the \"extensionInputMessage1\" extension.
+          containing the \"extensioninputMessage1-2\" extension.
           """
         )
       } else {
@@ -224,7 +226,7 @@ final class ReflectionServiceIntegrationTests: GRPCTestCase {
           dependentProtos.contains(fileDescriptorProto),
           """
           The \(fileDescriptorProto.name) is not a dependency of the \
-          proto file containing the \"extensionInputMessage1\" symbol.
+          proto file containing the \"extensioninputMessage1-2\" extension.
           """
         )
       }
@@ -236,4 +238,25 @@ final class ReflectionServiceIntegrationTests: GRPCTestCase {
     )
     XCTAssertEqual(dependenciesCount, 3)
   }
+
+  func testAllExtensionNumbersOfType() 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.allExtensionNumbersOfType = "packagebar2.inputMessage2"
+      }
+    )
+
+    serviceReflectionInfo.requestStream.finish()
+    var iterator = serviceReflectionInfo.responseStream.makeAsyncIterator()
+    guard let message = try await iterator.next() else {
+      return XCTFail("Could not get a response message.")
+    }
+    XCTAssertEqual(message.allExtensionNumbersResponse.baseTypeName, "packagebar2.inputMessage2")
+    XCTAssertEqual(message.allExtensionNumbersResponse.extensionNumber, [1, 2, 3, 4, 5])
+  }
 }

Tests/GRPCTests/GRPCReflectionServiceTests/ReflectionServiceUnitTests.swift

@@ -337,29 +337,9 @@ final class ReflectionServiceUnitTests: GRPCTestCase {
     let registry = try ReflectionServiceData(fileDescriptors: protos)
     for proto in protos {
       for `extension` in proto.extension {
+        let typeName = String(`extension`.extendee.drop(while: { $0 == "." }))
         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,
+          extendeeName: typeName,
           fieldNumber: `extension`.number
         )
         XCTAssertEqual(registryFileName, proto.name)
@@ -371,7 +351,7 @@ final class ReflectionServiceUnitTests: GRPCTestCase {
     let protos = makeProtosWithDependencies()
     let registry = try ReflectionServiceData(fileDescriptors: protos)
     let registryFileName = registry.nameOfFileContainingExtension(
-      named: "InvalidType",
+      extendeeName: "InvalidType",
       fieldNumber: 2
     )
     XCTAssertNil(registryFileName)
@@ -381,8 +361,8 @@ final class ReflectionServiceUnitTests: GRPCTestCase {
     let protos = makeProtosWithDependencies()
     let registry = try ReflectionServiceData(fileDescriptors: protos)
     let registryFileName = registry.nameOfFileContainingExtension(
-      named: protos[0].extension[0].extendee,
-      fieldNumber: 4
+      extendeeName: protos[0].extension[0].extendee,
+      fieldNumber: 9
     )
     XCTAssertNil(registryFileName)
   }
@@ -391,7 +371,7 @@ final class ReflectionServiceUnitTests: GRPCTestCase {
     var protos = makeProtosWithDependencies()
     protos[0].extension.append(
       .with {
-        $0.extendee = "inputMessage1"
+        $0.extendee = ".packagebar1.inputMessage1"
         $0.number = 2
       }
     )
@@ -404,11 +384,68 @@ final class ReflectionServiceUnitTests: GRPCTestCase {
           code: .alreadyExists,
           message:
             """
-            The extension of the inputMessage1 type with the field number equal to \
+            The extension of the packagebar1.inputMessage1 type with the field number equal to \
             2 from \(protos[0].name) already exists in \(protos[0].name).
             """
         )
       )
     }
   }
+
+  func testExtensionsFieldNumbersOfType() throws {
+    var protos = makeProtosWithDependencies()
+    protos[0].extension.append(
+      .with {
+        $0.extendee = ".packagebar1.inputMessage1"
+        $0.number = 120
+      }
+    )
+    let registry = try ReflectionServiceData(fileDescriptors: protos)
+    let extensionNumbers = registry.extensionsFieldNumbersOfType(named: "packagebar1.inputMessage1")
+    XCTAssertEqual(extensionNumbers, [1, 2, 3, 4, 5, 120])
+  }
+
+  func testExtensionsFieldNumbersOfTypeNoExtensionsType() throws {
+    var protos = makeProtosWithDependencies()
+    protos[0].messageType.append(
+      Google_Protobuf_DescriptorProto.with {
+        $0.name = "noExtensionMessage"
+        $0.field = [
+          Google_Protobuf_FieldDescriptorProto.with {
+            $0.name = "noExtensionField"
+            $0.type = .bool
+          }
+        ]
+      }
+    )
+    let registry = try ReflectionServiceData(fileDescriptors: protos)
+    let extensionNumbers = registry.extensionsFieldNumbersOfType(
+      named: "packagebar1.noExtensionMessage"
+    )
+    XCTAssertNil(extensionNumbers)
+    XCTAssertTrue(registry.containsMessageType(name: "packagebar1.noExtensionMessage"))
+  }
+
+  func testExtensionsFieldNumbersOfTypeInvalidTypeName() throws {
+    let protos = makeProtosWithDependencies()
+    let registry = try ReflectionServiceData(fileDescriptors: protos)
+    let extensionNumbers = registry.extensionsFieldNumbersOfType(
+      named: "packagebar1.invalidTypeMessage"
+    )
+    XCTAssertNil(extensionNumbers)
+    XCTAssertFalse(registry.containsMessageType(name: "packagebar1.noExtensionMessage"))
+  }
+
+  func testExtensionsFieldNumbersOfTypeExtensionsInDifferentProtoFiles() throws {
+    var protos = makeProtosWithDependencies()
+    protos[2].extension.append(
+      .with {
+        $0.extendee = ".packagebar1.inputMessage1"
+        $0.number = 130
+      }
+    )
+    let registry = try ReflectionServiceData(fileDescriptors: protos)
+    let extensionNumbers = registry.extensionsFieldNumbersOfType(named: "packagebar1.inputMessage1")
+    XCTAssertEqual(extensionNumbers, [1, 2, 3, 4, 5, 130])
+  }
 }

Tests/GRPCTests/GRPCReflectionServiceTests/Utils.swift

@@ -18,6 +18,23 @@ import Foundation
 import GRPC
 import SwiftProtobuf
 
+internal func makeExtensions(
+  forType typeName: String,
+  number: Int
+) -> [Google_Protobuf_FieldDescriptorProto] {
+  var extensions: [Google_Protobuf_FieldDescriptorProto] = []
+  for id in 1 ... number {
+    extensions.append(
+      Google_Protobuf_FieldDescriptorProto.with {
+        $0.name = "extension" + typeName + "-" + String(id)
+        $0.extendee = typeName
+        $0.number = Int32(id)
+      }
+    )
+  }
+  return extensions
+}
+
 internal func generateFileDescriptorProto(
   fileName name: String,
   suffix: String
@@ -32,11 +49,11 @@ internal func generateFileDescriptorProto(
     ]
   }
 
-  let inputMessageExtension = Google_Protobuf_FieldDescriptorProto.with {
-    $0.name = "extensionInputMessage" + suffix
-    $0.extendee = "inputMessage" + suffix
-    $0.number = 2
-  }
+  let packageName = "package" + name + suffix
+  let inputMessageExtensions = makeExtensions(
+    forType: "." + packageName + "." + "inputMessage" + suffix,
+    number: 5
+  )
 
   let outputMessage = Google_Protobuf_DescriptorProto.with {
     $0.name = "outputMessage" + suffix
@@ -77,7 +94,7 @@ internal func generateFileDescriptorProto(
     $0.package = "package" + name + suffix
     $0.messageType = [inputMessage, outputMessage]
     $0.enumType = [enumType]
-    $0.extension = [inputMessageExtension]
+    $0.extension = inputMessageExtensions
   }
 
   return fileDescriptorProto
@@ -109,6 +126,7 @@ internal func makeProtosWithComplexDependencies() -> [Google_Protobuf_FileDescri
       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)