From 5146bbe296c8af345c1ca22bb0825f642322ad13 Mon Sep 17 00:00:00 2001
From: Luke Wilde <luke@ladybird.org>
Date: Thu, 27 Feb 2025 14:50:22 +0000
Subject: [PATCH] LibGC: Visit the edges of the cells that must survive garbage
 collection

Previously, we would only keep the cell that must survive alive, but
none of it's edges.

This cropped up with a GC UAF in must_survive_garbage_collection of
WebSocket in .NET's SignalR frontend implementation, where an
out-of-scope WebSocket had it's underlying EventTarget properties
garbage collected, and must_survive_garbage_collection read from the
destroyed EventTarget properties.

See: https://github.com/dotnet/aspnetcore/blob/main/src/SignalR/clients/ts/signalr/src/WebSocketTransport.ts#L81
Found on https://www.formula1.com/ during a live session.

Co-Authored-By: Tim Flynn <trflynn89@pm.me>
---
 Libraries/LibGC/Heap.cpp                      | 12 +++++++++--
 .../Text/expected/WebSocket/WebSocket-gc.txt  |  1 +
 .../Text/input/WebSocket/WebSocket-gc.html    | 21 +++++++++++++++++++
 3 files changed, 32 insertions(+), 2 deletions(-)
 create mode 100644 Tests/LibWeb/Text/expected/WebSocket/WebSocket-gc.txt
 create mode 100644 Tests/LibWeb/Text/input/WebSocket/WebSocket-gc.html

diff --git a/Libraries/LibGC/Heap.cpp b/Libraries/LibGC/Heap.cpp
index f470c42f4c7..9617211726d 100644
--- a/Libraries/LibGC/Heap.cpp
+++ b/Libraries/LibGC/Heap.cpp
@@ -430,6 +430,14 @@ void Heap::mark_live_cells(HashMap<Cell*, HeapRoot> const& roots)
     for (auto& inverse_root : m_uprooted_cells)
         inverse_root->set_marked(false);
 
+    for_each_block([&](auto& block) {
+        block.template for_each_cell_in_state<Cell::State::Live>([&](Cell* cell) {
+            if (!cell->is_marked() && cell_must_survive_garbage_collection(*cell))
+                cell->visit_edges(visitor);
+        });
+        return IterationDecision::Continue;
+    });
+
     m_uprooted_cells.clear();
 }
 
@@ -444,7 +452,7 @@ void Heap::finalize_unmarked_cells()
 {
     for_each_block([&](auto& block) {
         block.template for_each_cell_in_state<Cell::State::Live>([](Cell* cell) {
-            if (!cell->is_marked() && !cell_must_survive_garbage_collection(*cell))
+            if (!cell->is_marked())
                 cell->finalize();
         });
         return IterationDecision::Continue;
@@ -466,7 +474,7 @@ void Heap::sweep_dead_cells(bool print_report, Core::ElapsedTimer const& measure
         bool block_has_live_cells = false;
         bool block_was_full = block.is_full();
         block.template for_each_cell_in_state<Cell::State::Live>([&](Cell* cell) {
-            if (!cell->is_marked() && !cell_must_survive_garbage_collection(*cell)) {
+            if (!cell->is_marked()) {
                 dbgln_if(HEAP_DEBUG, "  ~ {}", cell);
                 block.deallocate(cell);
                 ++collected_cells;
diff --git a/Tests/LibWeb/Text/expected/WebSocket/WebSocket-gc.txt b/Tests/LibWeb/Text/expected/WebSocket/WebSocket-gc.txt
new file mode 100644
index 00000000000..50586a4cbfb
--- /dev/null
+++ b/Tests/LibWeb/Text/expected/WebSocket/WebSocket-gc.txt
@@ -0,0 +1 @@
+PASS! (Didn't crash)
diff --git a/Tests/LibWeb/Text/input/WebSocket/WebSocket-gc.html b/Tests/LibWeb/Text/input/WebSocket/WebSocket-gc.html
new file mode 100644
index 00000000000..246911e4fbb
--- /dev/null
+++ b/Tests/LibWeb/Text/input/WebSocket/WebSocket-gc.html
@@ -0,0 +1,21 @@
+<script src="../include.js"></script>
+<script>
+    asyncTest((done) => {
+        {
+            const ws = new WebSocket('wss://echo.websocket.org');
+
+            ws.onopen = () => {
+            };
+            ws.onmessage = () => {
+            };
+            ws.onerror = () => {
+            };
+        }
+
+        setTimeout(() => {
+            internals.gc();
+            println("PASS! (Didn't crash)");
+            done();
+        }, 0);
+    });
+</script>