Skip to content

Commit

Permalink
Attempting to use heap instead of tls for large arrays
Browse files Browse the repository at this point in the history 
Workaround issue where older glibc cannot allocate large TLS
blocks after the program has started running ("cannot allocate memory
in static TLS block").

ref:
microsoft/LightGBM#6509
https://bugzilla.redhat.com/show_bug.cgi?id=1722181
  • Loading branch information
@ccotter
ccotter committed Aug 3, 2024
1 parent 37123e9 commit 4b4fd9e
Show file tree
Hide file tree
Showing 5 changed files with 94 additions and 35 deletions.
71 changes: 54 additions & 17 deletions compiler-rt/lib/msan/msan.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,8 @@
#include "ubsan/ubsan_flags.h"
#include "ubsan/ubsan_init.h"

#include <sys/mman.h> // For PROT_READ etc

// ACHTUNG! No system header includes in this file.

using namespace __sanitizer;
Expand All @@ -39,28 +41,36 @@ using namespace __sanitizer;
static THREADLOCAL int msan_expect_umr = 0;
static THREADLOCAL int msan_expected_umr_found = 0;

#define __msan_param_tls__sizeof (kMsanParamTlsSize / sizeof(u64))
#define __msan_param_origin_tls__sizeof (kMsanParamTlsSize / sizeof(u32))
#define __msan_retval_tls__sizeof (kMsanRetvalTlsSize / sizeof(u64))
#define __msan_va_arg_tls__sizeof (kMsanParamTlsSize / sizeof(u64))
#define __msan_va_arg_origin_tls__sizeof (kMsanParamTlsSize / sizeof(u32))

// Function argument shadow. Each argument starts at the next available 8-byte
// aligned address.
SANITIZER_INTERFACE_ATTRIBUTE
THREADLOCAL u64 __msan_param_tls[kMsanParamTlsSize / sizeof(u64)];
THREADLOCAL u64* __msan_param_tls;

// Function argument origin. Each argument starts at the same offset as the
// corresponding shadow in (__msan_param_tls). Slightly weird, but changing this
// would break compatibility with older prebuilt binaries.
SANITIZER_INTERFACE_ATTRIBUTE
THREADLOCAL u32 __msan_param_origin_tls[kMsanParamTlsSize / sizeof(u32)];
THREADLOCAL u32* __msan_param_origin_tls;



SANITIZER_INTERFACE_ATTRIBUTE
THREADLOCAL u64 __msan_retval_tls[kMsanRetvalTlsSize / sizeof(u64)];
THREADLOCAL u64* __msan_retval_tls;

SANITIZER_INTERFACE_ATTRIBUTE
THREADLOCAL u32 __msan_retval_origin_tls;

alignas(16) SANITIZER_INTERFACE_ATTRIBUTE THREADLOCAL u64
__msan_va_arg_tls[kMsanParamTlsSize / sizeof(u64)];
alignas(16) SANITIZER_INTERFACE_ATTRIBUTE THREADLOCAL u64*
__msan_va_arg_tls;

alignas(16) SANITIZER_INTERFACE_ATTRIBUTE THREADLOCAL u32
__msan_va_arg_origin_tls[kMsanParamTlsSize / sizeof(u32)];
alignas(16) SANITIZER_INTERFACE_ATTRIBUTE THREADLOCAL u32*
__msan_va_arg_origin_tls;

SANITIZER_INTERFACE_ATTRIBUTE
THREADLOCAL u64 __msan_va_arg_overflow_size_tls;
Expand Down Expand Up @@ -274,17 +284,17 @@ void ScopedThreadLocalStateBackup::Restore() {
// A lame implementation that only keeps essential state and resets the rest.
__msan_va_arg_overflow_size_tls = va_arg_overflow_size_tls;

internal_memset(__msan_param_tls, 0, sizeof(__msan_param_tls));
internal_memset(__msan_retval_tls, 0, sizeof(__msan_retval_tls));
internal_memset(__msan_va_arg_tls, 0, sizeof(__msan_va_arg_tls));
internal_memset(__msan_param_tls, 0, __msan_param_tls__sizeof);
internal_memset(__msan_retval_tls, 0, __msan_retval_tls__sizeof);
internal_memset(__msan_va_arg_tls, 0, __msan_va_arg_tls__sizeof);
internal_memset(__msan_va_arg_origin_tls, 0,
sizeof(__msan_va_arg_origin_tls));
__msan_va_arg_origin_tls__sizeof);

if (__msan_get_track_origins()) {
internal_memset(&__msan_retval_origin_tls, 0,
sizeof(__msan_retval_origin_tls));
internal_memset(__msan_param_origin_tls, 0,
sizeof(__msan_param_origin_tls));
__msan_param_origin_tls__sizeof);
}
}

