Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 12 additions & 2 deletions common/values/struct_value_builder.cc
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@
#include <cstdint>
#include <limits>
#include <memory>
#include <optional>
#include <string>
#include <utility>

Expand Down Expand Up @@ -84,8 +85,17 @@ absl::StatusOr<absl::optional<ErrorValue>> ProtoMessageCopy(
const google::protobuf::Message* absl_nonnull from_message) {
CEL_ASSIGN_OR_RETURN(const auto* from_descriptor,
GetDescriptor(*from_message));
if (to_descriptor == from_descriptor) {
// Same.
if (to_descriptor == from_descriptor &&
to_message->GetReflection()->GetMessageFactory() ==
from_message->GetReflection()->GetMessageFactory()) {
// Same type, use proto reflection copy.
//
// We use the slower serialization copy if the factory is different to avoid
// adding an implicit lifetime dependency on the other factory.
//
// This should only happen if the embedding application is calling the
// builder directly or attempting to set the field from an unsafe wrapped
// message.
to_message->CopyFrom(*from_message);
return std::nullopt;
}
Expand Down
112 changes: 78 additions & 34 deletions eval/eval/select_step.cc
Original file line number Diff line number Diff line change
Expand Up @@ -36,11 +36,9 @@ namespace {

using ::cel::BoolValue;
using ::cel::ErrorValue;
using ::cel::MapValue;
using ::cel::OptionalValue;
using ::cel::ProtoWrapperTypeOptions;
using ::cel::StringValue;
using ::cel::StructValue;
using ::cel::Value;
using ::cel::ValueKind;

Expand Down Expand Up @@ -424,6 +422,42 @@ class DirectSelectStep : public DirectExpressionStep {
bool enable_optional_types_;
};

bool CheckAttributeTrail(const std::string& field, ExecutionFrame* frame) {
if (!frame->attribute_tracking_enabled()) {
return false;
}
AttributeTrail& attr = frame->value_stack().PeekAttribute();
attr = attr.Step(&field);

absl::optional<Value> marked_attribute_check =
CheckForMarkedAttributes(attr, *frame);
if (marked_attribute_check.has_value()) {
frame->value_stack().Peek() = std::move(marked_attribute_check).value();
return true;
}

return false;
}

bool SupportsCachedFieldDescriptor(
const cel::ParsedMessageValue& parsed_message,
const google::protobuf::Descriptor* descriptor,
const google::protobuf::FieldDescriptor* field_descriptor) {
ABSL_DCHECK_EQ(field_descriptor->containing_type(), descriptor);
const google::protobuf::Descriptor* rt_descriptor = parsed_message.GetDescriptor();

if (rt_descriptor != descriptor) {
return false;
}

// Caller should have already checked this. Crash here if not instead of
// making proto crash.
ABSL_DCHECK_EQ(rt_descriptor->file()->pool(),
field_descriptor->file()->pool());

return true;
}

class ProtoSelectStep : public SelectStep {
public:
ProtoSelectStep(StringValue value, int64_t expr_id,
Expand All @@ -447,50 +481,40 @@ class ProtoSelectStep : public SelectStep {

const Value& arg = frame->value_stack().Peek();
if (auto unwrapped = arg.AsParsedMessage();
unwrapped.has_value() && unwrapped->GetDescriptor() == descriptor_) {
return EvaluateModernMessageGetField(frame, *unwrapped);
unwrapped.has_value() &&
SupportsCachedFieldDescriptor(*unwrapped, descriptor_,
field_descriptor_)) {
return EvaluateMessageFieldGet(frame, *unwrapped);
} else if (const google::protobuf::Message* legacy_message =
cel::interop_internal::GetLegacyMessage(arg);
legacy_message != nullptr &&
legacy_message->GetDescriptor() == descriptor_) {
frame->options().enable_use_new_field_select_implementation &&
legacy_message != nullptr) {
auto parsed_message = cel::UnsafeParsedMessageValue(legacy_message);
// A little unfortunate, but need to special case for legacy values so we
// can minimize back and forth interop conversions.
return EvaluateLegacyMessageGetField(frame, legacy_message);
if (SupportsCachedFieldDescriptor(parsed_message, descriptor_,
field_descriptor_)) {
return EvaluateMessageFieldGet(frame, legacy_message);
}
}
// If we get an unexpected value type, fall back to the generic
// implementation.
return SelectStep::Evaluate(frame);
}

private:
absl::Status EvaluateModernMessageGetField(
absl::Status EvaluateMessageFieldGet(
ExecutionFrame* frame,
const cel::ParsedMessageValue& parsed_message) const;
absl::Status EvaluateLegacyMessageGetField(
ExecutionFrame* frame, const google::protobuf::Message* legacy_message) const;
absl::Status EvaluateMessageFieldGet(
ExecutionFrame* frame,
const google::protobuf::Message* absl_nonnull legacy_message) const;

const google::protobuf::Descriptor* descriptor_;
const google::protobuf::FieldDescriptor* field_descriptor_;
};

bool CheckAttributeTrail(const std::string& field, ExecutionFrame* frame) {
if (!frame->attribute_tracking_enabled()) {
return false;
}
AttributeTrail& attr = frame->value_stack().PeekAttribute();
attr = attr.Step(&field);

absl::optional<Value> marked_attribute_check =
CheckForMarkedAttributes(attr, *frame);
if (marked_attribute_check.has_value()) {
frame->value_stack().Peek() = std::move(marked_attribute_check).value();
return true;
}

return false;
}

absl::Status ProtoSelectStep::EvaluateModernMessageGetField(
absl::Status ProtoSelectStep::EvaluateMessageFieldGet(
ExecutionFrame* frame,
const cel::ParsedMessageValue& parsed_message) const {
if (CheckAttributeTrail(field_, frame)) {
Expand All @@ -501,8 +525,10 @@ absl::Status ProtoSelectStep::EvaluateModernMessageGetField(
frame->message_factory(), frame->arena(), &frame->value_stack().Peek());
}

absl::Status ProtoSelectStep::EvaluateLegacyMessageGetField(
ExecutionFrame* frame, const google::protobuf::Message* legacy_message) const {
absl::Status ProtoSelectStep::EvaluateMessageFieldGet(
ExecutionFrame* frame,
const google::protobuf::Message* absl_nonnull legacy_message) const {
ABSL_DCHECK(legacy_message != nullptr);
if (CheckAttributeTrail(field_, frame)) {
return absl::OkStatus();
}
Expand Down Expand Up @@ -534,15 +560,19 @@ class ProtoHasStep : public SelectStep {

const Value& arg = frame->value_stack().Peek();
if (auto unwrapped = arg.AsParsedMessage();
unwrapped.has_value() && unwrapped->GetDescriptor() == descriptor_) {
unwrapped.has_value() &&
SupportsCachedFieldDescriptor(*unwrapped, descriptor_,
field_descriptor_)) {
return EvaluateHas(frame, *unwrapped);
} else if (const google::protobuf::Message* legacy_message =
cel::interop_internal::GetLegacyMessage(arg);
legacy_message != nullptr &&
legacy_message->GetDescriptor() == descriptor_) {
legacy_message != nullptr) {
cel::ParsedMessageValue parsed_message =
cel::UnsafeParsedMessageValue(legacy_message);
return EvaluateHas(frame, parsed_message);
if (SupportsCachedFieldDescriptor(parsed_message, descriptor_,
field_descriptor_)) {
return EvaluateHas(frame, parsed_message);
}
}
// If we get an unexpected value type, fall back to the generic
// implementation.
Expand Down Expand Up @@ -608,6 +638,20 @@ absl::StatusOr<std::unique_ptr<ExpressionStep>> CreateTypedSelectStep(
const google::protobuf::FieldDescriptor* field_descriptor =
resolved_field.GetMessage().descriptor();

if (field_descriptor->file()->pool() != descriptor->file()->pool()) {
// The field descriptor is not in the same pool as the operand type.
// (this should only happen if an overlay extends the type).
//
// We don't have a way to determine if a runtime message is compatible
// with the resolved extension and proto's reflection implementation may
// crash.
//
// Fallback to the generic implementation.
return CreateSelectStep(std::move(field), test_only, expr_id,
enable_wrapper_type_null_unboxing,
enable_optional_types);
}

if (test_only) {
return std::make_unique<ProtoHasStep>(
std::move(field), expr_id, enable_wrapper_type_null_unboxing,
Expand Down
1 change: 1 addition & 0 deletions extensions/protobuf/BUILD
Original file line number Diff line number Diff line change
Expand Up @@ -123,6 +123,7 @@ cc_test(
"//common:value_kind",
"//common:value_testing",
"//internal:testing",
"//testutil:test_external_extensions_descriptor_set",
"@com_google_absl//absl/log:absl_check",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:status_matchers",
Expand Down
112 changes: 79 additions & 33 deletions extensions/protobuf/value.h
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,70 @@

namespace cel::extensions {

namespace extensions_internal {
template <typename SkipCopyCheck>
absl::Status ProtoMessageFromValue(const cel::Value& value,
google::protobuf::Message& dest_message) {
const auto* dest_descriptor = dest_message.GetDescriptor();
const google::protobuf::Message* src_message = nullptr;
if (auto legacy_struct_value =
cel::common_internal::AsLegacyStructValue(value);
legacy_struct_value) {
src_message = legacy_struct_value->message_ptr();
}
if (auto parsed_message_value = value.AsParsedMessage();
parsed_message_value) {
src_message = cel::to_address(*parsed_message_value);
}

if (src_message == nullptr) {
return TypeConversionError(value.GetRuntimeType(),
MessageType(dest_descriptor))
.NativeValue();
}

const auto* src_descriptor = src_message->GetDescriptor();
if (dest_descriptor != src_descriptor) {
goto slow_path;
}

if constexpr (!SkipCopyCheck::value) {
// Try to catch cases where we'll take an implicit dependency on a
// dynamic message factory.
//
// This isn't exhaustive, but correctly checking requires fully
// traversing the source message which will approach the cost of the
// serialization round trip.
if (dest_message.GetReflection()->GetMessageFactory() !=
src_message->GetReflection()->GetMessageFactory()) {
goto slow_path;
}
}

dest_message.CopyFrom(*src_message);
return absl::OkStatus();

slow_path:
if (dest_descriptor->full_name() != src_descriptor->full_name()) {
return TypeConversionError(value.GetRuntimeType(),
MessageType(dest_descriptor))
.NativeValue();
}

absl::Cord serialized;
if (!src_message->SerializePartialToCord(&serialized)) {
return absl::UnknownError(absl::StrCat("failed to serialize message: ",
src_descriptor->full_name()));
}
if (!dest_message.ParsePartialFromCord(serialized)) {
return absl::UnknownError(absl::StrCat("failed to parse message: ",
dest_descriptor->full_name()));
}
return absl::OkStatus();
}

} // namespace extensions_internal

// Adapt a protobuf message to a cel::Value.
//
// Handles unwrapping message types with special meanings in CEL (WKTs).
Expand All @@ -56,41 +120,23 @@ ProtoMessageToValue(T&& value,
message_factory, arena);
}

// Unwraps a protobuf message from a cel::Value.
inline absl::Status ProtoMessageFromValue(const Value& value,
google::protobuf::Message& dest_message) {
const auto* dest_descriptor = dest_message.GetDescriptor();
const google::protobuf::Message* src_message = nullptr;
if (auto legacy_struct_value =
cel::common_internal::AsLegacyStructValue(value);
legacy_struct_value) {
src_message = legacy_struct_value->message_ptr();
}
if (auto parsed_message_value = value.AsParsedMessage();
parsed_message_value) {
src_message = cel::to_address(*parsed_message_value);
}
if (src_message != nullptr) {
const auto* src_descriptor = src_message->GetDescriptor();
if (dest_descriptor == src_descriptor) {
dest_message.CopyFrom(*src_message);
return absl::OkStatus();
}
if (dest_descriptor->full_name() == src_descriptor->full_name()) {
absl::Cord serialized;
if (!src_message->SerializePartialToCord(&serialized)) {
return absl::UnknownError(absl::StrCat("failed to serialize message: ",
src_descriptor->full_name()));
}
if (!dest_message.ParsePartialFromCord(serialized)) {
return absl::UnknownError(absl::StrCat("failed to parse message: ",
dest_descriptor->full_name()));
}
return absl::OkStatus();
}
}
return TypeConversionError(value.GetRuntimeType(),
MessageType(dest_descriptor))
.NativeValue();
return extensions_internal::ProtoMessageFromValue<std::false_type>(
value, dest_message);
}

// Unwraps a protobuf message from a cel::Value without checking for the
// presence of extensions.
//
// Warning: This function can lead to subtle use after free bugs if the caller
// is not careful to ensure that the source and destination message were created
// in a compatible way and do not outlive any implicit dependencies.
inline absl::Status ProtoMessageFromValueUnsafe(const Value& value,
google::protobuf::Message& dest_message) {
return extensions_internal::ProtoMessageFromValue<std::true_type>(
value, dest_message);
}

} // namespace cel::extensions
Expand Down
Loading
Loading