Raise "Proto Message has an Unknown Field" only when the extension field is known by the Python side.

PiperOrigin-RevId: 541546258

Author: pybind11_protobuf authors

pybind11_protobuf/BUILD

@@ -93,6 +93,7 @@ cc_library(
         "@com_google_absl//absl/meta:type_traits",
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/synchronization",
+        "@com_google_protobuf//:proto_api",
         "@com_google_protobuf//:protobuf",
     ],
 )

pybind11_protobuf/check_unknown_fields.cc

@@ -66,8 +66,9 @@ bool MessageMayContainExtensionsMemoized(const ::google::protobuf::Descriptor* d
 }
 
 struct HasUnknownFields {
-  HasUnknownFields(const ::google::protobuf::Descriptor* root_descriptor)
-      : root_descriptor(root_descriptor) {}
+  HasUnknownFields(const ::google::protobuf::python::PyProto_API* py_proto_api,
+                   const ::google::protobuf::Descriptor* root_descriptor)
+      : py_proto_api(py_proto_api), root_descriptor(root_descriptor) {}
 
   std::string FieldFQN() const { return absl::StrJoin(field_fqn_parts, "."); }
   std::string FieldFQNWithFieldNumber() const {
@@ -81,6 +82,7 @@ struct HasUnknownFields {
 
   std::string BuildErrorMessage() const;
 
+  const ::google::protobuf::python::PyProto_API* py_proto_api;
   const ::google::protobuf::Descriptor* root_descriptor = nullptr;
   const ::google::protobuf::Descriptor* unknown_field_parent_descriptor = nullptr;
   std::vector<std::string> field_fqn_parts;
@@ -97,9 +99,15 @@ bool HasUnknownFields::FindUnknownFieldsRecursive(
       reflection.GetUnknownFields(*sub_message);
   if (!unknown_field_set.empty()) {
     unknown_field_parent_descriptor = sub_message->GetDescriptor();
-    field_fqn_parts.resize(depth);
     unknown_field_number = unknown_field_set.field(0).number();
-    return true;
+
+    // Stop only if the extension is known by Python.
+    if (py_proto_api->GetDefaultDescriptorPool()->FindExtensionByNumber(
+            unknown_field_parent_descriptor,
+            unknown_field_number)) {
+      field_fqn_parts.resize(depth);
+      return true;
+    }
   }
 
   // If this message does not include submessages which allow extensions,
@@ -174,9 +182,10 @@ void AllowUnknownFieldsFor(absl::string_view top_message_descriptor_full_name,
 }
 
 std::optional<std::string> CheckAndBuildErrorMessageIfAny(
+    const ::google::protobuf::python::PyProto_API* py_proto_api,
     const ::google::protobuf::Message* message) {
   const auto* root_descriptor = message->GetDescriptor();
-  HasUnknownFields search{root_descriptor};
+  HasUnknownFields search{py_proto_api, root_descriptor};
   if (!search.FindUnknownFieldsRecursive(message, 0u)) {
     return std::nullopt;
   }

pybind11_protobuf/check_unknown_fields.h

@@ -4,6 +4,7 @@
 #include <optional>
 
 #include "google/protobuf/message.h"
+#include "python/google/protobuf/proto_api.h"
 #include "absl/strings/string_view.h"
 
 namespace pybind11_protobuf::check_unknown_fields {
@@ -12,6 +13,7 @@ void AllowUnknownFieldsFor(absl::string_view top_message_descriptor_full_name,
                            absl::string_view unknown_field_parent_message_fqn);
 
 std::optional<std::string> CheckAndBuildErrorMessageIfAny(
+    const ::google::protobuf::python::PyProto_API* py_proto_api,
     const ::google::protobuf::Message* top_message);
 
 }  // namespace pybind11_protobuf::check_unknown_fields

pybind11_protobuf/native_proto_caster.h

@@ -67,6 +67,7 @@ inline void AllowUnknownFieldsFor(
 }
 
 }  // namespace pybind11_protobuf
+
 namespace pybind11 {
 namespace detail {
 

pybind11_protobuf/proto_cast_util.cc

@@ -836,7 +836,8 @@ py::handle GenericProtoCast(Message* src, py::return_value_policy policy,
   }
 
   std::optional<std::string> emsg =
-      check_unknown_fields::CheckAndBuildErrorMessageIfAny(src);
+      check_unknown_fields::CheckAndBuildErrorMessageIfAny(
+          GlobalState::instance()->py_proto_api(), src);
   if (emsg) {
     throw py::value_error(*emsg);
   }

pybind11_protobuf/tests/extension_test.py

@@ -197,6 +197,12 @@ def test_reserialize_allow_unknown_outer(self):
           extension_in_other_file_pb2.AllowUnknownInnerExtension.hook].value
       self.assertEqual(97, b_inner_value)
 
+  def test_reserialize_allow_python_unknown_fields(self):
+    inner = get_allow_unknown_inner(63)
+    # Creates a message with only unknown fields.
+    a = extension_pb2.BaseMessage.FromString(inner.SerializeToString())
+    b = m.reserialize_base_message(a)
+    self.assertEqual(a.SerializeToString(), b.SerializeToString())
 
 if __name__ == '__main__':
   absltest.main()