Expand Down Expand Up @@ -441,12 +451,39 @@ static void CheckUnwind() {
stack.Print();
}

extern "C" void __msan_init_tls() {
if (__msan_param_tls)
return;

__msan_param_tls = (u64*)internal_mmap(nullptr, 0x1000, PROT_READ|PROT_WRITE, MAP_PRIVATE|MAP_ANONYMOUS, -1, 0);
__msan_param_origin_tls = (u32*)internal_mmap(nullptr, 0x1000, PROT_READ|PROT_WRITE, MAP_PRIVATE|MAP_ANONYMOUS, -1, 0);
__msan_retval_tls = (u64*)internal_mmap(nullptr, 0x1000, PROT_READ|PROT_WRITE, MAP_PRIVATE|MAP_ANONYMOUS, -1, 0);
__msan_va_arg_tls = (u64*)internal_mmap(nullptr, 0x1000, PROT_READ|PROT_WRITE, MAP_PRIVATE|MAP_ANONYMOUS, -1, 0);
__msan_va_arg_origin_tls = (u32*)internal_mmap(nullptr, 0x1000, PROT_READ|PROT_WRITE, MAP_PRIVATE|MAP_ANONYMOUS, -1, 0);
}

extern "C" void __msan_clear_tls() {
internal_munmap(__msan_param_tls, 0x1000);
internal_munmap(__msan_param_origin_tls, 0x1000);
internal_munmap(__msan_retval_tls, 0x1000);
internal_munmap(__msan_va_arg_tls, 0x1000);
internal_munmap(__msan_va_arg_origin_tls, 0x1000);

__msan_param_tls = 0;
__msan_param_origin_tls = 0;
__msan_retval_tls = 0;
__msan_va_arg_tls = 0;
__msan_va_arg_origin_tls = 0;
}

void __msan_init() {
CHECK(!msan_init_is_running);
if (msan_inited) return;
msan_init_is_running = 1;
SanitizerToolName = "MemorySanitizer";

__msan_init_tls();

AvoidCVE_2016_2143();

CacheBinaryName();
Expand Down Expand Up @@ -719,17 +756,17 @@ void __msan_finish_switch_fiber(const void **bottom_old, uptr *size_old) {
}
t->FinishSwitchFiber((uptr *)bottom_old, (uptr *)size_old);

internal_memset(__msan_param_tls, 0, sizeof(__msan_param_tls));
internal_memset(__msan_retval_tls, 0, sizeof(__msan_retval_tls));
internal_memset(__msan_va_arg_tls, 0, sizeof(__msan_va_arg_tls));
internal_memset(__msan_param_tls, 0, __msan_param_tls__sizeof);
internal_memset(__msan_retval_tls, 0, __msan_retval_tls__sizeof);
internal_memset(__msan_va_arg_tls, 0, __msan_va_arg_tls__sizeof);

if (__msan_get_track_origins()) {
internal_memset(__msan_param_origin_tls, 0,
sizeof(__msan_param_origin_tls));
__msan_param_origin_tls__sizeof);
internal_memset(&__msan_retval_origin_tls, 0,
sizeof(__msan_retval_origin_tls));
internal_memset(__msan_va_arg_origin_tls, 0,
sizeof(__msan_va_arg_origin_tls));
__msan_va_arg_origin_tls__sizeof);
}
}

