Support all_extension_numbers_of_type reflection requests

Sources/GRPCReflectionService/Server/ReflectionService.swift

@@ -52,12 +52,16 @@ internal struct ReflectionServiceData: Sendable {
   internal var serviceNames: [String]
   internal var fileNameBySymbol: [String: String]
   private var fileNameByExtensionDescriptor: [ExtensionDescriptor: String]
+  private var fieldNumbersByType: [String: [Int32]]
+  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 +96,13 @@ 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 = ReflectionServiceData.extractTypeNameFrom(
+          fullyQualifiedName: `extension`.extendee
+        )
         let extensionDescriptor = ExtensionDescriptor(
-          extendeeTypeName: `extension`.extendee,
+          extendeeTypeName: typeName,
           fieldNumber: `extension`.number
         )
         let oldFileName = self.fileNameByExtensionDescriptor.updateValue(
@@ -112,8 +119,28 @@ internal struct ReflectionServiceData: Sendable {
               """
           )
         }
+        if self.fieldNumbersByType[typeName] == nil {
+          self.fieldNumbersByType[typeName] = []
+        }
+        let numberPosition = self.fieldNumbersByType[typeName]!.count
+        self.fieldNumbersByType[typeName]?.insert(
+          `extension`.number,
+          at: numberPosition
+        )
       }
+      // Populating messageTypeNames array.
+      self.messageTypeNames.append(
+        contentsOf: fileDescriptorProto.qualifiedMessageTypes
+      )
+    }
+  }
+
+  internal static func extractTypeNameFrom(fullyQualifiedName name: String) -> String {
+    var nameCopy = name
+    if nameCopy.first == "." {
+      nameCopy.removeFirst()
     }
+    return nameCopy
   }
 
   internal func serialisedFileDescriptorProtosForDependenciesOfFile(
@@ -151,12 +178,20 @@ 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]
   }
+
+  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 +251,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 +263,29 @@ 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 {
+      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 +318,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 +354,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

@@ -330,36 +330,28 @@ final class ReflectionServiceUnitTests: GRPCTestCase {
     }
   }
 
+  func testExtractTypeNameFrom() throws {
+    let initialExtendeeName = ".package.extendeeName"
+    XCTAssertEqual(
+      ReflectionServiceData.extractTypeNameFrom(
+        fullyQualifiedName: initialExtendeeName
+      ),
+      "package.extendeeName"
+    )
+  }
+
   // 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
+        let typeName = ReflectionServiceData.extractTypeNameFrom(
+          fullyQualifiedName: `extension`.extendee
         )
-        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 +363,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 +373,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 +383,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 +396,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])
+  }
 }