Fixed a bug where FileDescriptor maps like message_types_by_name could return descriptors from other files.

haberman · copybara-github · commit c05b320d9c18 · 2023-12-01T15:12:41.000-08:00
Fixes: #13740 PiperOrigin-RevId: 587135054
diff --git a/python/descriptor.c b/python/descriptor.c
@@ -1248,26 +1248,34 @@ static const void* PyUpb_FileDescriptor_NestedLookup(
 
 static const void* PyUpb_FileDescriptor_LookupMessage(
     const upb_FileDef* filedef, const char* name) {
-  return PyUpb_FileDescriptor_NestedLookup(
+  const upb_MessageDef* m = PyUpb_FileDescriptor_NestedLookup(
       filedef, name, (void*)&upb_DefPool_FindMessageByName);
+  if (!m) return NULL;
+  return upb_MessageDef_File(m) == filedef ? m : NULL;
 }
 
 static const void* PyUpb_FileDescriptor_LookupEnum(const upb_FileDef* filedef,
                                                    const char* name) {
-  return PyUpb_FileDescriptor_NestedLookup(filedef, name,
-                                           (void*)&upb_DefPool_FindEnumByName);
+  const upb_EnumDef* e = PyUpb_FileDescriptor_NestedLookup(
+      filedef, name, (void*)&upb_DefPool_FindEnumByName);
+  if (!e) return NULL;
+  return upb_EnumDef_File(e) == filedef ? e : NULL;
 }
 
 static const void* PyUpb_FileDescriptor_LookupExtension(
     const upb_FileDef* filedef, const char* name) {
-  return PyUpb_FileDescriptor_NestedLookup(
+  const upb_FieldDef* f = PyUpb_FileDescriptor_NestedLookup(
       filedef, name, (void*)&upb_DefPool_FindExtensionByName);
+  if (!f) return NULL;
+  return upb_FieldDef_File(f) == filedef ? f : NULL;
 }
 
 static const void* PyUpb_FileDescriptor_LookupService(
     const upb_FileDef* filedef, const char* name) {
-  return PyUpb_FileDescriptor_NestedLookup(
+  const upb_ServiceDef* s = PyUpb_FileDescriptor_NestedLookup(
       filedef, name, (void*)&upb_DefPool_FindServiceByName);
+  if (!s) return NULL;
+  return upb_ServiceDef_File(s) == filedef ? s : NULL;
 }
 
 static PyObject* PyUpb_FileDescriptor_GetName(PyUpb_DescriptorBase* self,
diff --git a/python/google/protobuf/internal/descriptor_pool_test.py b/python/google/protobuf/internal/descriptor_pool_test.py
@@ -132,6 +132,51 @@ def testFindFileContainingSymbol(self):
     self.assertRaises(KeyError, self.pool.FindFileContainingSymbol,
                       'google.protobuf.python.internal.Factory1Message.none_field')
 
+  def testCrossFileMessageTypesByName(self):
+    self.assertIs(
+        descriptor_pool_test1_pb2.DescriptorPoolTest1.DESCRIPTOR,
+        descriptor_pool_test1_pb2.DESCRIPTOR.message_types_by_name[
+            'DescriptorPoolTest1'
+        ],
+    )
+    with self.assertRaises(KeyError):
+      descriptor_pool_test2_pb2.DESCRIPTOR.message_types_by_name[
+          'DescriptorPoolTest1'
+      ]
+
+  def testCrossFileEnumTypesByName(self):
+    self.assertIs(
+        descriptor_pool_test1_pb2.TopLevelEnumTest1.DESCRIPTOR,
+        descriptor_pool_test1_pb2.DESCRIPTOR.enum_types_by_name[
+            'TopLevelEnumTest1'
+        ],
+    )
+    with self.assertRaises(KeyError):
+      descriptor_pool_test2_pb2.DESCRIPTOR.enum_types_by_name[
+          'TopLevelEnumTest1'
+      ]
+
+  def testCrossFileExtensionsByName(self):
+    self.assertIs(
+        descriptor_pool_test1_pb2.top_level_extension_test1,
+        descriptor_pool_test1_pb2.DESCRIPTOR.extensions_by_name[
+            'top_level_extension_test1'
+        ],
+    )
+    with self.assertRaises(KeyError):
+      descriptor_pool_test2_pb2.DESCRIPTOR.extensions_by_name[
+          'top_level_extension_test1'
+      ]
+
+  def testCrossFileServicesByName(self):
+    descriptor_pool_test1_pb2.DESCRIPTOR.services_by_name[
+        'DescriporPoolTestService'
+    ],
+    with self.assertRaises(KeyError):
+      descriptor_pool_test2_pb2.DESCRIPTOR.services_by_name[
+          'DescriporPoolTestService'
+      ]
+
   def testFindFileContainingSymbolFailure(self):
     with self.assertRaises(KeyError):
       self.pool.FindFileContainingSymbol('Does not exist')
diff --git a/python/google/protobuf/internal/descriptor_pool_test1.proto b/python/google/protobuf/internal/descriptor_pool_test1.proto
@@ -9,6 +9,17 @@ syntax = "proto2";
 
 package google.protobuf.python.internal;
 
+enum TopLevelEnumTest1 {
+  TOP_LEVEL_ENUM_TEST_1_NONE = 0;
+  TOP_LEVEL_ENUM_TEST_1_ONE = 1;
+}
+
+extend DescriptorPoolTest1 {
+  optional TopLevelEnumTest1 top_level_extension_test1 = 1000;
+}
+
+service DescriporPoolTestService {}
+
 message DescriptorPoolTest1 {
   extensions 1000 to max;
 
@@ -63,7 +74,7 @@ message DescriptorPoolTest2 {
         LAMBDA = 11;
         MU = 12;
 
-        reserved -1;
+        reserved - 1;
         reserved 536870913;  // 0x20000001
       }
       optional NestedEnum nested_enum = 1 [default = MU];