Expand Down
3 changes: 3 additions & 0 deletions compiler-rt/lib/msan/msan.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -257,6 +257,9 @@ addr_is_type(uptr addr, int mapping_types) {
const int kMsanParamTlsSize = 800;
const int kMsanRetvalTlsSize = 800;

extern "C" void __msan_init_tls();
extern "C" void __msan_clear_tls();

namespace __msan {
extern int msan_inited;
extern bool msan_init_is_running;
Expand Down
9 changes: 9 additions & 0 deletions compiler-rt/lib/msan/msan_interceptors.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1154,6 +1154,15 @@ extern "C" int pthread_attr_init(void *attr);
extern "C" int pthread_attr_destroy(void *attr);

static void *MsanThreadStartFunc(void *arg) {
struct MsanTlsGuard {
MsanTlsGuard() {
__msan_init_tls();
}
~MsanTlsGuard() {
__msan_clear_tls();
}
} msanTlsGuard;

MsanThread *t = (MsanThread *)arg;
SetCurrentThread(t);
t->Init();
Expand Down
2 changes: 1 addition & 1 deletion compiler-rt/test/msan/vararg_shadow.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -82,7 +82,7 @@ __attribute__((noinline, no_sanitize("memory"))) void printtls(int line) {
#endif // DEBUG_VARARG_SHADOW_TEST

const int kMsanParamTlsSize = 800;
extern "C" __thread uint8_t __msan_va_arg_tls[];
extern "C" __thread uint8_t* __msan_va_arg_tls;

struct IntInt {
int a;
Expand Down
44 changes: 27 additions & 17 deletions llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -861,25 +861,25 @@ void MemorySanitizer::createUserspaceApi(Module &M, const TargetLibraryInfo &TLI
// Create the global TLS variables.
RetvalTLS =
getOrInsertGlobal(M, "__msan_retval_tls",
ArrayType::get(IRB.getInt64Ty(), kRetvalTLSSize / 8));
PointerType::get(IRB.getInt64Ty(), 0));

RetvalOriginTLS = getOrInsertGlobal(M, "__msan_retval_origin_tls", OriginTy);

ParamTLS =
getOrInsertGlobal(M, "__msan_param_tls",
ArrayType::get(IRB.getInt64Ty(), kParamTLSSize / 8));
PointerType::get(IRB.getInt64Ty(), 0));

ParamOriginTLS =
getOrInsertGlobal(M, "__msan_param_origin_tls",
ArrayType::get(OriginTy, kParamTLSSize / 4));
PointerType::get(OriginTy, 0));

VAArgTLS =
getOrInsertGlobal(M, "__msan_va_arg_tls",
ArrayType::get(IRB.getInt64Ty(), kParamTLSSize / 8));
PointerType::get(IRB.getInt64Ty(), 0));

VAArgOriginTLS =
getOrInsertGlobal(M, "__msan_va_arg_origin_tls",
ArrayType::get(OriginTy, kParamTLSSize / 4));
PointerType::get(OriginTy, 0));

VAArgOverflowSizeTLS =
getOrInsertGlobal(M, "__msan_va_arg_overflow_size_tls", IRB.getInt64Ty());
Expand Down Expand Up @@ -1850,7 +1850,7 @@ struct MemorySanitizerVisitor : public InstVisitor<MemorySanitizerVisitor> {
///
/// Shadow = ParamTLS+ArgOffset.
Value *getShadowPtrForArgument(IRBuilder<> &IRB, int ArgOffset) {
Value *Base = IRB.CreatePointerCast(MS.ParamTLS, MS.IntptrTy);
Value *Base = IRB.CreatePointerCast(IRB.CreateLoad(PointerType::get(MS.IntptrTy, 0), MS.ParamTLS), MS.IntptrTy);
if (ArgOffset)
Base = IRB.CreateAdd(Base, ConstantInt::get(MS.IntptrTy, ArgOffset));
return IRB.CreateIntToPtr(Base, IRB.getPtrTy(0), "_msarg");
Expand All @@ -1860,15 +1860,16 @@ struct MemorySanitizerVisitor : public InstVisitor<MemorySanitizerVisitor> {
Value *getOriginPtrForArgument(IRBuilder<> &IRB, int ArgOffset) {
if (!MS.TrackOrigins)
return nullptr;
Value *Base = IRB.CreatePointerCast(MS.ParamOriginTLS, MS.IntptrTy);
Value *Base = IRB.CreateLoad(PointerType::get(MS.IntptrTy, 0), MS.ParamOriginTLS);
//Value *Base = IRB.CreatePointerCast(MS.ParamOriginTLS, MS.IntptrTy);
if (ArgOffset)
Base = IRB.CreateAdd(Base, ConstantInt::get(MS.IntptrTy, ArgOffset));
return IRB.CreateIntToPtr(Base, IRB.getPtrTy(0), "_msarg_o");
}

/// Compute the shadow address for a retval.
Value *getShadowPtrForRetval(IRBuilder<> &IRB) {
return IRB.CreatePointerCast(MS.RetvalTLS, IRB.getPtrTy(0), "_msret");
return IRB.CreateLoad(PointerType::get(IRB.getInt64Ty(), 0), MS.RetvalTLS);
}

/// Compute the origin address for a retval.
Expand Down Expand Up @@ -4942,14 +4943,16 @@ struct VarArgHelperBase : public VarArgHelper {
: F(F), MS(MS), MSV(MSV), VAListTagSize(VAListTagSize) {}

Value *getShadowAddrForVAArgument(IRBuilder<> &IRB, unsigned ArgOffset) {
Value *Base = IRB.CreatePointerCast(MS.VAArgTLS, MS.IntptrTy);
Value *Base = IRB.CreateLoad(PointerType::get(MS.IntptrTy, 0), MS.VAArgTLS);
//Value *Base = IRB.CreatePointerCast(MS.VAArgTLS, MS.IntptrTy);
return IRB.CreateAdd(Base, ConstantInt::get(MS.IntptrTy, ArgOffset));
}

/// Compute the shadow address for a given va_arg.
Value *getShadowPtrForVAArgument(Type *Ty, IRBuilder<> &IRB,
unsigned ArgOffset) {
Value *Base = IRB.CreatePointerCast(MS.VAArgTLS, MS.IntptrTy);
Value *Base = IRB.CreateLoad(PointerType::get(MS.IntptrTy, 0), MS.VAArgTLS);
//Value *Base = IRB.CreatePointerCast(MS.VAArgTLS, MS.IntptrTy);
Base = IRB.CreateAdd(Base, ConstantInt::get(MS.IntptrTy, ArgOffset));
return IRB.CreateIntToPtr(Base, PointerType::get(MSV.getShadowTy(Ty), 0),
"_msarg_va_s");
Expand All @@ -4966,7 +4969,8 @@ struct VarArgHelperBase : public VarArgHelper {

/// Compute the origin address for a given va_arg.
Value *getOriginPtrForVAArgument(IRBuilder<> &IRB, int ArgOffset) {
Value *Base = IRB.CreatePointerCast(MS.VAArgOriginTLS, MS.IntptrTy);
Value *Base = IRB.CreateLoad(PointerType::get(MS.IntptrTy, 0), MS.VAArgOriginTLS);
//Value *Base = IRB.CreatePointerCast(MS.VAArgOriginTLS, MS.IntptrTy);
// getOriginPtrForVAArgument() is always called after
// getShadowPtrForVAArgument(), so __msan_va_arg_origin_tls can never
// overflow.
Expand Down Expand Up @@ -5185,13 +5189,15 @@ struct VarArgAMD64Helper : public VarArgHelperBase {
Value *SrcSize = IRB.CreateBinaryIntrinsic(
Intrinsic::umin, CopySize,
ConstantInt::get(MS.IntptrTy, kParamTLSSize));
IRB.CreateMemCpy(VAArgTLSCopy, kShadowTLSAlignment, MS.VAArgTLS,
IRB.CreateMemCpy(VAArgTLSCopy, kShadowTLSAlignment,
IRB.CreateLoad(PointerType::get(MS.IntptrTy, 0), MS.VAArgTLS),
kShadowTLSAlignment, SrcSize);
if (MS.TrackOrigins) {
VAArgTLSOriginCopy = IRB.CreateAlloca(Type::getInt8Ty(*MS.C), CopySize);
VAArgTLSOriginCopy->setAlignment(kShadowTLSAlignment);
IRB.CreateMemCpy(VAArgTLSOriginCopy, kShadowTLSAlignment,
MS.VAArgOriginTLS, kShadowTLSAlignment, SrcSize);
IRB.CreateLoad(PointerType::get(MS.IntptrTy, 0), MS.VAArgOriginTLS),
kShadowTLSAlignment, SrcSize);
}
}

Expand Down Expand Up @@ -5300,7 +5306,8 @@ struct VarArgMIPS64Helper : public VarArgHelperBase {
Value *SrcSize = IRB.CreateBinaryIntrinsic(
Intrinsic::umin, CopySize,
ConstantInt::get(MS.IntptrTy, kParamTLSSize));
IRB.CreateMemCpy(VAArgTLSCopy, kShadowTLSAlignment, MS.VAArgTLS,
IRB.CreateMemCpy(VAArgTLSCopy, kShadowTLSAlignment,
IRB.CreateLoad(PointerType::get(MS.IntptrTy, 0), MS.VAArgTLS),
kShadowTLSAlignment, SrcSize);
}

Expand Down Expand Up @@ -5471,7 +5478,8 @@ struct VarArgAArch64Helper : public VarArgHelperBase {
Value *SrcSize = IRB.CreateBinaryIntrinsic(
Intrinsic::umin, CopySize,
ConstantInt::get(MS.IntptrTy, kParamTLSSize));
IRB.CreateMemCpy(VAArgTLSCopy, kShadowTLSAlignment, MS.VAArgTLS,
IRB.CreateMemCpy(VAArgTLSCopy, kShadowTLSAlignment,
IRB.CreateLoad(PointerType::get(MS.IntptrTy, 0), MS.VAArgTLS),
kShadowTLSAlignment, SrcSize);
}

Expand Down Expand Up @@ -5684,7 +5692,8 @@ struct VarArgPowerPC64Helper : public VarArgHelperBase {
Value *SrcSize = IRB.CreateBinaryIntrinsic(
Intrinsic::umin, CopySize,
ConstantInt::get(MS.IntptrTy, kParamTLSSize));
IRB.CreateMemCpy(VAArgTLSCopy, kShadowTLSAlignment, MS.VAArgTLS,
IRB.CreateMemCpy(VAArgTLSCopy, kShadowTLSAlignment,
IRB.CreateLoad(PointerType::get(MS.IntptrTy, 0), MS.VAArgTLS),
kShadowTLSAlignment, SrcSize);
}

Expand Down Expand Up @@ -5974,7 +5983,8 @@ struct VarArgSystemZHelper : public VarArgHelperBase {
Value *SrcSize = IRB.CreateBinaryIntrinsic(
Intrinsic::umin, CopySize,
ConstantInt::get(MS.IntptrTy, kParamTLSSize));
IRB.CreateMemCpy(VAArgTLSCopy, kShadowTLSAlignment, MS.VAArgTLS,
IRB.CreateMemCpy(VAArgTLSCopy, kShadowTLSAlignment,
IRB.CreateLoad(PointerType::get(MS.IntptrTy, 0), MS.VAArgTLS),
kShadowTLSAlignment, SrcSize);
if (MS.TrackOrigins) {
VAArgTLSOriginCopy = IRB.CreateAlloca(Type::getInt8Ty(*MS.C), CopySize);
Expand Down

0 comments on commit 4b4fd9e

Please sign in to comment.