Fix crashes with unproper cleaning of repeated yield from

Author: bwoebi

Zend/tests/generators/repeated_yield_from_with_immediate_release.phpt

@@ -0,0 +1,19 @@
+--TEST--
+A generator can be yielded from multiple times, testing immediate release of the yield from'ing generator
+--FILE--
+<?php
+
+function gen() {
+	yield 42;
+}
+function yield_from($gen) {
+	yield from $gen;
+}
+$gen = gen();
+var_dump(yield_from($gen)->current());
+var_dump(yield_from($gen)->current());
+
+?>
+--EXPECT--
+int(42)
+int(42)

Zend/zend_generators.c

@@ -163,6 +163,52 @@ ZEND_API void zend_generator_close(zend_generator *generator, zend_bool finished
 
 static zend_generator *zend_generator_get_child(zend_generator_node *node, zend_generator *leaf);
 
+static void zend_generator_update_leaf_of_child(zend_generator_node *node, zend_generator *from_leaf, zend_generator *to_leaf)
+{
+	if (node->ptr.leaf == from_leaf) {
+		node->ptr.leaf = to_leaf;
+	}
+	if (node->children <= 1) {
+		node->child.single.leaf = to_leaf;
+	} else {
+		HashTable *ht = node->child.ht;
+		zend_generator *child = zend_hash_index_find_ptr(ht, (zend_ulong) from_leaf);
+		if (child) {
+			zend_hash_index_del(ht, (zend_ulong) from_leaf);
+			zend_hash_index_add_ptr(ht, (zend_ulong) to_leaf, child);
+		}
+	}
+}
+
+static void zend_generator_remove_leaf_child(zend_generator_node *node, zend_generator *leaf, zend_generator *replace_leaf) {
+	if (node->children > 1) {
+		HashTable *ht = node->child.ht;
+		zend_ulong child_leaf;
+		zend_generator *child_generator;
+		zend_hash_index_del(ht, (zend_ulong) leaf);
+		if (--node->children == 1) {
+			ZEND_HASH_FOREACH_NUM_KEY_PTR(ht, child_leaf, child_generator) {
+				node->child.single.leaf = (zend_generator *) child_leaf;
+				node->child.single.child = child_generator;
+				if (node->ptr.leaf == leaf) {
+					node->ptr.leaf = (zend_generator *) child_leaf;
+				}
+				break;
+			} ZEND_HASH_FOREACH_END();
+			zend_hash_destroy(ht);
+			efree(ht);
+		} else if (node->ptr.leaf == leaf) {
+			ZEND_HASH_FOREACH_NUM_KEY_PTR(ht, child_leaf, child_generator) {
+				node->ptr.leaf = (zend_generator *) child_leaf;
+				break;
+			} ZEND_HASH_FOREACH_END();
+		}
+	} else if (node->ptr.leaf == leaf) {
+		ZEND_ASSERT(replace_leaf != leaf);
+		node->ptr.leaf = replace_leaf;
+	}
+}
+
 static void zend_generator_dtor_storage(zend_object *object) /* {{{ */
 {
 	zend_generator *generator = (zend_generator*) object;
@@ -177,14 +223,43 @@ static void zend_generator_dtor_storage(zend_object *object) /* {{{ */
 	}
 
 	if (EXPECTED(generator->node.children == 0)) {
-		zend_generator *root = generator->node.ptr.root, *next;
-		while (UNEXPECTED(root != generator)) {
-			next = zend_generator_get_child(&root->node, generator);
-			generator->node.ptr.root = next;
-			next->node.parent = NULL;
-			OBJ_RELEASE(&root->std);
-			root = next;
+run_normal_destruction:
+		zend_generator_update_current(generator, generator); /* ensure we remove it from a *live* root */
+		zend_generator *root = generator->node.ptr.root, *parent = generator->node.parent, *next, *toproot = root;
+		if (parent) {
+			zend_bool parent_becomes_leaf = parent->node.children == 1;
+			while (UNEXPECTED(root != generator)) {
+				next = zend_generator_get_child(&root->node, generator);
+				if (parent_becomes_leaf) {
+					zend_generator_update_leaf_of_child(&root->node, generator, parent);
+				} else {
+					zend_generator_remove_leaf_child(&root->node, generator, parent->node.ptr.leaf);
+					OBJ_RELEASE(&root->std);
+				}
+				root = next;
+			}
+			if (parent_becomes_leaf) {
+				parent->node.ptr.root = toproot;
+				parent->node.children = 0;
+				OBJ_RELEASE(&parent->std);
+			}
+			/* Reset for resuming in finally */
+			generator->node.parent = NULL;
+			generator->node.ptr.root = generator; 
 		}
+	} else if (generator->node.parent) {
+		/* we're called out of order - this must only happen during shutdown sequence: we call our (direct) child nodes destructors first, to clean it from the bottom up */
+		while (generator->node.children != 0) {
+			zend_generator *child;
+			if (generator->node.children == 1) {
+				child = generator->node.child.single.child;
+			} else {
+				child = (zend_generator *) Z_PTR_P(zend_hash_get_current_data(generator->node.child.ht));
+			}
+			GC_ADD_FLAGS(&child->std, IS_OBJ_DESTRUCTOR_CALLED);
+			zend_generator_dtor_storage(&child->std);
+		}
+		goto run_normal_destruction;
 	}
 
 	if (EXPECTED(!ex) || EXPECTED(!(ex->func->op_array.fn_flags & ZEND_ACC_HAS_FINALLY_BLOCK))
@@ -465,10 +540,13 @@ static void zend_generator_add_single_child(zend_generator_node *node, zend_gene
 			node->child.ht = ht;
 		}
 
-		zend_hash_index_add_ptr(node->child.ht, (zend_ulong) leaf, child);
+		if (zend_hash_index_add_ptr(node->child.ht, (zend_ulong) leaf, child) == NULL) {
+			ZEND_ASSERT(node->children > 1);
+			return;
+		}
 	}
 
-	node->children++;
+	++node->children;
 }
 
 static void zend_generator_merge_child_nodes(zend_generator_node *dest, zend_generator_node *src, zend_generator *child)
@@ -507,7 +585,6 @@ static void zend_generator_add_child(zend_generator *generator, zend_generator *
 	} else if (generator->node.children == 1) {
 		multi_children_node = zend_generator_search_multi_children_node(&generator->node);
 		if (multi_children_node) {
-			generator->node.children = 0;
 			zend_generator_merge_child_nodes(&generator->node, multi_children_node, generator->node.child.single.child);
 		}
 	}
@@ -518,6 +595,7 @@ static void zend_generator_add_child(zend_generator *generator, zend_generator *
 		multi_children_node = (zend_generator_node *) 0x1;
 	}
 
+	/* for allowing zend_generator_get_child() to work, we need every multi children node to have ALL its leaf descendents present, linking to their respective child */
 	{
 		zend_generator *parent = generator->node.parent, *cur = generator;
 
@@ -574,7 +652,7 @@ ZEND_API zend_generator *zend_generator_update_current(zend_generator *generator
 
 	if (root->node.parent) {
 		if (root->node.parent->execute_data == NULL) {
-			if (EXPECTED(EG(exception) == NULL)) {
+			if (EXPECTED(EG(exception) == NULL) && EXPECTED((OBJ_FLAGS(&generator->std) & IS_OBJ_DESTRUCTOR_CALLED) == 0)) {
 				zend_op *yield_from = (zend_op *) root->execute_data->opline - 1;
 
 				if (yield_from->opcode == ZEND_YIELD_FROM) {