diff --git a/eval/public/structs/BUILD b/eval/public/structs/BUILD index 468867294..41c0d7047 100644 --- a/eval/public/structs/BUILD +++ b/eval/public/structs/BUILD @@ -30,10 +30,17 @@ cc_library( deps = [ ":cel_proto_wrap_util", ":proto_message_type_adapter", + ":trivial_legacy_type_info_internal", + "//common:value", "//eval/public:cel_value", "//eval/public:message_wrapper", "//internal:proto_time_encoding", - "@com_google_absl//absl/types:optional", + "@com_google_absl//absl/base:no_destructor", + "@com_google_absl//absl/base:nullability", + "@com_google_absl//absl/log:absl_check", + "@com_google_absl//absl/log:absl_log", + "@com_google_absl//absl/status", + "@com_google_absl//absl/status:statusor", "@com_google_protobuf//:duration_cc_proto", "@com_google_protobuf//:protobuf", "@com_google_protobuf//:timestamp_cc_proto", diff --git a/eval/public/structs/cel_proto_wrapper.cc b/eval/public/structs/cel_proto_wrapper.cc index 6fad6aee3..e220b64dd 100644 --- a/eval/public/structs/cel_proto_wrapper.cc +++ b/eval/public/structs/cel_proto_wrapper.cc @@ -14,11 +14,21 @@ #include "eval/public/structs/cel_proto_wrapper.h" -#include "absl/types/optional.h" +#include + +#include "absl/base/no_destructor.h" +#include "absl/base/nullability.h" +#include "absl/log/absl_check.h" +#include "absl/log/absl_log.h" +#include "absl/status/status.h" +#include "absl/status/statusor.h" +#include "common/legacy_value.h" +#include "common/value.h" #include "eval/public/cel_value.h" #include "eval/public/message_wrapper.h" #include "eval/public/structs/cel_proto_wrap_util.h" #include "eval/public/structs/proto_message_type_adapter.h" +#include "eval/public/structs/trivial_legacy_type_info_internal.h" #include "google/protobuf/arena.h" #include "google/protobuf/descriptor.h" #include "google/protobuf/message.h" @@ -27,9 +37,36 @@ namespace google::api::expr::runtime { namespace { +using ::cel::interop_internal::TrivialTypeInfo; using ::google::protobuf::Arena; using ::google::protobuf::Descriptor; +using ::google::protobuf::DescriptorPool; using ::google::protobuf::Message; +using ::google::protobuf::MessageFactory; + +// Returns the arena for the given message, or the fallback arena if the +// message does not have an arena. +// +// A global fallback arena is used to avoid allocation when the arena is not +// specified. This is effectively a memory leak, but won't trigger leak +// check analyzers. +// +// This emulates the old behavior of tolerating a nullptr arena without +// triggering a crash. +google::protobuf::Arena* GetArena(const Message* absl_nonnull message, + google::protobuf::Arena* absl_nullable arena) { + if (arena != nullptr) { + return arena; + } + if (message->GetArena() != nullptr) { + return message->GetArena(); + } + static absl::NoDestructor fallback_arena; + ABSL_LOG(WARNING) + << "CelValue: using fallback global arena for wrapping message: " + << message->GetTypeName(); + return fallback_arena.get(); +} } // namespace @@ -38,14 +75,50 @@ CelValue CelProtoWrapper::InternalWrapMessage(const Message* message) { MessageWrapper(message, &GetGenericProtoTypeInfoInstance())); } -// CreateMessage creates CelValue from google::protobuf::Message. -// As some of CEL basic types are subclassing google::protobuf::Message, -// this method contains type checking and downcasts. -CelValue CelProtoWrapper::CreateMessage(const Message* value, Arena* arena) { - return internal::UnwrapMessageToValue(value, &InternalWrapMessage, arena); +CelValue CelProtoWrapper::CreateMessage( + const Message* absl_nonnull value, + const google::protobuf::DescriptorPool* absl_nonnull pool, + MessageFactory* absl_nonnull factory, Arena* absl_nonnull arena) { + ABSL_DCHECK(value != nullptr); + if (value->GetDescriptor() == nullptr || value->GetReflection() == nullptr) { + // This only happens for custom google::protobuf::Message subclasses that CEL can't + // support. + return CelValue::CreateMessageWrapper( + MessageWrapper(value, TrivialTypeInfo::GetInstance())); + } + + auto modern_value = + cel::Value::WrapMessageUnsafe(value, pool, factory, arena); + + absl::StatusOr cel_value = cel::LegacyValue(arena, modern_value); + if (!cel_value.ok()) { + // This only happens for custom google::protobuf::Message subclasses that CEL can't + // support. + auto* status = + google::protobuf::Arena::Create(arena, cel_value.status()); + return CelValue::CreateError(status); + } + return *cel_value; +} + +CelValue CelProtoWrapper::CreateMessage(const Message* absl_nullable value, + Arena* absl_nullable arena) { + if (value == nullptr) { + return CelValue::CreateNull(); + } + + if (value->GetDescriptor() == nullptr || value->GetReflection() == nullptr) { + // This only happens for custom messages subclasses that CEL can't support. + return CelValue::CreateMessageWrapper( + MessageWrapper(value, TrivialTypeInfo::GetInstance())); + } + const auto* pool = value->GetDescriptor()->file()->pool(); + auto* factory = value->GetReflection()->GetMessageFactory(); + arena = GetArena(value, arena); + return CreateMessage(value, pool, factory, arena); } -absl::optional CelProtoWrapper::MaybeWrapValue( +std::optional CelProtoWrapper::MaybeWrapValue( const Descriptor* descriptor, google::protobuf::MessageFactory* factory, const CelValue& value, Arena* arena) { const Message* msg = diff --git a/eval/public/structs/cel_proto_wrapper.h b/eval/public/structs/cel_proto_wrapper.h index 73942c253..8d334961c 100644 --- a/eval/public/structs/cel_proto_wrapper.h +++ b/eval/public/structs/cel_proto_wrapper.h @@ -1,9 +1,11 @@ #ifndef THIRD_PARTY_CEL_CPP_EVAL_PUBLIC_STRUCTS_CEL_PROTO_WRAPPER_H_ #define THIRD_PARTY_CEL_CPP_EVAL_PUBLIC_STRUCTS_CEL_PROTO_WRAPPER_H_ +#include + #include "google/protobuf/duration.pb.h" #include "google/protobuf/timestamp.pb.h" -#include "absl/types/optional.h" +#include "absl/base/nullability.h" #include "eval/public/cel_value.h" #include "internal/proto_time_encoding.h" #include "google/protobuf/arena.h" @@ -17,8 +19,19 @@ class CelProtoWrapper { // CreateMessage creates CelValue from google::protobuf::Message. // As some of CEL basic types are subclassing google::protobuf::Message, // this method contains type checking and downcasts. - static CelValue CreateMessage(const google::protobuf::Message* value, - google::protobuf::Arena* arena); + static CelValue CreateMessage(const google::protobuf::Message* absl_nonnull value, + const google::protobuf::DescriptorPool* absl_nonnull pool, + google::protobuf::MessageFactory* absl_nonnull factory, + google::protobuf::Arena* absl_nonnull arena); + + // Prefer using the overload that takes an explicit descriptor pool and + // message factory instead. This overload will use the ones associated with + // the value. + // + // For backward compatibility, nullptr message is allowed and will result in + // the CEL null_type value. + static CelValue CreateMessage(const google::protobuf::Message* absl_nullable value, + google::protobuf::Arena* absl_nullable arena); // Internal utility for creating a CelValue wrapping a user defined type. // Assumes that the message has been properly unpacked. @@ -43,7 +56,7 @@ class CelProtoWrapper { // message to native CelValue representation during a protobuf field read. // Just as CreateMessage should only be used when reading protobuf values, // MaybeWrapValue should only be used when assigning protobuf fields. - static absl::optional MaybeWrapValue( + static std::optional MaybeWrapValue( const google::protobuf::Descriptor* descriptor, google::protobuf::MessageFactory* factory, const CelValue& value, google::protobuf::Arena* arena); }; diff --git a/eval/public/structs/cel_proto_wrapper_test.cc b/eval/public/structs/cel_proto_wrapper_test.cc index 408e33284..597943f7a 100644 --- a/eval/public/structs/cel_proto_wrapper_test.cc +++ b/eval/public/structs/cel_proto_wrapper_test.cc @@ -104,7 +104,7 @@ class CelProtoWrapperTest : public ::testing::Test { T dyn_value; CelValue cel_dyn_value = - CelProtoWrapper::CreateMessage(ReflectedCopy(message).get(), arena()); + CelProtoWrapper::CreateMessage(ReflectedCopy(message), arena()); EXPECT_THAT(cel_dyn_value.type(), Eq(cel_value.type())); EXPECT_TRUE(cel_dyn_value.GetValue(&dyn_value)); EXPECT_THAT(value, Eq(dyn_value)); @@ -121,10 +121,9 @@ class CelProtoWrapperTest : public ::testing::Test { EXPECT_THAT(cel_value.MessageOrDie(), testutil::EqualsProto(*result)); } - std::unique_ptr ReflectedCopy( - const google::protobuf::Message& message) { - std::unique_ptr dynamic_value( - factory_.GetPrototype(message.GetDescriptor())->New()); + google::protobuf::Message* ReflectedCopy(const google::protobuf::Message& message) { + google::protobuf::Message* dynamic_value = + factory_.GetPrototype(message.GetDescriptor())->New(&arena_); dynamic_value->CopyFrom(message); return dynamic_value; } @@ -213,7 +212,7 @@ TEST_F(CelProtoWrapperTest, UnwrapDynamicValueNull) { value_msg.set_null_value(protobuf::NULL_VALUE); CelValue value = - CelProtoWrapper::CreateMessage(ReflectedCopy(value_msg).get(), arena()); + CelProtoWrapper::CreateMessage(ReflectedCopy(value_msg), arena()); EXPECT_TRUE(value.IsNull()); } @@ -314,8 +313,8 @@ TEST_F(CelProtoWrapperTest, UnwrapDynamicStruct) { const std::string kFieldBool = "field_bool"; (*struct_msg.mutable_fields())[kFieldInt].set_number_value(1.); (*struct_msg.mutable_fields())[kFieldBool].set_bool_value(true); - CelValue value = - CelProtoWrapper::CreateMessage(ReflectedCopy(struct_msg).get(), arena()); + auto reflected_copy = ReflectedCopy(struct_msg); + CelValue value = CelProtoWrapper::CreateMessage(reflected_copy, arena()); EXPECT_TRUE(value.IsMap()); const CelMap* cel_map = value.MapOrDie(); ASSERT_TRUE(cel_map != nullptr); @@ -355,7 +354,7 @@ TEST_F(CelProtoWrapperTest, UnwrapDynamicValueStruct) { .set_number_value(2); CelValue value = - CelProtoWrapper::CreateMessage(ReflectedCopy(value_msg).get(), arena()); + CelProtoWrapper::CreateMessage(ReflectedCopy(value_msg), arena()); EXPECT_TRUE(value.IsMap()); EXPECT_TRUE( (*value.MapOrDie())[CelValue::CreateString(&kField1)].has_value()); @@ -398,7 +397,7 @@ TEST_F(CelProtoWrapperTest, UnwrapDynamicValueListValue) { value_msg.mutable_list_value()->add_values()->set_number_value(2.); CelValue value = - CelProtoWrapper::CreateMessage(ReflectedCopy(value_msg).get(), arena()); + CelProtoWrapper::CreateMessage(ReflectedCopy(value_msg), arena()); EXPECT_TRUE(value.IsList()); EXPECT_THAT((*value.ListOrDie())[0].DoubleOrDie(), testing::DoubleEq(1)); EXPECT_THAT((*value.ListOrDie())[1].DoubleOrDie(), testing::DoubleEq(2)); @@ -426,6 +425,47 @@ TEST_F(CelProtoWrapperTest, UnwrapInvalidAny) { ASSERT_TRUE(CelProtoWrapper::CreateMessage(&any, arena()).IsError()); } +TEST_F(CelProtoWrapperTest, CreateMessageExplicitPoolAndFactory) { + TestMessage test_message; + test_message.set_string_value("test"); + + CelValue value = CelProtoWrapper::CreateMessage( + &test_message, google::protobuf::DescriptorPool::generated_pool(), + google::protobuf::MessageFactory::generated_factory(), arena()); + ASSERT_TRUE(value.IsMessage()); + EXPECT_THAT(value.MessageOrDie(), testutil::EqualsProto(test_message)); +} + +TEST_F(CelProtoWrapperTest, CreateMessageExplicitPoolAndFactoryUnpackAny) { + TestMessage test_message; + test_message.set_string_value("test"); + + Any any; + any.PackFrom(test_message); + + google::protobuf::DynamicMessageFactory factory( + google::protobuf::DescriptorPool::generated_pool()); + CelValue value = CelProtoWrapper::CreateMessage( + &any, google::protobuf::DescriptorPool::generated_pool(), &factory, arena()); + ASSERT_TRUE(value.IsMessage()); + EXPECT_THAT(value.MessageOrDie(), testutil::EqualsProto(test_message)); +} + +TEST_F(CelProtoWrapperTest, + CreateMessageExplicitPoolAndFactoryUnpackAnyNotFound) { + TestMessage test_message; + test_message.set_string_value("test"); + + Any any; + any.PackFrom(test_message); + + google::protobuf::DescriptorPool empty_pool; + google::protobuf::DynamicMessageFactory factory(&empty_pool); + CelValue value = + CelProtoWrapper::CreateMessage(&any, &empty_pool, &factory, arena()); + EXPECT_TRUE(value.IsError()); +} + // Test support of google.protobuf.Value wrappers in CelValue. TEST_F(CelProtoWrapperTest, UnwrapBoolWrapper) { bool value = true;