[llvm] [SampleProfile] Fix UB in Demangler invocation. (PR #137659)

[Krzysztof Pszeniczny via llvm-commits]

https://github.com/amharc created https://github.com/llvm/llvm-project/pull/137659

Currently the backing buffer of a `std::vector<char>` is passed[1] to `Demangler.getFunctionBaseName`. However, deeply inside the call stack `OutputBuffer::grow` will call[2] `std::realloc` if it needs to grow the buffer, leading to UB.

The demangler APIs specify[3] that "`Buf` and `N` behave like the second and third parameters to `__cxa_demangle`" and the docs for the latter say[4] that the output buffer must be allocated with `malloc` (but can also be `NULL` and will then be realloced accordingly). 

Note: PR #135863 changed this from a stack array to a `std::vector` and increased the size to 65K, but this can still lead to a crash if the demangled name is longer than that - yes, I'm surprised that a >65K-long function name happens in practice...

[1]: https://github.com/llvm/llvm-project/blob/d7e631c7cd6d9c13b9519991ec6becf08bc6b8aa/llvm/lib/Transforms/IPO/SampleProfileMatcher.cpp#L744
[2]: https://github.com/llvm/llvm-project/blob/d7e631c7cd6d9c13b9519991ec6becf08bc6b8aa/llvm/include/llvm/Demangle/Utility.h#L50
[3]: https://github.com/llvm/llvm-project/blob/d7e631c7cd6d9c13b9519991ec6becf08bc6b8aa/llvm/include/llvm/Demangle/Demangle.h#L92-L93
[4]: https://gcc.gnu.org/onlinedocs/libstdc++/libstdc++-html-USERS-4.3/a01696.html

>From a4af7bd1d1cfe3066f4a2b9701fd2b8cfaf1579d Mon Sep 17 00:00:00 2001
From: Krzysztof Pszeniczny <kpszeniczny at google.com>
Date: Mon, 28 Apr 2025 17:38:46 +0200
Subject: [PATCH] [SampleProfile] Fix UB in Demangler invocation.

Currently the backing buffer of a std::vector<char> is passed[1] to
Demangler.getFunctionBaseName. However, deeply inside the call stack
OutputBuffer::grow will call[2] std::realloc if it needs to grow the
buffer, leading to UB.

The demangler APIs specify[3] that "Buf and N behave like the second
and third parameters to __cxa_demangle" and the docs for the latter
say[4] that the output buffer must be allocated with malloc (but can
also be NULL and will then be realloced accordingly).

[1]: https://github.com/llvm/llvm-project/blob/d7e631c7cd6d9c13b9519991ec6becf08bc6b8aa/llvm/lib/Transforms/IPO/SampleProfileMatcher.cpp#L744
[2]: https://github.com/llvm/llvm-project/blob/d7e631c7cd6d9c13b9519991ec6becf08bc6b8aa/llvm/include/llvm/Demangle/Utility.h#L50
[3]: https://github.com/llvm/llvm-project/blob/d7e631c7cd6d9c13b9519991ec6becf08bc6b8aa/llvm/include/llvm/Demangle/Demangle.h#L92-L93
[4]: https://gcc.gnu.org/onlinedocs/libstdc++/libstdc++-html-USERS-4.3/a01696.html
---
 llvm/lib/Transforms/IPO/SampleProfileMatcher.cpp | 15 +++++++--------
 1 file changed, 7 insertions(+), 8 deletions(-)

diff --git a/llvm/lib/Transforms/IPO/SampleProfileMatcher.cpp b/llvm/lib/Transforms/IPO/SampleProfileMatcher.cpp
index 963c321772d6e..fdb5631f415a0 100644
--- a/llvm/lib/Transforms/IPO/SampleProfileMatcher.cpp
+++ b/llvm/lib/Transforms/IPO/SampleProfileMatcher.cpp
@@ -737,14 +737,13 @@ bool SampleProfileMatcher::functionMatchesProfileHelper(
     auto FunctionName = FName.str();
     if (Demangler.partialDemangle(FunctionName.c_str()))
       return std::string();
-    constexpr size_t MaxBaseNameSize = 65536;
-    std::vector<char> BaseNameBuf(MaxBaseNameSize, 0);
-    size_t BaseNameSize = MaxBaseNameSize;
-    char *BaseNamePtr =
-        Demangler.getFunctionBaseName(BaseNameBuf.data(), &BaseNameSize);
-    return (BaseNamePtr && BaseNameSize)
-               ? std::string(BaseNamePtr, BaseNameSize)
-               : std::string();
+    size_t BaseNameSize = 0;
+    char *BaseNamePtr = Demangler.getFunctionBaseName(nullptr, &BaseNameSize);
+    std::string Result = (BaseNamePtr && BaseNameSize)
+                             ? std::string(BaseNamePtr, BaseNameSize)
+                             : std::string();
+    free(BaseNamePtr);
+    return Result;
   };
   auto IRBaseName = GetBaseName(IRFunc.getName());
   auto ProfBaseName = GetBaseName(ProfFunc.stringRef());