@@ -36,11 +36,9 @@ namespace {
3636
3737using ::cel::BoolValue;
3838using ::cel::ErrorValue;
39- using ::cel::MapValue;
4039using ::cel::OptionalValue;
4140using ::cel::ProtoWrapperTypeOptions;
4241using ::cel::StringValue;
43- using ::cel::StructValue;
4442using ::cel::Value;
4543using ::cel::ValueKind;
4644
@@ -424,6 +422,42 @@ class DirectSelectStep : public DirectExpressionStep {
424422 bool enable_optional_types_;
425423};
426424
425+ bool CheckAttributeTrail (const std::string& field, ExecutionFrame* frame) {
426+ if (!frame->attribute_tracking_enabled ()) {
427+ return false ;
428+ }
429+ AttributeTrail& attr = frame->value_stack ().PeekAttribute ();
430+ attr = attr.Step (&field);
431+
432+ absl::optional<Value> marked_attribute_check =
433+ CheckForMarkedAttributes (attr, *frame);
434+ if (marked_attribute_check.has_value ()) {
435+ frame->value_stack ().Peek () = std::move (marked_attribute_check).value ();
436+ return true ;
437+ }
438+
439+ return false ;
440+ }
441+
442+ bool SupportsCachedFieldDescriptor (
443+ const cel::ParsedMessageValue& parsed_message,
444+ const google::protobuf::Descriptor* descriptor,
445+ const google::protobuf::FieldDescriptor* field_descriptor) {
446+ ABSL_DCHECK_EQ (field_descriptor->containing_type (), descriptor);
447+ const google::protobuf::Descriptor* rt_descriptor = parsed_message.GetDescriptor ();
448+
449+ if (rt_descriptor != descriptor) {
450+ return false ;
451+ }
452+
453+ // Caller should have already checked this. Crash here if not instead of
454+ // making proto crash.
455+ ABSL_DCHECK_EQ (rt_descriptor->file ()->pool (),
456+ field_descriptor->file ()->pool ());
457+
458+ return true ;
459+ }
460+
427461class ProtoSelectStep : public SelectStep {
428462 public:
429463 ProtoSelectStep (StringValue value, int64_t expr_id,
@@ -447,50 +481,40 @@ class ProtoSelectStep : public SelectStep {
447481
448482 const Value& arg = frame->value_stack ().Peek ();
449483 if (auto unwrapped = arg.AsParsedMessage ();
450- unwrapped.has_value () && unwrapped->GetDescriptor () == descriptor_) {
451- return EvaluateModernMessageGetField (frame, *unwrapped);
484+ unwrapped.has_value () &&
485+ SupportsCachedFieldDescriptor (*unwrapped, descriptor_,
486+ field_descriptor_)) {
487+ return EvaluateMessageFieldGet (frame, *unwrapped);
452488 } else if (const google::protobuf::Message* legacy_message =
453489 cel::interop_internal::GetLegacyMessage (arg);
454- legacy_message != nullptr &&
455- legacy_message->GetDescriptor () == descriptor_) {
490+ frame->options ().enable_use_new_field_select_implementation &&
491+ legacy_message != nullptr ) {
492+ auto parsed_message = cel::UnsafeParsedMessageValue (legacy_message);
456493 // A little unfortunate, but need to special case for legacy values so we
457494 // can minimize back and forth interop conversions.
458- return EvaluateLegacyMessageGetField (frame, legacy_message);
495+ if (SupportsCachedFieldDescriptor (parsed_message, descriptor_,
496+ field_descriptor_)) {
497+ return EvaluateMessageFieldGet (frame, legacy_message);
498+ }
459499 }
460500 // If we get an unexpected value type, fall back to the generic
461501 // implementation.
462502 return SelectStep::Evaluate (frame);
463503 }
464504
465505 private:
466- absl::Status EvaluateModernMessageGetField (
506+ absl::Status EvaluateMessageFieldGet (
467507 ExecutionFrame* frame,
468508 const cel::ParsedMessageValue& parsed_message) const ;
469- absl::Status EvaluateLegacyMessageGetField (
470- ExecutionFrame* frame, const google::protobuf::Message* legacy_message) const ;
509+ absl::Status EvaluateMessageFieldGet (
510+ ExecutionFrame* frame,
511+ const google::protobuf::Message* absl_nonnull legacy_message) const ;
471512
472513 const google::protobuf::Descriptor* descriptor_;
473514 const google::protobuf::FieldDescriptor* field_descriptor_;
474515};
475516
476- bool CheckAttributeTrail (const std::string& field, ExecutionFrame* frame) {
477- if (!frame->attribute_tracking_enabled ()) {
478- return false ;
479- }
480- AttributeTrail& attr = frame->value_stack ().PeekAttribute ();
481- attr = attr.Step (&field);
482-
483- absl::optional<Value> marked_attribute_check =
484- CheckForMarkedAttributes (attr, *frame);
485- if (marked_attribute_check.has_value ()) {
486- frame->value_stack ().Peek () = std::move (marked_attribute_check).value ();
487- return true ;
488- }
489-
490- return false ;
491- }
492-
493- absl::Status ProtoSelectStep::EvaluateModernMessageGetField (
517+ absl::Status ProtoSelectStep::EvaluateMessageFieldGet (
494518 ExecutionFrame* frame,
495519 const cel::ParsedMessageValue& parsed_message) const {
496520 if (CheckAttributeTrail (field_, frame)) {
@@ -501,8 +525,10 @@ absl::Status ProtoSelectStep::EvaluateModernMessageGetField(
501525 frame->message_factory (), frame->arena (), &frame->value_stack ().Peek ());
502526}
503527
504- absl::Status ProtoSelectStep::EvaluateLegacyMessageGetField (
505- ExecutionFrame* frame, const google::protobuf::Message* legacy_message) const {
528+ absl::Status ProtoSelectStep::EvaluateMessageFieldGet (
529+ ExecutionFrame* frame,
530+ const google::protobuf::Message* absl_nonnull legacy_message) const {
531+ ABSL_DCHECK (legacy_message != nullptr );
506532 if (CheckAttributeTrail (field_, frame)) {
507533 return absl::OkStatus ();
508534 }
@@ -534,15 +560,19 @@ class ProtoHasStep : public SelectStep {
534560
535561 const Value& arg = frame->value_stack ().Peek ();
536562 if (auto unwrapped = arg.AsParsedMessage ();
537- unwrapped.has_value () && unwrapped->GetDescriptor () == descriptor_) {
563+ unwrapped.has_value () &&
564+ SupportsCachedFieldDescriptor (*unwrapped, descriptor_,
565+ field_descriptor_)) {
538566 return EvaluateHas (frame, *unwrapped);
539567 } else if (const google::protobuf::Message* legacy_message =
540568 cel::interop_internal::GetLegacyMessage (arg);
541- legacy_message != nullptr &&
542- legacy_message->GetDescriptor () == descriptor_) {
569+ legacy_message != nullptr ) {
543570 cel::ParsedMessageValue parsed_message =
544571 cel::UnsafeParsedMessageValue (legacy_message);
545- return EvaluateHas (frame, parsed_message);
572+ if (SupportsCachedFieldDescriptor (parsed_message, descriptor_,
573+ field_descriptor_)) {
574+ return EvaluateHas (frame, parsed_message);
575+ }
546576 }
547577 // If we get an unexpected value type, fall back to the generic
548578 // implementation.
@@ -608,6 +638,20 @@ absl::StatusOr<std::unique_ptr<ExpressionStep>> CreateTypedSelectStep(
608638 const google::protobuf::FieldDescriptor* field_descriptor =
609639 resolved_field.GetMessage ().descriptor ();
610640
641+ if (field_descriptor->file ()->pool () != descriptor->file ()->pool ()) {
642+ // The field descriptor is not in the same pool as the operand type.
643+ // (this should only happen if an overlay extends the type).
644+ //
645+ // We don't have a way to determine if a runtime message is compatible
646+ // with the resolved extension and proto's reflection implementation may
647+ // crash.
648+ //
649+ // Fallback to the generic implementation.
650+ return CreateSelectStep (std::move (field), test_only, expr_id,
651+ enable_wrapper_type_null_unboxing,
652+ enable_optional_types);
653+ }
654+
611655 if (test_only) {
612656 return std::make_unique<ProtoHasStep>(
613657 std::move (field), expr_id, enable_wrapper_type_null_unboxing,
0 commit comments