[llvm] [CSSPGO] Fix use-after-free in SampleContextTracker::getBaseSamplesFor (PR #210824)
via llvm-commits
llvm-commits at lists.llvm.org
Mon Jul 20 20:40:00 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-llvm-transforms
Author: Kunal Pathak (kunalspathak)
<details>
<summary>Changes</summary>
`getBaseSamplesFor()` iterates the pre-collected list `FuncToCtxtProfiles[Name]` and promotes each context profile's node into the base profile via `promoteMergeContextSamplesTree()`. Promotion mutates the trie:
`moveContextSamples()` relocates a node by copying it (and its whole subtree) into the destination and repointing each moved `FunctionSamples` in `ProfileToNodeMap` to the copy, and the merge path frees nodes via `ContextTrieNode::getAllChildContext().clear()` / `removeChildContext()`.
The bug is a stale entry in `ProfileToNodeMap`: when a node is destroyed during promotion, its `FunctionSamples*` is left mapping to the freed node. When a function appears in nested (e.g. recursive) contexts, promoting one context can free the trie node of another same-function context that is still a pending entry in the worklist. A later iteration calls `getContextNodeForProfile()` for that entry, gets back the dangling pointer, and `promoteMergeContextSamplesTree()` dereferences it -- a heap-use-after-free. It only manifests once enough
promotions relocate/free nodes, so it reproduces on large context-sensitive profiles (observed with AddressSanitizer on a full application CS profile) but not on small ones.
~Fix: Skip `SyntheticContext` profiles in `getBaseSamplesFor()`. Besides fixing the UAF, this is correct: a synthesized node has already been merged into the base and must not be re-merged (double counting).~
Keep `ProfileToNodeMap` consistent instead of trusting profile state. When a subtree is about to be destroyed (`clear()` of merged children, and `removeChildContext()` of a promoted node), `invalidateContextNodesInMap()`
drops the map entries for those nodes. It is guarded by an identity check (`entry->second == Node`) so a profile whose mapping was already repointed to a live relocated copy is preserved. `getBaseSamplesFor()` then skips any entry whose node has been freed (`getContextNodeForProfile()` returns null), which cannot happen for a live node still needing promotion.
An earlier version of this fix instead skipped all `SyntheticContext` profiles in `getBaseSamplesFor()`. That is incorrect: `moveContextSamples()` marks the *entire* relocated subtree `SyntheticContext`, including distinct callee contexts (e.g. `bar` nested under a promoted `foo`) that are still live and must get their own base profile. Skipping them drops those functions' samples and regressed `llvm-profgen/X86/cs-preinline.test` and `llvm-profgen/X86/profile-density.test`. The map-invalidation approach fixes the UAF without changing the promotion
behavior for live nodes.
Added a SampleContextTracker unit test covering nested/recursive contexts that triggers the relocate-then-re-promote path; it fails under ASAN on the unfixed code and passes with the fix.
---
Full diff: https://github.com/llvm/llvm-project/pull/210824.diff
4 Files Affected:
- (modified) llvm/include/llvm/Transforms/IPO/SampleContextTracker.h (+5)
- (modified) llvm/lib/Transforms/IPO/SampleContextTracker.cpp (+33-3)
- (modified) llvm/unittests/Transforms/IPO/CMakeLists.txt (+2)
- (added) llvm/unittests/Transforms/IPO/SampleContextTrackerTest.cpp (+106)
``````````diff
diff --git a/llvm/include/llvm/Transforms/IPO/SampleContextTracker.h b/llvm/include/llvm/Transforms/IPO/SampleContextTracker.h
index f8890899fedfb..d0f188be71a87 100644
--- a/llvm/include/llvm/Transforms/IPO/SampleContextTracker.h
+++ b/llvm/include/llvm/Transforms/IPO/SampleContextTracker.h
@@ -210,6 +210,11 @@ class SampleContextTracker {
void setContextNode(const FunctionSamples *FSample, ContextTrieNode *Node) {
ProfileToNodeMap[FSample] = Node;
}
+ // Drop ProfileToNodeMap entries for Node and its whole subtree that still
+ // point at these nodes, so getContextNodeForProfile() never returns a
+ // dangling pointer after the nodes are destroyed. Entries already repointed
+ // to a relocated copy (moveContextSamples) are left untouched.
+ void invalidateContextNodesInMap(ContextTrieNode &Node);
// Map from function name to context profiles (excluding base profile)
HashKeyMap<std::unordered_map, FunctionId, ContextSamplesTy>
FuncToCtxtProfiles;
diff --git a/llvm/lib/Transforms/IPO/SampleContextTracker.cpp b/llvm/lib/Transforms/IPO/SampleContextTracker.cpp
index c8f203e8ce3ee..b9d091a3459aa 100644
--- a/llvm/lib/Transforms/IPO/SampleContextTracker.cpp
+++ b/llvm/lib/Transforms/IPO/SampleContextTracker.cpp
@@ -102,6 +102,25 @@ SampleContextTracker::moveContextSamples(ContextTrieNode &ToNodeParent,
return NewNode;
}
+void SampleContextTracker::invalidateContextNodesInMap(ContextTrieNode &Node) {
+ std::queue<ContextTrieNode *> NodeToInvalidate;
+ NodeToInvalidate.push(&Node);
+ while (!NodeToInvalidate.empty()) {
+ ContextTrieNode *N = NodeToInvalidate.front();
+ NodeToInvalidate.pop();
+ if (FunctionSamples *FSamples = N->getFunctionSamples()) {
+ auto I = ProfileToNodeMap.find(FSamples);
+ // Only drop the mapping if it still points at the node being destroyed.
+ // A relocated profile already had its mapping repointed to the live copy
+ // by moveContextSamples/mergeContextNode and must be preserved.
+ if (I != ProfileToNodeMap.end() && I->second == N)
+ ProfileToNodeMap.erase(I);
+ }
+ for (auto &It : N->getAllChildContext())
+ NodeToInvalidate.push(&It.second);
+ }
+}
+
void ContextTrieNode::removeChildContext(const LineLocation &CallSite,
FunctionId CalleeName) {
uint64_t Hash = FunctionSamples::getCallSiteHash(CalleeName, CallSite);
@@ -340,7 +359,13 @@ FunctionSamples *SampleContextTracker::getBaseSamplesFor(FunctionId Name,
continue;
ContextTrieNode *FromNode = getContextNodeForProfile(CSamples);
- if (FromNode == Node)
+ // A previous promotion in this loop may have merged and freed this
+ // context's node (e.g. a self-recursive context nested under one already
+ // promoted). Its map entry is cleared when the node is destroyed, so a
+ // null result here means the node is gone; skip it to avoid a
+ // use-after-free. Distinct callee contexts merely relocated under another
+ // function's base stay live and are still promoted.
+ if (!FromNode || FromNode == Node)
continue;
ContextTrieNode &ToNode = promoteMergeContextSamplesTree(*FromNode);
@@ -608,13 +633,18 @@ ContextTrieNode &SampleContextTracker::promoteMergeContextSamplesTree(
promoteMergeContextSamplesTree(FromChildNode, *ToNode);
}
- // Remove children once they're all merged
+ // Remove children once they're all merged. Drop their map entries first so
+ // no FunctionSamples keeps pointing at a node that clear() destroys.
+ for (auto &It : FromNode.getAllChildContext())
+ invalidateContextNodesInMap(It.second);
FromNode.getAllChildContext().clear();
}
// For root of subtree, remove itself from old parent too
- if (MoveToRoot)
+ if (MoveToRoot) {
+ invalidateContextNodesInMap(FromNode);
FromNodeParent.removeChildContext(OldCallSiteLoc, ToNode->getFuncName());
+ }
return *ToNode;
}
diff --git a/llvm/unittests/Transforms/IPO/CMakeLists.txt b/llvm/unittests/Transforms/IPO/CMakeLists.txt
index 5b45191afc711..b7550c0f16c89 100644
--- a/llvm/unittests/Transforms/IPO/CMakeLists.txt
+++ b/llvm/unittests/Transforms/IPO/CMakeLists.txt
@@ -3,6 +3,7 @@ set(LLVM_LINK_COMPONENTS
AsmParser
Core
IPO
+ ProfileData
Support
TargetParser
TransformUtils
@@ -15,4 +16,5 @@ add_llvm_unittest(IPOTests
FunctionSpecializationTest.cpp
ImportIDTableTests.cpp
MergeFunctionsTest.cpp
+ SampleContextTrackerTest.cpp
)
diff --git a/llvm/unittests/Transforms/IPO/SampleContextTrackerTest.cpp b/llvm/unittests/Transforms/IPO/SampleContextTrackerTest.cpp
new file mode 100644
index 0000000000000..b987a395ced85
--- /dev/null
+++ b/llvm/unittests/Transforms/IPO/SampleContextTrackerTest.cpp
@@ -0,0 +1,106 @@
+//===- SampleContextTrackerTest.cpp ---------------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/Transforms/IPO/SampleContextTracker.h"
+#include "llvm/ProfileData/SampleProf.h"
+#include "gtest/gtest.h"
+#include <list>
+
+using namespace llvm;
+using namespace sampleprof;
+
+namespace {
+
+// SampleContextFrames are non-owning ArrayRefs into a name table, so the table
+// backing every parsed context string must outlive the profile map and tracker.
+static FunctionSamples &
+addProfile(SampleProfileMap &Profiles,
+ std::list<SampleContextFrameVector> &CSNameTable, StringRef CtxStr,
+ uint64_t Samples) {
+ FunctionSamples &FS = Profiles.create(SampleContext(CtxStr, CSNameTable));
+ FS.addTotalSamples(Samples);
+ FS.addBodySamples(1, 0, Samples);
+ return FS;
+}
+
+// Regression test for a heap-use-after-free in
+// SampleContextTracker::getBaseSamplesFor().
+//
+// When a function appears in nested (e.g. recursive) contexts, promoting the
+// outer context relocates the inner same-function node into the base subtree
+// and later merges/frees it via clear()/removeChildContext().
+// getBaseSamplesFor() iterates a pre-collected list of all context profiles for
+// the function; the profile of a freed node otherwise still maps to that node
+// in ProfileToNodeMap, so re-promoting it reads freed memory. The fix drops
+// such stale map entries when a node is destroyed, so
+// getContextNodeForProfile() returns null and the entry is skipped.
+//
+// This builds exactly that shape (foo nested under foo) and checks the base
+// profile is produced without crashing. Under an ASAN build this test fails on
+// the unfixed code with a heap-use-after-free.
+TEST(SampleContextTrackerTest, GetBaseSamplesForNestedRecursiveContext) {
+ std::list<SampleContextFrameVector> CSNameTable;
+ SampleProfileMap Profiles;
+
+ addProfile(Profiles, CSNameTable, "[main:1 @ foo]", 100);
+ addProfile(Profiles, CSNameTable, "[main:1 @ foo:2 @ foo]", 50);
+ addProfile(Profiles, CSNameTable, "[main:1 @ foo:2 @ foo:2 @ foo]", 25);
+
+ SampleContextTracker Tracker(Profiles, /*GUIDToFuncNameMap=*/nullptr);
+
+ // Must not crash (heap-use-after-free on the unfixed code) and must return a
+ // non-null merged base profile carrying samples.
+ FunctionSamples *Base =
+ Tracker.getBaseSamplesFor(FunctionId("foo"), /*MergeContext=*/true);
+ ASSERT_NE(Base, nullptr);
+ EXPECT_GE(Base->getTotalSamples(), 100u);
+}
+
+// A single (non-nested) context should also merge cleanly and is a sanity check
+// that fixing the use-after-free does not drop a legitimate first-time
+// promotion.
+TEST(SampleContextTrackerTest, GetBaseSamplesForSingleContext) {
+ std::list<SampleContextFrameVector> CSNameTable;
+ SampleProfileMap Profiles;
+
+ addProfile(Profiles, CSNameTable, "[main:1 @ bar]", 42);
+
+ SampleContextTracker Tracker(Profiles, /*GUIDToFuncNameMap=*/nullptr);
+ FunctionSamples *Base =
+ Tracker.getBaseSamplesFor(FunctionId("bar"), /*MergeContext=*/true);
+ ASSERT_NE(Base, nullptr);
+ EXPECT_EQ(Base->getTotalSamples(), 42u);
+}
+
+// A distinct callee nested under another function must still get its own base
+// profile. Promoting foo relocates bar's node into foo's base subtree (marking
+// bar SyntheticContext); the relocated node stays live, so getBaseSamplesFor
+// for bar must still promote it. A too-broad "skip all SyntheticContext"
+// use-after-free workaround would drop bar's samples here.
+TEST(SampleContextTrackerTest, GetBaseSamplesForDistinctNestedCallee) {
+ std::list<SampleContextFrameVector> CSNameTable;
+ SampleProfileMap Profiles;
+
+ addProfile(Profiles, CSNameTable, "[main:1 @ foo]", 100);
+ addProfile(Profiles, CSNameTable, "[main:1 @ foo:3 @ bar]", 42);
+
+ SampleContextTracker Tracker(Profiles, /*GUIDToFuncNameMap=*/nullptr);
+
+ // Promote foo first; this relocates bar's node under foo's base.
+ FunctionSamples *FooBase =
+ Tracker.getBaseSamplesFor(FunctionId("foo"), /*MergeContext=*/true);
+ ASSERT_NE(FooBase, nullptr);
+
+ // bar must still be promotable to its own base with its samples intact.
+ FunctionSamples *BarBase =
+ Tracker.getBaseSamplesFor(FunctionId("bar"), /*MergeContext=*/true);
+ ASSERT_NE(BarBase, nullptr);
+ EXPECT_EQ(BarBase->getTotalSamples(), 42u);
+}
+
+} // namespace
``````````
</details>
https://github.com/llvm/llvm-project/pull/210824
More information about the llvm-commits
mailing list