diff --git a/common/value.cc b/common/value.cc index bc8a3d7a0..8421dfeb4 100644 --- a/common/value.cc +++ b/common/value.cc @@ -935,6 +935,9 @@ void BoolRepeatedFieldAccessor( *result = BoolValue(reflection->GetRepeatedBool(*message, field, index)); } +// If `Unsafe` is `std::true_type`, elements are borrowed rather than copied to +// `arena`, as with `Value::WrapRepeatedFieldUnsafe`. +template void StringRepeatedFieldAccessor( int index, const google::protobuf::Message* absl_nonnull message, const google::protobuf::FieldDescriptor* absl_nonnull field, @@ -964,10 +967,15 @@ void StringRepeatedFieldAccessor( string.size() == scratch.size()) { *result = StringValue::From(std::move(scratch), arena); } else { - if (message->GetArena() == nullptr) { - *result = StringValue::From(string, arena); + if constexpr (Unsafe::value) { + *result = StringValue::WrapUnsafe(string); } else { - *result = StringValue::Wrap(string, message->GetArena()); + if (google::protobuf::Arena* message_arena = message->GetArena(); + message_arena != nullptr) { + *result = StringValue::Wrap(string, message_arena); + } else { + *result = StringValue::From(string, arena); + } } } }, @@ -975,9 +983,10 @@ void StringRepeatedFieldAccessor( *result = StringValue::From(std::move(cord), arena); }), well_known_types::AsVariant(well_known_types::GetRepeatedStringField( - *message, field, index, scratch))); + reflection, *message, field, index, scratch))); } +template void MessageRepeatedFieldAccessor( int index, const google::protobuf::Message* absl_nonnull message, const google::protobuf::FieldDescriptor* absl_nonnull field, @@ -999,11 +1008,18 @@ void MessageRepeatedFieldAccessor( ABSL_DCHECK_GE(index, 0); ABSL_DCHECK_LT(index, reflection->FieldSize(*message, field)); - *result = Value::WrapMessage( - &reflection->GetRepeatedMessage(*message, field, index), descriptor_pool, - message_factory, arena); + if constexpr (Unsafe::value) { + *result = Value::WrapMessageUnsafe( + &reflection->GetRepeatedMessage(*message, field, index), + descriptor_pool, message_factory, arena); + } else { + *result = Value::WrapMessage( + &reflection->GetRepeatedMessage(*message, field, index), + descriptor_pool, message_factory, arena); + } } +template void BytesRepeatedFieldAccessor( int index, const google::protobuf::Message* absl_nonnull message, const google::protobuf::FieldDescriptor* absl_nonnull field, @@ -1033,10 +1049,15 @@ void BytesRepeatedFieldAccessor( string.size() == scratch.size()) { *result = BytesValue::From(std::move(scratch), arena); } else { - if (message->GetArena() == nullptr) { - *result = BytesValue::From(string, arena); + if constexpr (Unsafe::value) { + *result = BytesValue::WrapUnsafe(string); } else { - *result = BytesValue::Wrap(string, message->GetArena()); + if (google::protobuf::Arena* message_arena = message->GetArena(); + message_arena != nullptr) { + *result = BytesValue::Wrap(string, message_arena); + } else { + *result = BytesValue::From(string, arena); + } } } }, @@ -1044,7 +1065,7 @@ void BytesRepeatedFieldAccessor( *result = BytesValue::From(std::move(cord), arena); }), well_known_types::AsVariant(well_known_types::GetRepeatedBytesField( - *message, field, index, scratch))); + reflection, *message, field, index, scratch))); } void EnumRepeatedFieldAccessor( @@ -1101,7 +1122,7 @@ void NullRepeatedFieldAccessor( } // namespace absl::StatusOr RepeatedFieldAccessorFor( - const google::protobuf::FieldDescriptor* absl_nonnull field) { + const google::protobuf::FieldDescriptor* absl_nonnull field, bool unsafe) { switch (field->type()) { case google::protobuf::FieldDescriptor::TYPE_DOUBLE: return &DoubleRepeatedFieldAccessor; @@ -1126,13 +1147,22 @@ absl::StatusOr RepeatedFieldAccessorFor( case google::protobuf::FieldDescriptor::TYPE_BOOL: return &BoolRepeatedFieldAccessor; case google::protobuf::FieldDescriptor::TYPE_STRING: - return &StringRepeatedFieldAccessor; + if (unsafe) { + return &StringRepeatedFieldAccessor; + } + return &StringRepeatedFieldAccessor; case google::protobuf::FieldDescriptor::TYPE_GROUP: ABSL_FALLTHROUGH_INTENDED; case google::protobuf::FieldDescriptor::TYPE_MESSAGE: - return &MessageRepeatedFieldAccessor; + if (unsafe) { + return &MessageRepeatedFieldAccessor; + } + return &MessageRepeatedFieldAccessor; case google::protobuf::FieldDescriptor::TYPE_BYTES: - return &BytesRepeatedFieldAccessor; + if (unsafe) { + return &BytesRepeatedFieldAccessor; + } + return &BytesRepeatedFieldAccessor; case google::protobuf::FieldDescriptor::TYPE_FIXED32: ABSL_FALLTHROUGH_INTENDED; case google::protobuf::FieldDescriptor::TYPE_UINT32: diff --git a/common/value.h b/common/value.h index fc7992c49..1757c5b03 100644 --- a/common/value.h +++ b/common/value.h @@ -2888,8 +2888,11 @@ using RepeatedFieldAccessor = google::protobuf::MessageFactory* absl_nonnull, google::protobuf::Arena* absl_nonnull, Value* absl_nonnull); +// If `unsafe` is true, message, string, and bytes elements are borrowed rather +// than copied, as with `Value::WrapRepeatedFieldUnsafe`. Caller is responsible +// for ensuring the resulting values do not outlive the input message. absl::StatusOr RepeatedFieldAccessorFor( - const google::protobuf::FieldDescriptor* absl_nonnull field); + const google::protobuf::FieldDescriptor* absl_nonnull field, bool unsafe = false); } // namespace common_internal diff --git a/common/values/parsed_repeated_field_value.cc b/common/values/parsed_repeated_field_value.cc index b716962cc..d9c79c271 100644 --- a/common/values/parsed_repeated_field_value.cc +++ b/common/values/parsed_repeated_field_value.cc @@ -249,7 +249,7 @@ absl::Status ParsedRepeatedFieldValue::Get( *result = IndexOutOfBoundsError(index, arena); return absl::OkStatus(); } - if (arena_ == nullptr) { + if (is_unsafe()) { *result = Value::WrapRepeatedFieldUnsafe(static_cast(index), message_, field_, descriptor_pool, message_factory, arena); @@ -274,7 +274,8 @@ absl::Status ParsedRepeatedFieldValue::ForEach( const int size = reflection->FieldSize(*message_, field_); if (size > 0) { CEL_ASSIGN_OR_RETURN(auto accessor, - common_internal::RepeatedFieldAccessorFor(field_)); + common_internal::RepeatedFieldAccessorFor( + field_, /*unsafe=*/is_unsafe())); Value scratch; for (int i = 0; i < size; ++i) { (*accessor)(i, message_, field_, reflection, descriptor_pool, @@ -377,8 +378,8 @@ ParsedRepeatedFieldValue::NewIterator() const { if (ABSL_PREDICT_FALSE(field_ == nullptr)) { return NewEmptyValueIterator(); } - CEL_ASSIGN_OR_RETURN(auto accessor, - common_internal::RepeatedFieldAccessorFor(field_)); + CEL_ASSIGN_OR_RETURN(auto accessor, common_internal::RepeatedFieldAccessorFor( + field_, /*unsafe=*/is_unsafe())); return std::make_unique(message_, field_, accessor); } @@ -397,7 +398,8 @@ absl::Status ParsedRepeatedFieldValue::Contains( const int size = reflection->FieldSize(*message_, field_); if (size > 0) { CEL_ASSIGN_OR_RETURN(auto accessor, - common_internal::RepeatedFieldAccessorFor(field_)); + common_internal::RepeatedFieldAccessorFor( + field_, /*unsafe=*/is_unsafe())); Value scratch; for (int i = 0; i < size; ++i) { (*accessor)(i, message_, field_, reflection, descriptor_pool, diff --git a/common/values/parsed_repeated_field_value.h b/common/values/parsed_repeated_field_value.h index e345c8ffa..45bcfaa98 100644 --- a/common/values/parsed_repeated_field_value.h +++ b/common/values/parsed_repeated_field_value.h @@ -176,7 +176,7 @@ class ParsedRepeatedFieldValue final ParsedRepeatedFieldValue(const google::protobuf::Message* absl_nonnull message, const google::protobuf::FieldDescriptor* absl_nonnull field) - : message_(message), field_(field), arena_(message->GetArena()) { + : message_(message), field_(field), arena_(nullptr) { ABSL_DCHECK(message != nullptr); ABSL_DCHECK(field != nullptr); ABSL_DCHECK(field_->is_repeated() && !field_->is_map()) @@ -193,10 +193,15 @@ class ParsedRepeatedFieldValue final return absl::OkStatus(); } + bool is_unsafe() const { return arena_ == nullptr; } + const google::protobuf::Reflection* absl_nonnull GetReflection() const; const google::protobuf::Message* absl_nullable message_ = nullptr; const google::protobuf::FieldDescriptor* absl_nullable field_ = nullptr; + // The arena attributed as owning this value. Null if the value is created by + // UnsafeParsedRepeatedFieldValue(), indicating an externally managed message + // whose elements should be borrowed rather than copied. google::protobuf::Arena* absl_nullable arena_ = nullptr; }; diff --git a/common/values/parsed_repeated_field_value_test.cc b/common/values/parsed_repeated_field_value_test.cc index b14e8753f..fd08558e3 100644 --- a/common/values/parsed_repeated_field_value_test.cc +++ b/common/values/parsed_repeated_field_value_test.cc @@ -14,6 +14,7 @@ #include #include +#include #include #include @@ -23,6 +24,7 @@ #include "absl/status/statusor.h" #include "absl/strings/string_view.h" #include "absl/time/time.h" +#include "absl/types/optional.h" #include "cel/expr/conformance/proto3/test_all_types.pb.h" #include "common/memory.h" #include "common/type.h" @@ -416,6 +418,110 @@ TEST_F(ParsedRepeatedFieldValueTest, NewIterator2) { IsOkAndHolds(Eq(std::nullopt))); } +TEST_F(ParsedRepeatedFieldValueTest, GetUnsafeDifferentArenaBorrowsMessages) { + google::protobuf::Arena source_arena; + auto* message = google::protobuf::Arena::Create(&source_arena); + message->add_repeated_nested_message()->set_bb(1); + ParsedRepeatedFieldValue value = UnsafeParsedRepeatedFieldValue( + message, TestAllTypesProto3::descriptor()->FindFieldByName( + "repeated_nested_message")); + ASSERT_OK_AND_ASSIGN(Value element, value.Get(0, descriptor_pool(), + message_factory(), arena())); + ASSERT_TRUE(element.IsParsedMessage()); + EXPECT_EQ(cel::to_address(element.GetParsedMessage()), + &message->repeated_nested_message(0)); + EXPECT_EQ(arena()->SpaceUsed(), 0); +} + +TEST_F(ParsedRepeatedFieldValueTest, NewIteratorUnsafeBorrowsMessages) { + TestAllTypesProto3 heap_message; + heap_message.add_repeated_nested_message()->set_bb(1); + heap_message.add_repeated_nested_message()->set_bb(2); + + google::protobuf::Arena source_arena; + auto* arena_message = + google::protobuf::Arena::Create(&source_arena); + arena_message->CopyFrom(heap_message); + + for (const TestAllTypesProto3* message : {&heap_message, arena_message}) { + ParsedRepeatedFieldValue value = UnsafeParsedRepeatedFieldValue( + message, TestAllTypesProto3::descriptor()->FindFieldByName( + "repeated_nested_message")); + ASSERT_OK_AND_ASSIGN(auto iterator, value.NewIterator()); + for (int i = 0; i < 2; ++i) { + SCOPED_TRACE(i); + ASSERT_OK_AND_ASSIGN( + auto element, + iterator->Next1(descriptor_pool(), message_factory(), arena())); + ASSERT_TRUE(element.has_value()); + ASSERT_TRUE(element->IsParsedMessage()); + EXPECT_EQ(cel::to_address(element->GetParsedMessage()), + &message->repeated_nested_message(i)); + } + } + EXPECT_EQ(arena()->SpaceUsed(), 0); +} + +TEST_F(ParsedRepeatedFieldValueTest, ForEachUnsafeBorrowsStringsAndBytes) { + TestAllTypesProto3 heap_message; + heap_message.add_repeated_string(std::string(64, 'a')); + heap_message.add_repeated_bytes(std::string(64, 'b')); + + google::protobuf::Arena source_arena; + auto* arena_message = + google::protobuf::Arena::Create(&source_arena); + arena_message->CopyFrom(heap_message); + + for (const TestAllTypesProto3* message : {&heap_message, arena_message}) { + std::vector> flat; + for (absl::string_view field_name : {"repeated_string", "repeated_bytes"}) { + ParsedRepeatedFieldValue value = UnsafeParsedRepeatedFieldValue( + message, + TestAllTypesProto3::descriptor()->FindFieldByName(field_name)); + EXPECT_THAT(value.ForEach( + [&](const Value& element) -> absl::StatusOr { + if (element.IsString()) { + flat.push_back(element.GetString().TryFlat()); + } else { + flat.push_back(element.GetBytes().TryFlat()); + } + return true; + }, + descriptor_pool(), message_factory(), arena()), + IsOk()); + } + ASSERT_THAT(flat, ElementsAre(Optional(_), Optional(_))); + EXPECT_EQ(flat[0]->data(), message->repeated_string(0).data()); + EXPECT_EQ(flat[1]->data(), message->repeated_bytes(0).data()); + } + EXPECT_EQ(arena()->SpaceUsed(), 0); +} + +TEST_F(ParsedRepeatedFieldValueTest, ContainsUnsafeBorrowsElements) { + google::protobuf::Arena source_arena; + auto* message = google::protobuf::Arena::Create(&source_arena); + message->add_repeated_string(std::string(64, 'a')); + message->add_repeated_nested_message()->set_bb(42); + + ParsedRepeatedFieldValue strings = UnsafeParsedRepeatedFieldValue( + message, + TestAllTypesProto3::descriptor()->FindFieldByName("repeated_string")); + const std::string target_str(64, 'a'); + EXPECT_THAT(strings.Contains(StringValue::WrapUnsafe(target_str), + descriptor_pool(), message_factory(), arena()), + IsOkAndHolds(BoolValueIs(true))); + + TestAllTypesProto3::NestedMessage target_msg; + target_msg.set_bb(42); + ParsedRepeatedFieldValue messages = UnsafeParsedRepeatedFieldValue( + message, TestAllTypesProto3::descriptor()->FindFieldByName( + "repeated_nested_message")); + EXPECT_THAT(messages.Contains(UnsafeParsedMessageValue(&target_msg), + descriptor_pool(), message_factory(), arena()), + IsOkAndHolds(BoolValueIs(true))); + EXPECT_EQ(arena()->SpaceUsed(), 0); +} + TEST_F(ParsedRepeatedFieldValueTest, Contains) { ParsedRepeatedFieldValue value( DynamicParseTextProto(R"pb(repeated_bool: true)pb"), diff --git a/eval/public/BUILD b/eval/public/BUILD index d3340c893..d361ed45b 100644 --- a/eval/public/BUILD +++ b/eval/public/BUILD @@ -204,9 +204,10 @@ cc_library( ], deps = [ ":activation", + "//common:value", "//eval/public/containers:field_access", - "//eval/public/containers:field_backed_list_impl", "//eval/public/containers:field_backed_map_impl", + "//internal:status_macros", "@com_google_absl//absl/status", ], ) @@ -793,11 +794,24 @@ cc_test( deps = [ ":activation", ":activation_bind_helper", + ":builtin_func_registrar", + ":cel_expr_builder_factory", + ":cel_expression", + ":cel_function", + ":cel_options", + ":cel_value", "//eval/testutil:test_message_cc_proto", "//internal:status_macros", "//internal:testing", + "//parser", "//testutil:util", "@com_google_absl//absl/status", + "@com_google_absl//absl/status:status_matchers", + "@com_google_absl//absl/status:statusor", + "@com_google_absl//absl/strings:string_view", + "@com_google_absl//absl/types:span", + "@com_google_cel_spec//proto/cel/expr:syntax_cc_proto", + "@com_google_protobuf//:protobuf", ], ) diff --git a/eval/public/activation_bind_helper.cc b/eval/public/activation_bind_helper.cc index 1e8004003..e1e01789d 100644 --- a/eval/public/activation_bind_helper.cc +++ b/eval/public/activation_bind_helper.cc @@ -1,9 +1,11 @@ #include "eval/public/activation_bind_helper.h" #include "absl/status/status.h" +#include "common/legacy_value.h" +#include "common/value.h" #include "eval/public/containers/field_access.h" -#include "eval/public/containers/field_backed_list_impl.h" #include "eval/public/containers/field_backed_map_impl.h" +#include "internal/status_macros.h" namespace google { namespace api { @@ -25,8 +27,9 @@ absl::Status CreateValueFromField(const google::protobuf::Message* msg, arena, msg, field_desc, arena)); return absl::OkStatus(); } else if (field_desc->is_repeated()) { - *result = CelValue::CreateList(google::protobuf::Arena::Create( - arena, msg, field_desc, arena)); + CEL_ASSIGN_OR_RETURN( + *result, cel::LegacyValue(arena, cel::UnsafeParsedRepeatedFieldValue( + msg, field_desc))); return absl::OkStatus(); } else { return CreateValueFromSingleField(msg, field_desc, arena, result); diff --git a/eval/public/activation_bind_helper_test.cc b/eval/public/activation_bind_helper_test.cc index f644f1669..2d2e197b5 100644 --- a/eval/public/activation_bind_helper_test.cc +++ b/eval/public/activation_bind_helper_test.cc @@ -1,11 +1,27 @@ #include "eval/public/activation_bind_helper.h" +#include +#include + +#include "cel/expr/syntax.pb.h" #include "absl/status/status.h" +#include "absl/status/status_matchers.h" +#include "absl/status/statusor.h" +#include "absl/strings/string_view.h" +#include "absl/types/span.h" #include "eval/public/activation.h" +#include "eval/public/builtin_func_registrar.h" +#include "eval/public/cel_expr_builder_factory.h" +#include "eval/public/cel_expression.h" +#include "eval/public/cel_function.h" +#include "eval/public/cel_options.h" +#include "eval/public/cel_value.h" #include "eval/testutil/test_message.pb.h" #include "internal/status_macros.h" #include "internal/testing.h" +#include "parser/parser.h" #include "testutil/util.h" +#include "google/protobuf/arena.h" namespace google { namespace api { @@ -14,6 +30,7 @@ namespace runtime { namespace { +using ::absl_testing::IsOk; using testutil::EqualsProto; TEST(ActivationBindHelperTest, TestSingleBoolBind) { @@ -138,6 +155,125 @@ TEST(ActivationBindHelperTest, RejectsNullArena) { "arena must not be null for BindProtoToActivation.")); } +class RecordListPtrFunction final : public CelFunction { + public: + explicit RecordListPtrFunction(const CelList** recorded) + : CelFunction(CelFunctionDescriptor("record_list_ptr", + /*receiver_style=*/false, + {CelValue::Type::kList}, + /*is_strict=*/true)), + recorded_(recorded) {} + + absl::Status Evaluate(absl::Span arguments, CelValue* result, + google::protobuf::Arena*) const override { + *recorded_ = arguments[0].ListOrDie(); + *result = CelValue::CreateInt64(arguments[0].ListOrDie()->size()); + return absl::OkStatus(); + } + + private: + const CelList** recorded_; +}; + +absl::StatusOr> CompileExpr( + CelExpressionBuilder& builder, absl::string_view expr) { + CEL_ASSIGN_OR_RETURN(auto parsed, parser::Parse(expr)); + return builder.CreateExpression(&parsed.expr(), &parsed.source_info()); +} + +TEST(ActivationBindHelperTest, RepeatedFieldsBorrowElementsAndEvaluateCleanly) { + TestMessage heap_message; + heap_message.add_int32_list(3); + heap_message.add_int32_list(7); + heap_message.add_string_list(std::string(64, 'a')); + heap_message.add_string_list(std::string(64, 'b')); + TestMessage* nested0 = heap_message.add_message_list(); + nested0->set_int32_value(1); + nested0->set_string_value(std::string(64, 'x')); + TestMessage* nested1 = heap_message.add_message_list(); + nested1->set_int32_value(7); + nested1->set_string_value(std::string(64, 'y')); + + google::protobuf::Arena source_arena; + auto* arena_message = google::protobuf::Arena::Create(&source_arena); + arena_message->CopyFrom(heap_message); + + const CelList* recorded_list = nullptr; + InterpreterOptions options; + std::unique_ptr builder = + CreateCelExpressionBuilder(options); + ASSERT_THAT(RegisterBuiltinFunctions(builder->GetRegistry(), options), + IsOk()); + ASSERT_THAT(builder->GetRegistry()->Register( + std::make_unique(&recorded_list)), + IsOk()); + + ASSERT_OK_AND_ASSIGN( + auto exists_expr, + CompileExpr(*builder, + "message_list.exists(m, m.int32_value == 7 && " + "m.string_value == '" + + std::string(64, 'y') + + "') && string_list.exists(s, s == '" + + std::string(64, 'b') + "')")); + ASSERT_OK_AND_ASSIGN(auto index_expr, + CompileExpr(*builder, "message_list[1]")); + ASSERT_OK_AND_ASSIGN( + auto contains_expr, + CompileExpr(*builder, "7 in int32_list && '" + std::string(64, 'a') + + "' in string_list")); + ASSERT_OK_AND_ASSIGN(auto legacy_fn_expr, + CompileExpr(*builder, "record_list_ptr(message_list)")); + + for (const TestMessage* msg : {&heap_message, arena_message}) { + google::protobuf::Arena eval_arena; + Activation activation; + ASSERT_THAT(BindProtoToActivation(msg, &eval_arena, &activation), IsOk()); + + auto bound_msg_list = activation.FindValue("message_list", &eval_arena); + ASSERT_TRUE(bound_msg_list.has_value()); + ASSERT_TRUE(bound_msg_list->IsList()); + + // First evaluation may initialize steady-state frame scratch; subsequent + // evaluations over borrowed repeated messages/strings allocate zero arena + // memory for elements. + ASSERT_OK_AND_ASSIGN(CelValue exists_val, + exists_expr->Evaluate(activation, &eval_arena)); + ASSERT_TRUE(exists_val.IsBool()); + EXPECT_TRUE(exists_val.BoolOrDie()); + + google::protobuf::Arena iter_arena; + ASSERT_OK_AND_ASSIGN(exists_val, + exists_expr->Evaluate(activation, &iter_arena)); + EXPECT_TRUE(exists_val.BoolOrDie()); + EXPECT_EQ(iter_arena.SpaceUsed(), 0); + + // Indexing returns a legacy MessageWrapper pointing directly to the + // backing repeated submessage. + ASSERT_OK_AND_ASSIGN(CelValue indexed_val, + index_expr->Evaluate(activation, &eval_arena)); + ASSERT_TRUE(indexed_val.IsMessage()); + EXPECT_EQ(indexed_val.MessageOrDie(), &msg->message_list(1)); + + // Membership checks succeed without copying onto a fresh evaluation arena. + google::protobuf::Arena contains_arena; + ASSERT_OK_AND_ASSIGN(CelValue contains_val, + contains_expr->Evaluate(activation, &contains_arena)); + ASSERT_TRUE(contains_val.IsBool()); + EXPECT_TRUE(contains_val.BoolOrDie()); + EXPECT_EQ(contains_arena.SpaceUsed(), 0); + + // Passing a bound repeated field to a legacy CelFunction preserves the + // exact CelList pointer with no re-wrapping. + recorded_list = nullptr; + ASSERT_OK_AND_ASSIGN(CelValue fn_val, + legacy_fn_expr->Evaluate(activation, &eval_arena)); + ASSERT_TRUE(fn_val.IsInt64()); + EXPECT_EQ(fn_val.Int64OrDie(), 2); + EXPECT_EQ(recorded_list, bound_msg_list->ListOrDie()); + } +} + } // namespace } // namespace runtime