From 9175989431ef264f38fb2ec4dd3bf6d199a530de Mon Sep 17 00:00:00 2001
From: Joshua Haberman <jhaberman@gmail.com>
Date: Sat, 9 Jan 2021 18:03:13 -0800
Subject: [PATCH] Bugfix for arena cleanup list when passing to upb_decode().

---
 tests/test_generated_code.c | 69 +++++++++++++++++++++++++++++++++++++
 upb/decode.c                |  2 ++
 2 files changed, 71 insertions(+)

diff --git a/tests/test_generated_code.c b/tests/test_generated_code.c
index bd1066b233..4b38599f0d 100644
--- a/tests/test_generated_code.c
+++ b/tests/test_generated_code.c
@@ -410,6 +410,73 @@ void test_status_truncation(void) {
   }
 }
 
+void decrement_int(void *ptr) {
+  int* iptr = ptr;
+  (*iptr)--;
+}
+
+void test_arena_fuse(void) {
+  int i1 = 5;
+  int i2 = 5;
+  int i3 = 5;
+  int i4 = 5;
+
+  upb_arena *arena1 = upb_arena_new();
+  upb_arena *arena2 = upb_arena_new();
+
+  upb_arena_addcleanup(arena1, &i1, decrement_int);
+  upb_arena_addcleanup(arena2, &i2, decrement_int);
+
+  upb_arena_fuse(arena1, arena2);
+
+  upb_arena_addcleanup(arena1, &i3, decrement_int);
+  upb_arena_addcleanup(arena2, &i4, decrement_int);
+
+  upb_arena_free(arena1);
+  ASSERT(i1 == 5);
+  ASSERT(i2 == 5);
+  ASSERT(i3 == 5);
+  ASSERT(i4 == 5);
+  upb_arena_free(arena2);
+  ASSERT(i1 == 4);
+  ASSERT(i2 == 4);
+  ASSERT(i3 == 4);
+  ASSERT(i4 == 4);
+}
+
+void test_arena_decode(void) {
+  // Tests against a bug that previously existed when passing an arena to
+  // upb_decode().
+  char large_string[1024] = {0};
+  upb_strview large_string_view = {large_string, sizeof(large_string)};
+  upb_arena *tmp = upb_arena_new();
+
+  protobuf_test_messages_proto3_TestAllTypesProto3 *msg =
+      protobuf_test_messages_proto3_TestAllTypesProto3_new(tmp);
+
+  protobuf_test_messages_proto3_TestAllTypesProto3_set_optional_bytes(
+      msg, large_string_view);
+
+  upb_strview serialized;
+  serialized.data = protobuf_test_messages_proto3_TestAllTypesProto3_serialize(
+      msg, tmp, &serialized.size);
+
+  upb_arena *arena = upb_arena_new();
+  // Parse the large payload, forcing an arena block to be allocated. This used
+  // to corrupt the cleanup list, preventing subsequent upb_arena_addcleanup()
+  // calls from working properly.
+  protobuf_test_messages_proto3_TestAllTypesProto3_parse(
+      serialized.data, serialized.size, arena);
+
+  int i1 = 5;
+  upb_arena_addcleanup(arena, &i1, decrement_int);
+  ASSERT(i1 == 5);
+  upb_arena_free(arena);
+  ASSERT(i1 == 4);
+
+  upb_arena_free(tmp);
+}
+
 int run_tests(int argc, char *argv[]) {
   test_scalars();
   test_utf8();
@@ -419,5 +486,7 @@ int run_tests(int argc, char *argv[]) {
   test_repeated();
   test_null_decode_buf();
   test_status_truncation();
+  test_arena_fuse();
+  test_arena_decode();
   return 0;
 }
diff --git a/upb/decode.c b/upb/decode.c
index a5f0666fa8..cc18d7dc17 100644
--- a/upb/decode.c
+++ b/upb/decode.c
@@ -679,6 +679,7 @@ bool _upb_decode(const char *buf, size_t size, void *msg,
   state.end_group = DECODE_NOGROUP;
   state.arena.head = arena->head;
   state.arena.last_size = arena->last_size;
+  state.arena.cleanups = arena->cleanups;
   state.arena.parent = arena;
 
   if (UPB_UNLIKELY(UPB_SETJMP(state.err))) {
@@ -692,6 +693,7 @@ bool _upb_decode(const char *buf, size_t size, void *msg,
 
   arena->head.ptr = state.arena.head.ptr;
   arena->head.end = state.arena.head.end;
+  arena->cleanups = state.arena.cleanups;
   return ok;
 }