Support all_extension_numbers_of_type reflection requests (#1680)

Motivation:

The reflection service should provide the possibility for users to request the list with all the field numbers of
the extensions of a type.

Modifications:

- Implemented the dictionary that stores the arrays of integers representing the field numbers of the extensions
for each type that has extensions.
- Implemented the methods of the Reflection Service that create the specific response with the array of integers
representing the field numbers or an empty array for the case that the type doesn't have any extensions (but is a valid type).
- Implemented integration and Unit tests.

Result:

The Reflection Service will enable users to find all the extension field numbers for a specific type they requested.

Author: stefanadranca

Sources/GRPCReflectionService/Server/ReflectionService.swift

@@ -51,13 +51,18 @@ 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]]
 
   internal init(fileDescriptors: [Google_Protobuf_FileDescriptorProto]) throws {
     self.serviceNames = []
     self.fileDescriptorDataByFilename = [:]
     self.fileNameBySymbol = [:]
     self.fileNameByExtensionDescriptor = [:]
+    self.fieldNumbersByType = [:]
 
     for fileDescriptorProto in fileDescriptors {
       let serializedFileDescriptorProto: Data
@@ -92,10 +97,15 @@ internal struct ReflectionServiceData: Sendable {
         }
       }
 
-      // Populating the <extension descriptor, file name> dictionary.
+      for typeName in fileDescriptorProto.qualifiedMessageTypes {
+        self.fieldNumbersByType[typeName] = []
+      }
+
+      // 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,6 +122,7 @@ internal struct ReflectionServiceData: Sendable {
               """
           )
         }
+        self.fieldNumbersByType[typeName, default: []].append(`extension`.number)
       }
     }
   }
@@ -151,12 +162,23 @@ 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 an empty array if the type has no extensions.
+  internal func extensionsFieldNumbersOfType(named typeName: String) throws -> [Int32] {
+    guard let fieldNumbers = self.fieldNumbersByType[typeName] else {
+      throw GRPCStatus(
+        code: .invalidArgument,
+        message: "The provided type is invalid."
+      )
+    }
+    return fieldNumbers
+  }
 }
 
 @available(macOS 10.15, iOS 13, tvOS 13, watchOS 6, *)
@@ -216,7 +238,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 +250,20 @@ 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 {
+    let fieldNumbers = try self.protoRegistry.extensionsFieldNumbersOfType(named: typeName)
+    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 +296,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 +332,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,81 @@ 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).
             """
         )
       )
     }
   }
+
+  // Testing the extensionsFieldNumbersOfType() method.
+
+  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 = try 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 = try registry.extensionsFieldNumbersOfType(
+      named: "packagebar1.noExtensionMessage"
+    )
+    XCTAssertEqual(extensionNumbers, [])
+  }
+
+  func testExtensionsFieldNumbersOfTypeInvalidTypeName() throws {
+    let protos = makeProtosWithDependencies()
+    let registry = try ReflectionServiceData(fileDescriptors: protos)
+    XCTAssertThrowsError(
+      try registry.extensionsFieldNumbersOfType(
+        named: "packagebar1.invalidTypeMessage"
+      )
+    ) { error in
+      XCTAssertEqual(
+        error as? GRPCStatus,
+        GRPCStatus(
+          code: .invalidArgument,
+          message: "The provided type is invalid."
+        )
+      )
+    }
+  }
+
+  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 = try 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)