Fix crash when forging containers into a short buffer

2 files changed, 137 insertions, 6 deletions

diff --git a/lv2/atom/forge-overflow-test.c b/lv2/atom/forge-overflow-test.c
index 353ad79..a580d69 100644
--- a/lv2/atom/forge-overflow-test.c
+++ b/lv2/atom/forge-overflow-test.c
@@ -20,6 +20,7 @@
 #include "lv2/atom/util.h"
 #include "lv2/urid/urid.h"
 
+#include <assert.h>
 #include <stdint.h>
 #include <stdlib.h>
 #include <string.h>
@@ -29,8 +30,8 @@ test_string_overflow(void)
 {
 #define MAX_CHARS 15
 
-	static const size_t capacity  = sizeof(LV2_Atom_String) + MAX_CHARS + 1;
-	static const char*  str       = "ABCDEFGHIJKLMNOPQRSTUVWXYZ";
+	static const size_t capacity = sizeof(LV2_Atom_String) + MAX_CHARS + 1;
+	static const char*  str      = "ABCDEFGHIJKLMNOPQRSTUVWXYZ";
 
 	uint8_t*       buf = (uint8_t*)malloc(capacity);
 	LV2_URID_Map   map = { NULL, urid_map };
@@ -103,10 +104,132 @@ test_literal_overflow(void)
 	return 0;
 }
 
+static int
+test_sequence_overflow(void)
+{
+	static const size_t size = sizeof(LV2_Atom_Sequence) + 6 * sizeof(LV2_Atom);
+	LV2_URID_Map        map  = { NULL, urid_map };
+
+	// Test over a range that fails in the sequence header and event components
+	for (size_t capacity = 1; capacity < size; ++capacity) {
+		uint8_t* buf = (uint8_t*)malloc(capacity);
+
+		LV2_Atom_Forge forge;
+		lv2_atom_forge_init(&forge, &map);
+		lv2_atom_forge_set_buffer(&forge, buf, capacity);
+
+		LV2_Atom_Forge_Frame frame;
+		LV2_Atom_Forge_Ref   ref =
+		        lv2_atom_forge_sequence_head(&forge, &frame, 0);
+
+		assert(capacity >= sizeof(LV2_Atom_Sequence) || !frame.ref);
+		assert(capacity >= sizeof(LV2_Atom_Sequence) || !ref);
+
+		lv2_atom_forge_frame_time(&forge, 0);
+		lv2_atom_forge_int(&forge, 42);
+		lv2_atom_forge_pop(&forge, &frame);
+
+		free(buf);
+	}
+
+	return 0;
+}
+
+static int
+test_vector_head_overflow(void)
+{
+	static const size_t size = sizeof(LV2_Atom_Vector) + 3 * sizeof(LV2_Atom);
+	LV2_URID_Map        map  = { NULL, urid_map };
+
+	// Test over a range that fails in the vector header and elements
+	for (size_t capacity = 1; capacity < size; ++capacity) {
+		uint8_t* buf = (uint8_t*)malloc(capacity);
+
+		LV2_Atom_Forge forge;
+		lv2_atom_forge_init(&forge, &map);
+		lv2_atom_forge_set_buffer(&forge, buf, capacity);
+
+		LV2_Atom_Forge_Frame frame;
+		LV2_Atom_Forge_Ref   ref = lv2_atom_forge_vector_head(
+                &forge, &frame, sizeof(int32_t), forge.Int);
+
+		assert(capacity >= sizeof(LV2_Atom_Vector) || !frame.ref);
+		assert(capacity >= sizeof(LV2_Atom_Vector) || !ref);
+
+		lv2_atom_forge_int(&forge, 1);
+		lv2_atom_forge_int(&forge, 2);
+		lv2_atom_forge_int(&forge, 3);
+		lv2_atom_forge_pop(&forge, &frame);
+
+		free(buf);
+	}
+
+	return 0;
+}
+
+static int
+test_vector_overflow(void)
+{
+	static const size_t  size  = sizeof(LV2_Atom_Vector) + 3 * sizeof(LV2_Atom);
+	static const int32_t vec[] = { 1, 2, 3 };
+	LV2_URID_Map         map   = { NULL, urid_map };
+
+	// Test over a range that fails in the vector header and elements
+	for (size_t capacity = 1; capacity < size; ++capacity) {
+		uint8_t* buf = (uint8_t*)malloc(capacity);
+
+		LV2_Atom_Forge forge;
+		lv2_atom_forge_init(&forge, &map);
+		lv2_atom_forge_set_buffer(&forge, buf, capacity);
+
+		LV2_Atom_Forge_Ref ref = lv2_atom_forge_vector(
+		        &forge, sizeof(int32_t), forge.Int, 3, vec);
+
+		assert(capacity >= sizeof(LV2_Atom_Vector) || !ref);
+
+		free(buf);
+	}
+
+	return 0;
+}
+
+static int
+test_tuple_overflow(void)
+{
+	static const size_t size = sizeof(LV2_Atom_Tuple) + 3 * sizeof(LV2_Atom);
+	LV2_URID_Map        map  = { NULL, urid_map };
+
+	// Test over a range that fails in the tuple header and elements
+	for (size_t capacity = 1; capacity < size; ++capacity) {
+		uint8_t* buf = (uint8_t*)malloc(capacity);
+
+		LV2_Atom_Forge forge;
+		lv2_atom_forge_init(&forge, &map);
+		lv2_atom_forge_set_buffer(&forge, buf, capacity);
+
+		LV2_Atom_Forge_Frame frame;
+		LV2_Atom_Forge_Ref   ref = lv2_atom_forge_tuple(&forge, &frame);
+
+		assert(capacity >= sizeof(LV2_Atom_Tuple) || !frame.ref);
+		assert(capacity >= sizeof(LV2_Atom_Tuple) || !ref);
+
+		lv2_atom_forge_int(&forge, 1);
+		lv2_atom_forge_float(&forge, 2.0f);
+		lv2_atom_forge_string(&forge, "three", 5);
+		lv2_atom_forge_pop(&forge, &frame);
+
+		free(buf);
+	}
+
+	return 0;
+}
+
 int
 main(void)
 {
-	const int ret = test_string_overflow() || test_literal_overflow();
+	const int ret = test_string_overflow() || test_literal_overflow() ||
+	                test_sequence_overflow() || test_vector_head_overflow() ||
+	                test_vector_overflow() || test_tuple_overflow();
 
 	free_urid_map();
 
diff --git a/lv2/atom/forge.h b/lv2/atom/forge.h
index e866b38..5bcd816 100644
--- a/lv2/atom/forge.h
+++ b/lv2/atom/forge.h
@@ -176,7 +176,11 @@ lv2_atom_forge_push(LV2_Atom_Forge*       forge,
 {
 	frame->parent = forge->stack;
 	frame->ref    = ref;
-	forge->stack  = frame;
+
+	if (ref) {
+		forge->stack = frame; // Don't push, so walking the stack is always safe
+	}
+
 	return ref;
 }
 
@@ -184,8 +188,12 @@ lv2_atom_forge_push(LV2_Atom_Forge*       forge,
 static inline void
 lv2_atom_forge_pop(LV2_Atom_Forge* forge, LV2_Atom_Forge_Frame* frame)
 {
-	assert(frame == forge->stack);
-	forge->stack = frame->parent;
+	if (frame->ref) {
+		// If frame has a valid ref, it must be the top of the stack
+		assert(frame == forge->stack);
+		forge->stack = frame->parent;
+	}
+	// Otherwise, frame was not pushed because of overflow, do nothing
 }
 
 /** Return true iff the top of the stack has the given type. */