Depth-limit upb deep clone/copy and propagate the limit to non-canonical extension encoding. PiperOrigin-RevId: 1002638556
diff --git a/ruby/src/main/java/com/google/protobuf/jruby/RubyMessage.java b/ruby/src/main/java/com/google/protobuf/jruby/RubyMessage.java index cf671ed..06bd685 100644 --- a/ruby/src/main/java/com/google/protobuf/jruby/RubyMessage.java +++ b/ruby/src/main/java/com/google/protobuf/jruby/RubyMessage.java
@@ -879,7 +879,7 @@ } protected DynamicMessage build(ThreadContext context, int depth, int recursionLimit) { - if (depth >= recursionLimit) { + if (depth > recursionLimit) { throw context.runtime.newRuntimeError("Recursion limit exceeded during encoding."); }
diff --git a/ruby/tests/encode_decode_test.rb b/ruby/tests/encode_decode_test.rb index 3cfaf7f..2ccd934 100755 --- a/ruby/tests/encode_decode_test.rb +++ b/ruby/tests/encode_decode_test.rb
@@ -206,10 +206,10 @@ assert_match msg.to_json, msg_out.to_json assert_raises RuntimeError do - A::B::C::TestMessage.encode(msg, { recursion_limit: 5 }) + A::B::C::TestMessage.encode(msg, { recursion_limit: 4 }) end - msg_encoded = A::B::C::TestMessage.encode(msg, { recursion_limit: 6 }) + msg_encoded = A::B::C::TestMessage.encode(msg, { recursion_limit: 5 }) msg_out = A::B::C::TestMessage.decode(msg_encoded) assert_match msg.to_json, msg_out.to_json end
diff --git a/upb/conformance/conformance_upb_harness_test.cc b/upb/conformance/conformance_upb_harness_test.cc index 5e7c393..18b8567 100644 --- a/upb/conformance/conformance_upb_harness_test.cc +++ b/upb/conformance/conformance_upb_harness_test.cc
@@ -288,13 +288,12 @@ "conformance_upb: parse of ConformanceRequest failed."); } -TEST_P(ConformanceUpbHarnessTest, ProtobufOutputSerializeErrorIsReported) { +TEST_P(ConformanceUpbHarnessTest, ProtobufOutputAtDepthLimitRoundTrips) { upb::Arena arena; conformance_ConformanceRequest* request = NewRequest(arena, conformance_PROTOBUF); - // recursive_message (27) nested 100 deep: upb's decoder accepts exactly its - // depth limit (100) of nested messages, but its encoder rejects the 100th, - // so this parses and then fails to serialize. + // recursive_message (27) nested 100 deep: exactly upb's default depth limit. + // The decoder accepts it, so the encoder must be able to serialize it again. Wire nested; for (int i = 0; i < 100; ++i) { nested = LengthPrefixedField(27, nested); @@ -304,6 +303,27 @@ const conformance_ConformanceResponse* response = Run(request, arena); ASSERT_NE(response, nullptr); + ASSERT_TRUE(conformance_ConformanceResponse_has_protobuf_payload(response)); + EXPECT_EQ( + ToString(conformance_ConformanceResponse_protobuf_payload(response)), + nested.data()); +} + +TEST_P(ConformanceUpbHarnessTest, ProtobufOutputSerializeErrorIsReported) { + upb::Arena arena; + conformance_ConformanceRequest* request = + NewRequest(arena, conformance_PROTOBUF); + // optional_value with 60 nested JSON arrays. The JSON parser limits JSON + // nesting (64 levels), so this parses. But each nested array is a Value + // holding a ListValue, i.e. two levels of message nesting, so the message is + // 120 levels deep and exceeds the wire encoder's limit (100). + constexpr int kNestedArrays = 60; + std::string json = R"({"optionalValue":)" + std::string(kNestedArrays, '[') + + std::string(kNestedArrays, ']') + "}"; + conformance_ConformanceRequest_set_json_payload(request, FromString(json)); + + const conformance_ConformanceResponse* response = Run(request, arena); + ASSERT_NE(response, nullptr); ASSERT_TRUE(conformance_ConformanceResponse_has_serialize_error(response)); EXPECT_EQ(ToString(conformance_ConformanceResponse_serialize_error(response)), "Error serializing.");
diff --git a/upb/message/BUILD b/upb/message/BUILD index 717b96c..1fbf189 100644 --- a/upb/message/BUILD +++ b/upb/message/BUILD
@@ -483,6 +483,8 @@ "//src/google/protobuf", "//upb/base", "//upb/mem", + "//upb/mini_descriptor", + "//upb/mini_descriptor:internal", "//upb/mini_table", "//upb/port", "//upb/test:test_messages_proto2_upb_minitable",
diff --git a/upb/message/convert_test.cc b/upb/message/convert_test.cc index 304137c..378a0eb 100644 --- a/upb/message/convert_test.cc +++ b/upb/message/convert_test.cc
@@ -2765,8 +2765,8 @@ // Set recursive_message (field 27) with a submessage that: // 1) Has a large repeated_string (field 44, encoded first) forcing // upb_BackAlloc to allocate a standalone block. - // 2) Has another recursive_message (field 27, encoded after field 44) - // that exceeds the max depth limit during _upb_Encode_Field. + // 2) Has two more levels of recursive_message (field 27, encoded after + // field 44) that exceed the max depth limit during _upb_Encode_Field. TestMsg* sub = protobuf_test_messages_proto3_TestAllTypesProto3_new(src_arena.ptr()); protobuf_test_messages_proto3_TestAllTypesProto3_set_recursive_message(msg, @@ -2775,8 +2775,11 @@ protobuf_test_messages_proto3_TestAllTypesProto3_add_repeated_string( sub, upb_StringView_FromDataAndSize(large_str.data(), large_str.size()), src_arena.ptr()); + TestMsg* sub2 = + protobuf_test_messages_proto3_TestAllTypesProto3_mutable_recursive_message( + sub, src_arena.ptr()); protobuf_test_messages_proto3_TestAllTypesProto3_mutable_recursive_message( - sub, src_arena.ptr()); + sub2, src_arena.ptr()); upb::Arena dst_arena; const upb_MiniTable* empty_mt = &upb_0test__EmptyMessage_msg_init;
diff --git a/upb/message/copy.c b/upb/message/copy.c index 5233fd4..faa779a 100644 --- a/upb/message/copy.c +++ b/upb/message/copy.c
@@ -34,6 +34,7 @@ #include "upb/mini_table/message.h" #include "upb/wire/encode.h" #include "upb/wire/encode_extension.h" +#include "upb/wire/internal/constants.h" // Must be last. #include "upb/port/def.inc" @@ -51,8 +52,13 @@ return upb_StringView_FromDataAndSize(cloned_data, str.size); } +static upb_Message* upb_Message_DeepCloneWithDepth(const upb_Message* msg, + const upb_MiniTable* m, + upb_Arena* arena, int depth); + static bool upb_Clone_MessageValue(void* value, upb_CType value_type, - const upb_MiniTable* sub, upb_Arena* arena) { + const upb_MiniTable* sub, upb_Arena* arena, + int depth) { switch (value_type) { case kUpb_CType_Bool: case kUpb_CType_Float: @@ -80,7 +86,8 @@ UPB_ASSERT(sub); const upb_Message* source = *(upb_Message**)value; UPB_ASSERT(source); - upb_Message* clone = upb_Message_DeepClone(source, sub, arena); + upb_Message* clone = + upb_Message_DeepCloneWithDepth(source, sub, arena, depth - 1); *(upb_Message**)value = clone; return clone != NULL; } break; @@ -91,10 +98,11 @@ static bool upb_Clone_MapValue(upb_value* tabval, size_t map_val_size, upb_CType value_field_type, const upb_MiniTable* value_sub, upb_Arena* arena, - upb_value* cloned_tabval) { + int depth, upb_value* cloned_tabval) { upb_MessageValue val; _upb_map_fromvalue(*tabval, &val, map_val_size); - if (!upb_Clone_MessageValue(&val, value_field_type, value_sub, arena)) { + if (!upb_Clone_MessageValue(&val, value_field_type, value_sub, arena, + depth)) { return false; } *cloned_tabval = (upb_value){0}; // Initialize @@ -104,10 +112,9 @@ return true; } -upb_Map* upb_Map_DeepClone(const upb_Map* map, upb_CType key_type, - upb_CType value_type, - const upb_MiniTable* map_entry_table, - upb_Arena* arena) { +static upb_Map* upb_Map_DeepCloneWithDepth(const upb_Map* map, + const upb_MiniTable* map_entry_table, + upb_Arena* arena, int depth) { upb_Map* cloned_map = upb_Arena_Malloc(arena, sizeof(upb_Map)); if (cloned_map == NULL) { return NULL; @@ -138,7 +145,7 @@ upb_strtable_next2(&cloned_map->t.strtable, &key, &tabval, &iter)) { upb_value cloned_tabval; if (!upb_Clone_MapValue(&tabval, map->val_size, value_field_type, - value_sub, arena, &cloned_tabval)) { + value_sub, arena, depth, &cloned_tabval)) { return NULL; } upb_strtable_setentryvalue(&cloned_map->t.strtable, iter, @@ -156,7 +163,7 @@ while (upb_inttable_next(&cloned_map->t.inttable, &key, &tabval, &iter)) { upb_value cloned_tabval; if (!upb_Clone_MapValue(&tabval, map->val_size, value_field_type, - value_sub, arena, &cloned_tabval)) { + value_sub, arena, depth, &cloned_tabval)) { return NULL; } upb_inttable_setentryvalue(&cloned_map->t.inttable, iter, @@ -167,22 +174,25 @@ return cloned_map; } +upb_Map* upb_Map_DeepClone(const upb_Map* map, upb_CType key_type, + upb_CType value_type, + const upb_MiniTable* map_entry_table, + upb_Arena* arena) { + return upb_Map_DeepCloneWithDepth(map, map_entry_table, arena, + kUpb_WireFormat_DefaultDepthLimit); +} + static upb_Map* upb_Message_Map_DeepClone(const upb_Map* map, const upb_MiniTable* mini_table, const upb_MiniTableField* f, - upb_Message* clone, - upb_Arena* arena) { + upb_Message* clone, upb_Arena* arena, + int depth) { UPB_ASSERT(!upb_Message_IsFrozen(clone)); const upb_MiniTable* map_entry_table = upb_MiniTable_MapEntrySubMessage(f); UPB_ASSERT(map_entry_table); - const upb_MiniTableField* key_field = upb_MiniTable_MapKey(map_entry_table); - const upb_MiniTableField* value_field = - upb_MiniTable_MapValue(map_entry_table); - - upb_Map* cloned_map = upb_Map_DeepClone( - map, upb_MiniTableField_CType(key_field), - upb_MiniTableField_CType(value_field), map_entry_table, arena); + upb_Map* cloned_map = + upb_Map_DeepCloneWithDepth(map, map_entry_table, arena, depth); if (!cloned_map) { return NULL; } @@ -190,8 +200,10 @@ return cloned_map; } -upb_Array* upb_Array_DeepClone(const upb_Array* array, upb_CType value_type, - const upb_MiniTable* sub, upb_Arena* arena) { +static upb_Array* upb_Array_DeepCloneWithDepth(const upb_Array* array, + upb_CType value_type, + const upb_MiniTable* sub, + upb_Arena* arena, int depth) { const size_t size = upb_Array_Size(array); const int lg2 = UPB_PRIVATE(_upb_CType_SizeLg2)(value_type); upb_Array* cloned_array = UPB_PRIVATE(_upb_Array_New)(arena, size, lg2); @@ -203,7 +215,7 @@ } for (size_t i = 0; i < size; ++i) { upb_MessageValue val = upb_Array_Get(array, i); - if (!upb_Clone_MessageValue(&val, value_type, sub, arena)) { + if (!upb_Clone_MessageValue(&val, value_type, sub, arena, depth)) { return NULL; } upb_Array_Set(cloned_array, i, val); @@ -211,18 +223,25 @@ return cloned_array; } +upb_Array* upb_Array_DeepClone(const upb_Array* array, upb_CType value_type, + const upb_MiniTable* sub, upb_Arena* arena) { + return upb_Array_DeepCloneWithDepth(array, value_type, sub, arena, + kUpb_WireFormat_DefaultDepthLimit); +} + static bool upb_Message_Array_DeepClone(const upb_Array* array, const upb_MiniTable* mini_table, const upb_MiniTableField* field, - upb_Message* clone, upb_Arena* arena) { + upb_Message* clone, upb_Arena* arena, + int depth) { UPB_ASSERT(!upb_Message_IsFrozen(clone)); UPB_PRIVATE(_upb_MiniTableField_CheckIsArray)(field); - upb_Array* cloned_array = - upb_Array_DeepClone(array, upb_MiniTableField_CType(field), - upb_MiniTableField_CType(field) == kUpb_CType_Message - ? upb_MiniTable_GetSubMessageTable(field) - : NULL, - arena); + upb_Array* cloned_array = upb_Array_DeepCloneWithDepth( + array, upb_MiniTableField_CType(field), + upb_MiniTableField_CType(field) == kUpb_CType_Message + ? upb_MiniTable_GetSubMessageTable(field) + : NULL, + arena, depth); if (!cloned_array) return false; // Clear out upb_Array* due to parent memcpy. @@ -232,20 +251,22 @@ static bool upb_Clone_ExtensionValue(const upb_MiniTableField* field, const upb_Extension* source, - upb_Extension* dest, upb_Arena* arena) { + upb_Extension* dest, upb_Arena* arena, + int depth) { upb_MessageValue val = upb_Extension_GetValue(source); if (!upb_Clone_MessageValue(&val, upb_MiniTableField_CType(field), - upb_MiniTable_SubMessage(field), arena)) { + upb_MiniTable_SubMessage(field), arena, depth)) { return false; } upb_Extension_SetValue(dest, val); return true; } -upb_Message* _upb_Message_Copy(upb_Message* dst, const upb_Message* src, - const upb_MiniTable* mini_table, - upb_Arena* arena) { +static upb_Message* _upb_Message_Copy(upb_Message* dst, const upb_Message* src, + const upb_MiniTable* mini_table, + upb_Arena* arena, int depth) { UPB_ASSERT(!upb_Message_IsFrozen(dst)); + UPB_ASSERT(depth >= 0); upb_StringView empty_string = upb_StringView_FromDataAndSize(NULL, 0); // Only copy message area skipping upb_Message_Internal. memcpy(dst + 1, src + 1, mini_table->UPB_PRIVATE(size) - sizeof(upb_Message)); @@ -259,8 +280,8 @@ if (sub_message != NULL) { const upb_MiniTable* sub_message_table = upb_MiniTable_GetSubMessageTable(field); - upb_Message* dst_sub_message = - upb_Message_DeepClone(sub_message, sub_message_table, arena); + upb_Message* dst_sub_message = upb_Message_DeepCloneWithDepth( + sub_message, sub_message_table, arena, depth - 1); if (dst_sub_message == NULL) goto err; upb_Message_SetBaseFieldMessage(dst, field, dst_sub_message); @@ -284,15 +305,16 @@ if (upb_MiniTableField_IsMap(field)) { const upb_Map* map = upb_Message_GetMap(src, field); if (map != NULL) { - if (!upb_Message_Map_DeepClone(map, mini_table, field, dst, arena)) { + if (!upb_Message_Map_DeepClone(map, mini_table, field, dst, arena, + depth)) { goto err; } } } else { const upb_Array* array = upb_Message_GetArray(src, field); if (array != NULL) { - if (!upb_Message_Array_DeepClone(array, mini_table, field, dst, - arena)) { + if (!upb_Message_Array_DeepClone(array, mini_table, field, dst, arena, + depth)) { goto err; } } @@ -314,8 +336,16 @@ // destination arena and store it as unknown bytes, preserving order. const upb_Extension* msg_ext = upb_TaggedAuxPtr_NonCanonicalExtension(tagged_ptr); + // `upb_EncodeOptions_MaxDepth(0)` selects the default limit, so a zero + // budget must be enforced here: no further sub-message nesting fits. + if (depth == 0 && upb_MiniTableField_IsSubMessage( + upb_Extension_MiniTableField(msg_ext))) { + goto err; + } upb_StringView encoded; - if (upb_EncodeExtension(msg_ext, arena, &encoded, /*encode_options=*/0) != + if (upb_EncodeExtension( + msg_ext, arena, &encoded, + upb_EncodeOptions_MaxDepth((uint16_t)(depth > 0 ? depth : 1))) != kUpb_EncodeStatus_Ok) { goto err; } @@ -333,7 +363,7 @@ if (!dst_ext) goto err; if (upb_MiniTableField_IsScalar(field)) { - if (!upb_Clone_ExtensionValue(field, msg_ext, dst_ext, arena)) { + if (!upb_Clone_ExtensionValue(field, msg_ext, dst_ext, arena, depth)) { goto err; } } else { @@ -341,9 +371,9 @@ UPB_ASSERT(upb_MiniTableField_IsArray(field)); const upb_Array* msg_array = upb_Extension_GetArray(msg_ext); UPB_ASSERT(msg_array); - upb_Array* cloned_array = - upb_Array_DeepClone(msg_array, upb_MiniTableField_CType(field), - upb_MiniTable_SubMessage(field), arena); + upb_Array* cloned_array = upb_Array_DeepCloneWithDepth( + msg_array, upb_MiniTableField_CType(field), + upb_MiniTable_SubMessage(field), arena, depth); if (!cloned_array) goto err; upb_Extension_SetArray(dst_ext, cloned_array); @@ -365,11 +395,22 @@ return NULL; } +static upb_Message* upb_Message_DeepCloneWithDepth(const upb_Message* msg, + const upb_MiniTable* m, + upb_Arena* arena, + int depth) { + if (depth < 0) return NULL; + upb_Message* clone = upb_Message_New(m, arena); + if (!clone) return NULL; + return _upb_Message_Copy(clone, msg, m, arena, depth); +} + bool upb_Message_DeepCopy(upb_Message* dst, const upb_Message* src, const upb_MiniTable* mini_table, upb_Arena* arena) { UPB_ASSERT(!upb_Message_IsFrozen(dst)); upb_Message_Clear(dst, mini_table); - return _upb_Message_Copy(dst, src, mini_table, arena) != NULL; + return _upb_Message_Copy(dst, src, mini_table, arena, + kUpb_WireFormat_DefaultDepthLimit) != NULL; } // Deep clones a message using the provided target arena. @@ -377,9 +418,8 @@ // Returns NULL on failure. upb_Message* upb_Message_DeepClone(const upb_Message* msg, const upb_MiniTable* m, upb_Arena* arena) { - upb_Message* clone = upb_Message_New(m, arena); - if (!clone) return NULL; - return _upb_Message_Copy(clone, msg, m, arena); + return upb_Message_DeepCloneWithDepth(msg, m, arena, + kUpb_WireFormat_DefaultDepthLimit); } // Performs a shallow copy.
diff --git a/upb/message/copy.h b/upb/message/copy.h index 4a356fa..df49b97 100644 --- a/upb/message/copy.h +++ b/upb/message/copy.h
@@ -31,7 +31,8 @@ // serialized into unknown fields, since their minitables may belong to a schema // that does not outlive the result. // -// Returns NULL on failure. +// Returns NULL on failure, including when sub-message nesting exceeds the +// default wire format depth limit (100). UPB_NODISCARD upb_Message* upb_Message_DeepClone(const upb_Message* msg, const upb_MiniTable* m, upb_Arena* arena); @@ -43,13 +44,15 @@ const upb_MiniTable* m, upb_Arena* arena); -// Deep clones array contents. +// Deep clones array contents. Returns NULL on failure, including when +// sub-message nesting exceeds the default wire format depth limit (100). UPB_NODISCARD upb_Array* upb_Array_DeepClone(const upb_Array* array, upb_CType value_type, const upb_MiniTable* sub, upb_Arena* arena); -// Deep clones map contents. +// Deep clones map contents. Returns NULL on failure, including when +// sub-message nesting exceeds the default wire format depth limit (100). UPB_NODISCARD upb_Map* upb_Map_DeepClone(const upb_Map* map, upb_CType key_type, upb_CType value_type, const upb_MiniTable* map_entry_table, @@ -64,7 +67,8 @@ // extensions are serialized into unknown fields, since their minitables may // belong to a schema that does not outlive `dst`. // -// Returns false on failure, in which case `dst` is cleared. +// Returns false on failure, including when sub-message nesting exceeds the +// default wire format depth limit (100), in which case `dst` is cleared. UPB_NODISCARD bool upb_Message_DeepCopy(upb_Message* dst, const upb_Message* src, const upb_MiniTable* m,
diff --git a/upb/message/copy_test.cc b/upb/message/copy_test.cc index fd8c569..90ceefc 100644 --- a/upb/message/copy_test.cc +++ b/upb/message/copy_test.cc
@@ -23,6 +23,8 @@ #include <gtest/gtest.h> #include "google/protobuf/test_messages_proto2.upb.h" #include "google/protobuf/test_messages_proto2.upb_minitable.h" +#include "upb/base/descriptor_constants.h" +#include "upb/base/status.h" #include "upb/base/string_view.h" #include "upb/base/upcast.h" #include "upb/mem/arena.h" @@ -33,12 +35,17 @@ #include "upb/message/map.h" #include "upb/message/message.h" #include "upb/message/unknown_fields.h" +#include "upb/mini_descriptor/decode.h" +#include "upb/mini_descriptor/internal/encode.hpp" +#include "upb/mini_table/extension.h" #include "upb/mini_table/extension_registry.h" #include "upb/mini_table/field.h" #include "upb/mini_table/message.h" #include "upb/test/test.upb.h" #include "upb/test/test.upb_minitable.h" +#include "upb/wire/decode.h" #include "upb/wire/encode.h" +#include "upb/wire/internal/constants.h" // Must be last. #include "upb/port/def.inc" @@ -679,6 +686,163 @@ upb_Arena_Free(arena); } +// Returns a TestAllTypesProto2 with `depth` levels of `recursive_message` +// nested below it. If `bottom` is non-null, it is set to the innermost message. +protobuf_test_messages_proto2_TestAllTypesProto2* MakeRecursiveChain( + upb_Arena* arena, int depth, + protobuf_test_messages_proto2_TestAllTypesProto2** bottom = nullptr) { + protobuf_test_messages_proto2_TestAllTypesProto2* top = + protobuf_test_messages_proto2_TestAllTypesProto2_new(arena); + protobuf_test_messages_proto2_TestAllTypesProto2* msg = top; + for (int i = 0; i < depth; ++i) { + msg = + protobuf_test_messages_proto2_TestAllTypesProto2_mutable_recursive_message( + msg, arena); + } + protobuf_test_messages_proto2_TestAllTypesProto2_set_optional_int32(msg, 1); + if (bottom != nullptr) *bottom = msg; + return top; +} + +bool EncodesWithDefaultDepthLimit(const upb_Message* msg, + const upb_MiniTable* mt, upb_Arena* arena) { + char* buf; + size_t size; + return upb_Encode(msg, mt, /*options=*/0, arena, &buf, &size) == + kUpb_EncodeStatus_Ok; +} + +bool DecodesWithDefaultDepthLimit(const upb_Message* msg, + const upb_MiniTable* mt, upb_Arena* arena) { + char* buf; + size_t size; + if (upb_Encode( + msg, mt, + upb_EncodeOptions_MaxDepth(kUpb_WireFormat_DefaultDepthLimit + 10), + arena, &buf, &size) != kUpb_EncodeStatus_Ok) { + ADD_FAILURE() << "Encoding with a generous depth limit failed"; + return false; + } + upb_Message* decoded = upb_Message_New(mt, arena); + return decoded != nullptr && + upb_Decode(buf, size, decoded, mt, nullptr, /*options=*/0, arena) == + kUpb_DecodeStatus_Ok; +} + +TEST(GeneratedCode, DeepCloneRespectsDepthLimit) { + const upb_MiniTable* mt = + &protobuf_0test_0messages__proto2__TestAllTypesProto2_msg_init; + for (int depth : {99, 100, 101, 102}) { + upb_Arena* arena = upb_Arena_New(); + protobuf_test_messages_proto2_TestAllTypesProto2* msg = + MakeRecursiveChain(arena, depth); + + bool cloned = upb_Message_DeepClone(UPB_UPCAST(msg), mt, arena) != nullptr; + EXPECT_EQ(cloned, depth <= kUpb_WireFormat_DefaultDepthLimit) << depth; + // Deep clone enforces the same limit as the decoder and encoder. + EXPECT_EQ(cloned, DecodesWithDefaultDepthLimit(UPB_UPCAST(msg), mt, arena)) + << depth; + EXPECT_EQ(cloned, EncodesWithDefaultDepthLimit(UPB_UPCAST(msg), mt, arena)) + << depth; + + upb_Message* dst = upb_Message_New(mt, arena); + ASSERT_NE(dst, nullptr); + EXPECT_EQ(upb_Message_DeepCopy(dst, UPB_UPCAST(msg), mt, arena), cloned) + << depth; + + upb_Arena_Free(arena); + } +} + +// A non-canonical extension is serialized during deep clone. The nesting +// inside it must count against the depth already consumed by the enclosing +// messages, not start over with a fresh limit. +TEST(GeneratedCode, DeepCloneNonCanonicalExtensionRespectsDepthLimit) { + const upb_MiniTable* mt = + &protobuf_0test_0messages__proto2__TestAllTypesProto2_msg_init; + upb_Arena* schema_arena = upb_Arena_New(); + // A message-typed extension whose value is itself a TestAllTypesProto2. + upb::MtDataEncoder enc; + ASSERT_TRUE(enc.EncodeExtension(kUpb_FieldType_Message, 150, 0)); + upb_Status status; + upb_Status_Clear(&status); + upb_MiniTableExtension* ext = upb_MiniTableExtension_BuildMessage( + enc.data().data(), enc.data().size(), mt, mt, schema_arena, &status); + ASSERT_NE(ext, nullptr) << upb_Status_ErrorMessage(&status); + + // The non-canonical extension sits `kOuterDepth` levels below the top, and + // its value has `value_depth` more levels below it, so the total nesting is + // kOuterDepth + 1 + value_depth. With a fresh limit for the extension, only + // 1 + value_depth would count, and every case below would succeed. + constexpr int kOuterDepth = 50; + for (int value_depth : {48, 49, 50, 51}) { + upb_Arena* arena = upb_Arena_New(); + protobuf_test_messages_proto2_TestAllTypesProto2* bottom; + protobuf_test_messages_proto2_TestAllTypesProto2* top = + MakeRecursiveChain(arena, kOuterDepth, &bottom); + protobuf_test_messages_proto2_TestAllTypesProto2* value = + MakeRecursiveChain(arena, value_depth); + ASSERT_TRUE(UPB_PRIVATE(_upb_Message_SetNonCanonicalExtension)( + UPB_UPCAST(bottom), ext, &value, arena)); + + int total_depth = kOuterDepth + 1 + value_depth; + bool cloned = upb_Message_DeepClone(UPB_UPCAST(top), mt, arena) != nullptr; + EXPECT_EQ(cloned, total_depth <= kUpb_WireFormat_DefaultDepthLimit) + << value_depth; + EXPECT_EQ(cloned, EncodesWithDefaultDepthLimit(UPB_UPCAST(top), mt, arena)) + << value_depth; + + upb_Arena_Free(arena); + } + upb_Arena_Free(schema_arena); +} + +// A non-canonical extension on a message that is already at the depth limit: +// a scalar value still fits, but a message value would exceed the limit. +TEST(GeneratedCode, DeepCloneNonCanonicalExtensionAtDepthLimit) { + const upb_MiniTable* mt = + &protobuf_0test_0messages__proto2__TestAllTypesProto2_msg_init; + upb_Arena* schema_arena = upb_Arena_New(); + upb_Status status; + upb_Status_Clear(&status); + upb::MtDataEncoder msg_enc; + ASSERT_TRUE(msg_enc.EncodeExtension(kUpb_FieldType_Message, 150, 0)); + upb_MiniTableExtension* msg_ext = upb_MiniTableExtension_BuildMessage( + msg_enc.data().data(), msg_enc.data().size(), mt, mt, schema_arena, + &status); + ASSERT_NE(msg_ext, nullptr) << upb_Status_ErrorMessage(&status); + upb::MtDataEncoder int_enc; + ASSERT_TRUE(int_enc.EncodeExtension(kUpb_FieldType_Int32, 151, 0)); + upb_MiniTableExtension* int_ext = upb_MiniTableExtension_Build( + int_enc.data().data(), int_enc.data().size(), mt, schema_arena, &status); + ASSERT_NE(int_ext, nullptr) << upb_Status_ErrorMessage(&status); + + for (bool message_value : {false, true}) { + SCOPED_TRACE(message_value ? "message value" : "scalar value"); + upb_Arena* arena = upb_Arena_New(); + protobuf_test_messages_proto2_TestAllTypesProto2* bottom; + protobuf_test_messages_proto2_TestAllTypesProto2* top = + MakeRecursiveChain(arena, kUpb_WireFormat_DefaultDepthLimit, &bottom); + if (message_value) { + protobuf_test_messages_proto2_TestAllTypesProto2* value = + protobuf_test_messages_proto2_TestAllTypesProto2_new(arena); + ASSERT_TRUE(UPB_PRIVATE(_upb_Message_SetNonCanonicalExtension)( + UPB_UPCAST(bottom), msg_ext, &value, arena)); + } else { + int32_t value = 42; + ASSERT_TRUE(UPB_PRIVATE(_upb_Message_SetNonCanonicalExtension)( + UPB_UPCAST(bottom), int_ext, &value, arena)); + } + + bool cloned = upb_Message_DeepClone(UPB_UPCAST(top), mt, arena) != nullptr; + EXPECT_EQ(cloned, !message_value); + EXPECT_EQ(cloned, EncodesWithDefaultDepthLimit(UPB_UPCAST(top), mt, arena)); + + upb_Arena_Free(arena); + } + upb_Arena_Free(schema_arena); +} + } // namespace #include "upb/port/undef.inc"
diff --git a/upb/wire/encode_test.cc b/upb/wire/encode_test.cc index abcdfa9..785d0c5 100644 --- a/upb/wire/encode_test.cc +++ b/upb/wire/encode_test.cc
@@ -6,6 +6,7 @@ #include <cstdint> #include <gtest/gtest.h> +#include "upb/base/upcast.h" #include "upb/mem/arena.h" #include "upb/message/array.h" #include "upb/message/internal/accessors.h" @@ -22,6 +23,7 @@ #include "upb/wire/encode_test.upb.h" #include "upb/wire/encode_test.upb_minitable.h" #include "upb/wire/internal/back_alloc.h" +#include "upb/wire/internal/constants.h" #include "upb/wire/internal/encoder.h" // Must be last. @@ -120,8 +122,11 @@ jmp_buf err; UPB_PRIVATE(_upb_encstate_init)(&e, &err, arena); + // Two levels of nesting below `msg` exceed a depth limit of 1. upb_wire_test_TestRecursive* sub_msg = upb_wire_test_TestRecursive_new(arena); upb_wire_test_TestRecursive_set_recursive(msg, sub_msg); + upb_wire_test_TestRecursive_set_recursive( + sub_msg, upb_wire_test_TestRecursive_new(arena)); const upb_MiniTable* mt = &upb_0wire_0test__TestRecursive_msg_init; const upb_MiniTableField* field = upb_MiniTable_FindFieldByNumber(mt, 1); @@ -134,6 +139,67 @@ upb_Arena_Free(arena); } +// Returns a TestRecursive with `depth` levels of `recursive` nested below it. +upb_wire_test_TestRecursive* MakeRecursiveChain(upb_Arena* arena, int depth) { + upb_wire_test_TestRecursive* top = upb_wire_test_TestRecursive_new(arena); + upb_wire_test_TestRecursive* msg = top; + for (int i = 0; i < depth; ++i) { + upb_wire_test_TestRecursive* child = upb_wire_test_TestRecursive_new(arena); + upb_wire_test_TestRecursive_set_recursive(msg, child); + msg = child; + } + return top; +} + +// The encoder and decoder must agree on the depth limit, so that any message +// the decoder accepts can be re-encoded with the same limit. +TEST(EncodeTest, DepthLimitMatchesDecoder) { + const upb_MiniTable* mt = &upb_0wire_0test__TestRecursive_msg_init; + // 0 selects the default limit. + for (int limit : {0, 1, 2, 10}) { + int effective_limit = + limit == 0 ? kUpb_WireFormat_DefaultDepthLimit : limit; + for (int depth : + {effective_limit - 1, effective_limit, effective_limit + 1}) { + SCOPED_TRACE(testing::Message() + << "limit=" << limit << " depth=" << depth); + upb_Arena* arena = upb_Arena_New(); + upb_wire_test_TestRecursive* msg = MakeRecursiveChain(arena, depth); + bool within_limit = depth <= effective_limit; + + char* buf; + size_t size; + EXPECT_EQ( + upb_Encode(UPB_UPCAST(msg), mt, upb_EncodeOptions_MaxDepth(limit), + arena, &buf, &size), + within_limit ? kUpb_EncodeStatus_Ok + : kUpb_EncodeStatus_MaxDepthExceeded); + + // Serialize with a generous limit, then decode with `limit`. + ASSERT_EQ(upb_Encode(UPB_UPCAST(msg), mt, + upb_EncodeOptions_MaxDepth(effective_limit + 10), + arena, &buf, &size), + kUpb_EncodeStatus_Ok); + upb_Message* decoded = upb_Message_New(mt, arena); + ASSERT_NE(decoded, nullptr); + EXPECT_EQ(upb_Decode(buf, size, decoded, mt, nullptr, + upb_DecodeOptions_MaxDepth(limit), arena), + within_limit ? kUpb_DecodeStatus_Ok + : kUpb_DecodeStatus_MaxDepthExceeded); + + // A message that decodes with `limit` re-encodes with the same limit. + if (within_limit) { + char* reencoded; + size_t reencoded_size; + EXPECT_EQ(upb_Encode(decoded, mt, upb_EncodeOptions_MaxDepth(limit), + arena, &reencoded, &reencoded_size), + kUpb_EncodeStatus_Ok); + } + upb_Arena_Free(arena); + } + } +} + TEST(EncodeTest, EncodeExtensionSuccess) { upb_Arena* arena = upb_Arena_New(); @@ -202,8 +268,12 @@ UPB_PRIVATE(_upb_encstate_init)(&e, &err, arena); const upb_MiniTableExtension* ext = upb_wire_test_ext_recursive_ext; + // The extension value plus one level of nesting exceed a depth limit of 1. + upb_wire_test_TestRecursive* value = upb_wire_test_TestRecursive_new(arena); + upb_wire_test_TestRecursive_set_recursive( + value, upb_wire_test_TestRecursive_new(arena)); upb_MessageValue ext_val; - ext_val.msg_val = (upb_Message*)upb_wire_test_TestRecursive_new(arena); + ext_val.msg_val = (upb_Message*)value; char* buf = e.alloc.limit; size_t size;
diff --git a/upb/wire/internal/encoder.c b/upb/wire/internal/encoder.c index 316909f..6026517 100644 --- a/upb/wire/internal/encoder.c +++ b/upb/wire/internal/encoder.c
@@ -423,7 +423,7 @@ if (submsg == 0) { return ptr; } - if (--e->depth == 0) encode_err(e, kUpb_EncodeStatus_MaxDepthExceeded); + if (--e->depth < 0) encode_err(e, kUpb_EncodeStatus_MaxDepthExceeded); ptr = encode_tag(ptr, e, upb_MiniTableField_Number(f), kUpb_WireType_EndGroup); ptr = encode_message(ptr, e, submsg, subm, &size); @@ -438,7 +438,7 @@ if (submsg == 0) { return ptr; } - if (--e->depth == 0) encode_err(e, kUpb_EncodeStatus_MaxDepthExceeded); + if (--e->depth < 0) encode_err(e, kUpb_EncodeStatus_MaxDepthExceeded); ptr = encode_message(ptr, e, submsg, subm, &size); e->depth++; size_t max_size = UPB_PB_VARINT32_MAX_LEN + UPB_PB_VARINT32_MAX_LEN; @@ -536,7 +536,7 @@ const upb_Message* const* start = upb_Array_DataPtr(arr); const upb_Message* const* arr_ptr = start + upb_Array_Size(arr); const upb_MiniTable* subm = upb_MiniTable_GetSubMessageTable(f); - if (--e->depth == 0) encode_err(e, kUpb_EncodeStatus_MaxDepthExceeded); + if (--e->depth < 0) encode_err(e, kUpb_EncodeStatus_MaxDepthExceeded); do { size_t size; arr_ptr--; @@ -553,7 +553,7 @@ const upb_Message* const* start = upb_Array_DataPtr(arr); const upb_Message* const* arr_ptr = start + upb_Array_Size(arr); const upb_MiniTable* subm = upb_MiniTable_GetSubMessageTable(f); - if (--e->depth == 0) encode_err(e, kUpb_EncodeStatus_MaxDepthExceeded); + if (--e->depth < 0) encode_err(e, kUpb_EncodeStatus_MaxDepthExceeded); do { size_t size; arr_ptr--;