Implemented file-containing-extension request for Reflection Service

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).
+            """
+        )
+      )
+    }
+  }
 }