upb: copy the wire decode recursion-depth-checking code to the wire encoder

PiperOrigin-RevId: 524873082

python/message.c

@@ -1528,7 +1528,7 @@ PyObject* PyUpb_Message_SerializeInternal(PyObject* _self, PyObject* args,
   const upb_MiniTable* layout = upb_MessageDef_MiniTable(msgdef);
   size_t size = 0;
   // Python does not currently have any effective limit on serialization depth.
-  int options = UPB_ENCODE_MAXDEPTH(UINT32_MAX);
+  int options = upb_EncodeOptions_MaxDepth(UINT16_MAX);
   if (check_required) options |= kUpb_EncodeOption_CheckRequired;
   if (deterministic) options |= kUpb_EncodeOption_Deterministic;
   char* pb;

upb/message/test.cc

@@ -507,12 +507,17 @@ TEST(MessageTest, MapField) {
 // static void DecodeEncodeArbitrarySchemaAndPayload(
 //     const upb::fuzz::MiniTableFuzzInput& input, std::string_view proto_payload,
 //     int decode_options, int encode_options) {
+//   // The value of 80 used here is empirical and intended to roughly represent
+//   // the tiny 64K stack size used by the test framework. We still see the
+//   // occasional stack overflow at 90, so far 80 has worked 100% of the time.
+//   decode_options = upb_Decode_LimitDepth(decode_options, 80);
+//   encode_options = upb_Encode_LimitDepth(encode_options, 80);
+//
 //   upb::Arena arena;
 //   upb_ExtensionRegistry* exts;
 //   const upb_MiniTable* mini_table =
 //       upb::fuzz::BuildMiniTable(input, &exts, arena.ptr());
 //   if (!mini_table) return;
-//   decode_options = upb_Decode_LimitDepth(decode_options, 80);
 //   upb_Message* msg = upb_Message_New(mini_table, arena.ptr());
 //   upb_Decode(proto_payload.data(), proto_payload.size(), msg, mini_table, exts,
 //              decode_options, arena.ptr());

upb/wire/encode.h

@@ -48,24 +48,37 @@ enum {
    * memory during encode. */
   kUpb_EncodeOption_Deterministic = 1,
 
-  /* When set, unknown fields are not printed. */
+  // When set, unknown fields are not printed.
   kUpb_EncodeOption_SkipUnknown = 2,
 
-  /* When set, the encode will fail if any required fields are missing. */
+  // When set, the encode will fail if any required fields are missing.
   kUpb_EncodeOption_CheckRequired = 4,
 };
 
-#define UPB_ENCODE_MAXDEPTH(depth) ((depth) << 16)
-
 typedef enum {
   kUpb_EncodeStatus_Ok = 0,
   kUpb_EncodeStatus_OutOfMemory = 1,  // Arena alloc failed
-  kUpb_EncodeStatus_MaxDepthExceeded = 2,  // Exceeded UPB_ENCODE_MAXDEPTH
+  kUpb_EncodeStatus_MaxDepthExceeded = 2,
 
   // kUpb_EncodeOption_CheckRequired failed but the parse otherwise succeeded.
   kUpb_EncodeStatus_MissingRequired = 3,
 } upb_EncodeStatus;
 
+UPB_INLINE uint32_t upb_EncodeOptions_MaxDepth(uint16_t depth) {
+  return (uint32_t)depth << 16;
+}
+
+UPB_INLINE uint16_t upb_EncodeOptions_GetMaxDepth(uint32_t options) {
+  return options >> 16;
+}
+
+// Enforce an upper bound on recursion depth.
+UPB_INLINE int upb_Encode_LimitDepth(uint32_t encode_options, uint32_t limit) {
+  uint32_t max_depth = upb_EncodeOptions_GetMaxDepth(encode_options);
+  if (max_depth > limit) max_depth = limit;
+  return upb_EncodeOptions_MaxDepth(max_depth) | (encode_options & 0xffff);
+}
+
 upb_EncodeStatus upb_Encode(const void* msg, const upb_MiniTable* l,
                             int options, upb_Arena* arena, char** buf,
                             size_t* size);