From fd1898b70aa40b9bbbb20e98c731fc9ad10a4a89 Mon Sep 17 00:00:00 2001
From: Walther Chen <hwchen@users.noreply.github.com>
Date: Fri, 25 Oct 2024 19:57:00 -0400
Subject: [PATCH] copy out keys also in HashMap.copy_keys (#1569)

* copy out keys also in HashMap.copy_keys
* test for copying keys out
---
 lib/std/collections/hashmap.c3           |  6 +++++-
 releasenotes.md                          |  1 +
 test/unit/stdlib/collections/copy_map.c3 | 21 ++++++++++++++++++++-
 3 files changed, 26 insertions(+), 2 deletions(-)

diff --git a/lib/std/collections/hashmap.c3 b/lib/std/collections/hashmap.c3
index 8f0c446fd..c97a44b32 100644
--- a/lib/std/collections/hashmap.c3
+++ b/lib/std/collections/hashmap.c3
@@ -313,7 +313,11 @@ fn Key[] HashMap.copy_keys(&map, Allocator allocator = allocator::heap())
 	{
 		while (entry)
 		{
-			list[index++] = entry.key;
+			$if COPY_KEYS:
+				list[index++] = entry.key.copy(allocator);
+			$else
+				list[index++] = entry.key;
+			$endif
 			entry = entry.next;
 		}
 	}
diff --git a/releasenotes.md b/releasenotes.md
index d051cb396..c9611f0dc 100644
--- a/releasenotes.md
+++ b/releasenotes.md
@@ -46,6 +46,7 @@
 - `&self` argument not implicitly null checked. #1556.
 - `(uptr)&((Foo*)null).a` incorrectly inserts a null check. #1544
 - Incorrect error message when `$eval` is provided an invalid string. #1570
+- `HashMap.copy_keys` did not properly copy keys which needed to be allocated #1569
 
 ### Stdlib changes
 - Remove unintended print of `char[]` as String
diff --git a/test/unit/stdlib/collections/copy_map.c3 b/test/unit/stdlib/collections/copy_map.c3
index 3fac5f3fd..55dfe1622 100644
--- a/test/unit/stdlib/collections/copy_map.c3
+++ b/test/unit/stdlib/collections/copy_map.c3
@@ -28,4 +28,23 @@ fn void! copy_map() @test
 		y.free();
 		assert(alloc.allocated() == 0);
 	};
-}
\ No newline at end of file
+}
+
+/*
+	Some Keys (including Strings) are deep_copied into the hashmap on insertion.
+	When copying keys out, the keys themselves must also be deep-copied out.
+	Otherwise when the map is freed, the copied-in keys will also be freed,
+	resulting in use-after-free.
+*/
+fn void! copy_keys() @test
+{
+	String[] y;
+	@pool() {
+		IntMap x;
+		x.temp_init();
+		x.set("hello", 0); // keys copied into temp hashmap
+		y = x.copy_keys(allocator::heap()); // keys copied out
+		// end of pool: hashmap and its copied-in keys dropped
+	};
+	assert(y == {"hello"});
+}