558 lines · cpp
1//===- AllocToken.cpp - Allocation token instrumentation ------------------===//2//3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.4// See https://llvm.org/LICENSE.txt for license information.5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception6//7//===----------------------------------------------------------------------===//8//9// This file implements AllocToken, an instrumentation pass that10// replaces allocation calls with token-enabled versions.11//12//===----------------------------------------------------------------------===//13 14#include "llvm/Transforms/Instrumentation/AllocToken.h"15#include "llvm/ADT/DenseMap.h"16#include "llvm/ADT/SmallPtrSet.h"17#include "llvm/ADT/SmallVector.h"18#include "llvm/ADT/Statistic.h"19#include "llvm/ADT/StringExtras.h"20#include "llvm/ADT/StringRef.h"21#include "llvm/Analysis/MemoryBuiltins.h"22#include "llvm/Analysis/OptimizationRemarkEmitter.h"23#include "llvm/Analysis/TargetLibraryInfo.h"24#include "llvm/IR/Analysis.h"25#include "llvm/IR/Attributes.h"26#include "llvm/IR/Constants.h"27#include "llvm/IR/DerivedTypes.h"28#include "llvm/IR/Function.h"29#include "llvm/IR/GlobalValue.h"30#include "llvm/IR/IRBuilder.h"31#include "llvm/IR/InstIterator.h"32#include "llvm/IR/InstrTypes.h"33#include "llvm/IR/Instructions.h"34#include "llvm/IR/IntrinsicInst.h"35#include "llvm/IR/Metadata.h"36#include "llvm/IR/Module.h"37#include "llvm/IR/PassManager.h"38#include "llvm/IR/Type.h"39#include "llvm/Support/AllocToken.h"40#include "llvm/Support/Casting.h"41#include "llvm/Support/CommandLine.h"42#include "llvm/Support/Compiler.h"43#include "llvm/Support/ErrorHandling.h"44#include "llvm/Support/RandomNumberGenerator.h"45#include "llvm/Support/SipHash.h"46#include "llvm/Support/raw_ostream.h"47#include <cassert>48#include <cstddef>49#include <cstdint>50#include <limits>51#include <memory>52#include <optional>53#include <string>54#include <utility>55#include <variant>56 57using namespace llvm;58using TokenMode = AllocTokenMode;59 60#define DEBUG_TYPE "alloc-token"61 62namespace {63 64//===--- Command-line options ---------------------------------------------===//65 66cl::opt<std::string> ClFuncPrefix("alloc-token-prefix",67 cl::desc("The allocation function prefix"),68 cl::Hidden, cl::init("__alloc_token_"));69 70cl::opt<uint64_t>71 ClMaxTokens("alloc-token-max",72 cl::desc("Maximum number of tokens (0 = target SIZE_MAX)"),73 cl::Hidden, cl::init(0));74 75cl::opt<bool>76 ClFastABI("alloc-token-fast-abi",77 cl::desc("The token ID is encoded in the function name"),78 cl::Hidden, cl::init(false));79 80// Instrument libcalls only by default - compatible allocators only need to take81// care of providing standard allocation functions. With extended coverage, also82// instrument non-libcall allocation function calls with !alloc_token83// metadata.84cl::opt<bool>85 ClExtended("alloc-token-extended",86 cl::desc("Extend coverage to custom allocation functions"),87 cl::Hidden, cl::init(false));88 89// C++ defines ::operator new (and variants) as replaceable (vs. standard90// library versions), which are nobuiltin, and are therefore not covered by91// isAllocationFn(). Cover by default, as users of AllocToken are already92// required to provide token-aware allocation functions (no defaults).93cl::opt<bool> ClCoverReplaceableNew("alloc-token-cover-replaceable-new",94 cl::desc("Cover replaceable operator new"),95 cl::Hidden, cl::init(true));96 97cl::opt<uint64_t> ClFallbackToken(98 "alloc-token-fallback",99 cl::desc("The default fallback token where none could be determined"),100 cl::Hidden, cl::init(0));101 102//===--- Statistics -------------------------------------------------------===//103 104STATISTIC(NumFunctionsModified, "Functions modified");105STATISTIC(NumAllocationsInstrumented, "Allocations instrumented");106 107//===----------------------------------------------------------------------===//108 109/// Returns the !alloc_token metadata if available.110///111/// Expected format is: !{<type-name>, <contains-pointer>}112MDNode *getAllocTokenMetadata(const CallBase &CB) {113 MDNode *Ret = nullptr;114 if (auto *II = dyn_cast<IntrinsicInst>(&CB);115 II && II->getIntrinsicID() == Intrinsic::alloc_token_id) {116 auto *MDV = cast<MetadataAsValue>(II->getArgOperand(0));117 Ret = cast<MDNode>(MDV->getMetadata());118 // If the intrinsic has an empty MDNode, type inference failed.119 if (Ret->getNumOperands() == 0)120 return nullptr;121 } else {122 Ret = CB.getMetadata(LLVMContext::MD_alloc_token);123 if (!Ret)124 return nullptr;125 }126 assert(Ret->getNumOperands() == 2 && "bad !alloc_token");127 assert(isa<MDString>(Ret->getOperand(0)));128 assert(isa<ConstantAsMetadata>(Ret->getOperand(1)));129 return Ret;130}131 132bool containsPointer(const MDNode *MD) {133 ConstantAsMetadata *C = cast<ConstantAsMetadata>(MD->getOperand(1));134 auto *CI = cast<ConstantInt>(C->getValue());135 return CI->getValue().getBoolValue();136}137 138class ModeBase {139public:140 explicit ModeBase(const IntegerType &TokenTy, uint64_t MaxTokens)141 : MaxTokens(MaxTokens ? MaxTokens : TokenTy.getBitMask()) {142 assert(MaxTokens <= TokenTy.getBitMask());143 }144 145protected:146 uint64_t boundedToken(uint64_t Val) const {147 assert(MaxTokens != 0);148 return Val % MaxTokens;149 }150 151 const uint64_t MaxTokens;152};153 154/// Implementation for TokenMode::Increment.155class IncrementMode : public ModeBase {156public:157 using ModeBase::ModeBase;158 159 uint64_t operator()(const CallBase &CB, OptimizationRemarkEmitter &) {160 return boundedToken(Counter++);161 }162 163private:164 uint64_t Counter = 0;165};166 167/// Implementation for TokenMode::Random.168class RandomMode : public ModeBase {169public:170 RandomMode(const IntegerType &TokenTy, uint64_t MaxTokens,171 std::unique_ptr<RandomNumberGenerator> RNG)172 : ModeBase(TokenTy, MaxTokens), RNG(std::move(RNG)) {}173 uint64_t operator()(const CallBase &CB, OptimizationRemarkEmitter &) {174 return boundedToken((*RNG)());175 }176 177private:178 std::unique_ptr<RandomNumberGenerator> RNG;179};180 181/// Implementation for TokenMode::TypeHash. The implementation ensures182/// hashes are stable across different compiler invocations. Uses SipHash as the183/// hash function.184class TypeHashMode : public ModeBase {185public:186 using ModeBase::ModeBase;187 188 uint64_t operator()(const CallBase &CB, OptimizationRemarkEmitter &ORE) {189 190 if (MDNode *N = getAllocTokenMetadata(CB)) {191 MDString *S = cast<MDString>(N->getOperand(0));192 AllocTokenMetadata Metadata{S->getString(), containsPointer(N)};193 if (auto Token = getAllocToken(TokenMode::TypeHash, Metadata, MaxTokens))194 return *Token;195 }196 // Fallback.197 remarkNoMetadata(CB, ORE);198 return ClFallbackToken;199 }200 201protected:202 /// Remark that there was no precise type information.203 static void remarkNoMetadata(const CallBase &CB,204 OptimizationRemarkEmitter &ORE) {205 ORE.emit([&] {206 ore::NV FuncNV("Function", CB.getParent()->getParent());207 const Function *Callee = CB.getCalledFunction();208 ore::NV CalleeNV("Callee", Callee ? Callee->getName() : "<unknown>");209 return OptimizationRemark(DEBUG_TYPE, "NoAllocToken", &CB)210 << "Call to '" << CalleeNV << "' in '" << FuncNV211 << "' without source-level type token";212 });213 }214};215 216/// Implementation for TokenMode::TypeHashPointerSplit.217class TypeHashPointerSplitMode : public TypeHashMode {218public:219 using TypeHashMode::TypeHashMode;220 221 uint64_t operator()(const CallBase &CB, OptimizationRemarkEmitter &ORE) {222 if (MDNode *N = getAllocTokenMetadata(CB)) {223 MDString *S = cast<MDString>(N->getOperand(0));224 AllocTokenMetadata Metadata{S->getString(), containsPointer(N)};225 if (auto Token = getAllocToken(TokenMode::TypeHashPointerSplit, Metadata,226 MaxTokens))227 return *Token;228 }229 // Pick the fallback token (ClFallbackToken), which by default is 0, meaning230 // it'll fall into the pointer-less bucket. Override by setting231 // -alloc-token-fallback if that is the wrong choice.232 remarkNoMetadata(CB, ORE);233 return ClFallbackToken;234 }235};236 237// Apply opt overrides.238AllocTokenOptions transformOptionsFromCl(AllocTokenOptions Opts) {239 if (!Opts.MaxTokens.has_value())240 Opts.MaxTokens = ClMaxTokens;241 Opts.FastABI |= ClFastABI;242 Opts.Extended |= ClExtended;243 return Opts;244}245 246class AllocToken {247public:248 explicit AllocToken(AllocTokenOptions Opts, Module &M,249 ModuleAnalysisManager &MAM)250 : Options(transformOptionsFromCl(std::move(Opts))), Mod(M),251 FAM(MAM.getResult<FunctionAnalysisManagerModuleProxy>(M).getManager()),252 Mode(IncrementMode(*IntPtrTy, *Options.MaxTokens)) {253 switch (Options.Mode) {254 case TokenMode::Increment:255 break;256 case TokenMode::Random:257 Mode.emplace<RandomMode>(*IntPtrTy, *Options.MaxTokens,258 M.createRNG(DEBUG_TYPE));259 break;260 case TokenMode::TypeHash:261 Mode.emplace<TypeHashMode>(*IntPtrTy, *Options.MaxTokens);262 break;263 case TokenMode::TypeHashPointerSplit:264 Mode.emplace<TypeHashPointerSplitMode>(*IntPtrTy, *Options.MaxTokens);265 break;266 }267 }268 269 bool instrumentFunction(Function &F);270 271private:272 /// Returns the LibFunc (or NotLibFunc) if this call should be instrumented.273 std::optional<LibFunc>274 shouldInstrumentCall(const CallBase &CB, const TargetLibraryInfo &TLI) const;275 276 /// Returns true for functions that are eligible for instrumentation.277 static bool isInstrumentableLibFunc(LibFunc Func, const CallBase &CB,278 const TargetLibraryInfo &TLI);279 280 /// Returns true for isAllocationFn() functions that we should ignore.281 static bool ignoreInstrumentableLibFunc(LibFunc Func);282 283 /// Replace a call/invoke with a call/invoke to the allocation function284 /// with token ID.285 bool replaceAllocationCall(CallBase *CB, LibFunc Func,286 OptimizationRemarkEmitter &ORE,287 const TargetLibraryInfo &TLI);288 289 /// Return replacement function for a LibFunc that takes a token ID.290 FunctionCallee getTokenAllocFunction(const CallBase &CB, uint64_t TokenID,291 LibFunc OriginalFunc);292 293 /// Lower alloc_token_* intrinsics.294 void replaceIntrinsicInst(IntrinsicInst *II, OptimizationRemarkEmitter &ORE);295 296 /// Return the token ID from metadata in the call.297 uint64_t getToken(const CallBase &CB, OptimizationRemarkEmitter &ORE) {298 return std::visit([&](auto &&Mode) { return Mode(CB, ORE); }, Mode);299 }300 301 const AllocTokenOptions Options;302 Module &Mod;303 IntegerType *IntPtrTy = Mod.getDataLayout().getIntPtrType(Mod.getContext());304 FunctionAnalysisManager &FAM;305 // Cache for replacement functions.306 DenseMap<std::pair<LibFunc, uint64_t>, FunctionCallee> TokenAllocFunctions;307 // Selected mode.308 std::variant<IncrementMode, RandomMode, TypeHashMode,309 TypeHashPointerSplitMode>310 Mode;311};312 313bool AllocToken::instrumentFunction(Function &F) {314 // Do not apply any instrumentation for naked functions.315 if (F.hasFnAttribute(Attribute::Naked))316 return false;317 // Don't touch available_externally functions, their actual body is elsewhere.318 if (F.getLinkage() == GlobalValue::AvailableExternallyLinkage)319 return false;320 321 auto &ORE = FAM.getResult<OptimizationRemarkEmitterAnalysis>(F);322 auto &TLI = FAM.getResult<TargetLibraryAnalysis>(F);323 SmallVector<std::pair<CallBase *, LibFunc>, 4> AllocCalls;324 SmallVector<IntrinsicInst *, 4> IntrinsicInsts;325 326 // Only instrument functions that have the sanitize_alloc_token attribute.327 const bool InstrumentFunction =328 F.hasFnAttribute(Attribute::SanitizeAllocToken) &&329 !F.hasFnAttribute(Attribute::DisableSanitizerInstrumentation);330 331 // Collect all allocation calls to avoid iterator invalidation.332 for (Instruction &I : instructions(F)) {333 // Collect all alloc_token_* intrinsics.334 if (auto *II = dyn_cast<IntrinsicInst>(&I);335 II && II->getIntrinsicID() == Intrinsic::alloc_token_id) {336 IntrinsicInsts.emplace_back(II);337 continue;338 }339 340 if (!InstrumentFunction)341 continue;342 343 auto *CB = dyn_cast<CallBase>(&I);344 if (!CB)345 continue;346 if (std::optional<LibFunc> Func = shouldInstrumentCall(*CB, TLI))347 AllocCalls.emplace_back(CB, Func.value());348 }349 350 bool Modified = false;351 352 if (!AllocCalls.empty()) {353 for (auto &[CB, Func] : AllocCalls)354 Modified |= replaceAllocationCall(CB, Func, ORE, TLI);355 if (Modified)356 NumFunctionsModified++;357 }358 359 if (!IntrinsicInsts.empty()) {360 for (auto *II : IntrinsicInsts)361 replaceIntrinsicInst(II, ORE);362 Modified = true;363 NumFunctionsModified++;364 }365 366 return Modified;367}368 369std::optional<LibFunc>370AllocToken::shouldInstrumentCall(const CallBase &CB,371 const TargetLibraryInfo &TLI) const {372 const Function *Callee = CB.getCalledFunction();373 if (!Callee)374 return std::nullopt;375 376 // Ignore nobuiltin of the CallBase, so that we can cover nobuiltin libcalls377 // if requested via isInstrumentableLibFunc(). Note that isAllocationFn() is378 // returning false for nobuiltin calls.379 LibFunc Func;380 if (TLI.getLibFunc(*Callee, Func)) {381 if (isInstrumentableLibFunc(Func, CB, TLI))382 return Func;383 } else if (Options.Extended && CB.getMetadata(LLVMContext::MD_alloc_token)) {384 return NotLibFunc;385 }386 387 return std::nullopt;388}389 390bool AllocToken::isInstrumentableLibFunc(LibFunc Func, const CallBase &CB,391 const TargetLibraryInfo &TLI) {392 if (ignoreInstrumentableLibFunc(Func))393 return false;394 395 if (isAllocationFn(&CB, &TLI))396 return true;397 398 switch (Func) {399 // These libfuncs don't return normal pointers, and are therefore not handled400 // by isAllocationFn().401 case LibFunc_posix_memalign:402 case LibFunc_size_returning_new:403 case LibFunc_size_returning_new_hot_cold:404 case LibFunc_size_returning_new_aligned:405 case LibFunc_size_returning_new_aligned_hot_cold:406 return true;407 408 // See comment above ClCoverReplaceableNew.409 case LibFunc_Znwj:410 case LibFunc_ZnwjRKSt9nothrow_t:411 case LibFunc_ZnwjSt11align_val_t:412 case LibFunc_ZnwjSt11align_val_tRKSt9nothrow_t:413 case LibFunc_Znwm:414 case LibFunc_Znwm12__hot_cold_t:415 case LibFunc_ZnwmRKSt9nothrow_t:416 case LibFunc_ZnwmRKSt9nothrow_t12__hot_cold_t:417 case LibFunc_ZnwmSt11align_val_t:418 case LibFunc_ZnwmSt11align_val_t12__hot_cold_t:419 case LibFunc_ZnwmSt11align_val_tRKSt9nothrow_t:420 case LibFunc_ZnwmSt11align_val_tRKSt9nothrow_t12__hot_cold_t:421 case LibFunc_Znaj:422 case LibFunc_ZnajRKSt9nothrow_t:423 case LibFunc_ZnajSt11align_val_t:424 case LibFunc_ZnajSt11align_val_tRKSt9nothrow_t:425 case LibFunc_Znam:426 case LibFunc_Znam12__hot_cold_t:427 case LibFunc_ZnamRKSt9nothrow_t:428 case LibFunc_ZnamRKSt9nothrow_t12__hot_cold_t:429 case LibFunc_ZnamSt11align_val_t:430 case LibFunc_ZnamSt11align_val_t12__hot_cold_t:431 case LibFunc_ZnamSt11align_val_tRKSt9nothrow_t:432 case LibFunc_ZnamSt11align_val_tRKSt9nothrow_t12__hot_cold_t:433 return ClCoverReplaceableNew;434 435 default:436 return false;437 }438}439 440bool AllocToken::ignoreInstrumentableLibFunc(LibFunc Func) {441 switch (Func) {442 case LibFunc_strdup:443 case LibFunc_dunder_strdup:444 case LibFunc_strndup:445 case LibFunc_dunder_strndup:446 return true;447 default:448 return false;449 }450}451 452bool AllocToken::replaceAllocationCall(CallBase *CB, LibFunc Func,453 OptimizationRemarkEmitter &ORE,454 const TargetLibraryInfo &TLI) {455 uint64_t TokenID = getToken(*CB, ORE);456 457 FunctionCallee TokenAlloc = getTokenAllocFunction(*CB, TokenID, Func);458 if (!TokenAlloc)459 return false;460 NumAllocationsInstrumented++;461 462 if (Options.FastABI) {463 assert(TokenAlloc.getFunctionType()->getNumParams() == CB->arg_size());464 CB->setCalledFunction(TokenAlloc);465 return true;466 }467 468 IRBuilder<> IRB(CB);469 // Original args.470 SmallVector<Value *, 4> NewArgs{CB->args()};471 // Add token ID, truncated to IntPtrTy width.472 NewArgs.push_back(ConstantInt::get(IntPtrTy, TokenID));473 assert(TokenAlloc.getFunctionType()->getNumParams() == NewArgs.size());474 475 // Preserve invoke vs call semantics for exception handling.476 CallBase *NewCall;477 if (auto *II = dyn_cast<InvokeInst>(CB)) {478 NewCall = IRB.CreateInvoke(TokenAlloc, II->getNormalDest(),479 II->getUnwindDest(), NewArgs);480 } else {481 NewCall = IRB.CreateCall(TokenAlloc, NewArgs);482 cast<CallInst>(NewCall)->setTailCall(CB->isTailCall());483 }484 NewCall->setCallingConv(CB->getCallingConv());485 NewCall->copyMetadata(*CB);486 NewCall->setAttributes(CB->getAttributes());487 488 // Replace all uses and delete the old call.489 CB->replaceAllUsesWith(NewCall);490 CB->eraseFromParent();491 return true;492}493 494FunctionCallee AllocToken::getTokenAllocFunction(const CallBase &CB,495 uint64_t TokenID,496 LibFunc OriginalFunc) {497 std::optional<std::pair<LibFunc, uint64_t>> Key;498 if (OriginalFunc != NotLibFunc) {499 Key = std::make_pair(OriginalFunc, Options.FastABI ? TokenID : 0);500 auto It = TokenAllocFunctions.find(*Key);501 if (It != TokenAllocFunctions.end())502 return It->second;503 }504 505 const Function *Callee = CB.getCalledFunction();506 if (!Callee)507 return FunctionCallee();508 const FunctionType *OldFTy = Callee->getFunctionType();509 if (OldFTy->isVarArg())510 return FunctionCallee();511 // Copy params, and append token ID type.512 Type *RetTy = OldFTy->getReturnType();513 SmallVector<Type *, 4> NewParams{OldFTy->params()};514 std::string TokenAllocName = ClFuncPrefix;515 if (Options.FastABI)516 TokenAllocName += utostr(TokenID) + "_";517 else518 NewParams.push_back(IntPtrTy); // token ID519 TokenAllocName += Callee->getName();520 FunctionType *NewFTy = FunctionType::get(RetTy, NewParams, false);521 FunctionCallee TokenAlloc = Mod.getOrInsertFunction(TokenAllocName, NewFTy);522 if (Function *F = dyn_cast<Function>(TokenAlloc.getCallee()))523 F->copyAttributesFrom(Callee); // preserve attrs524 525 if (Key.has_value())526 TokenAllocFunctions[*Key] = TokenAlloc;527 return TokenAlloc;528}529 530void AllocToken::replaceIntrinsicInst(IntrinsicInst *II,531 OptimizationRemarkEmitter &ORE) {532 assert(II->getIntrinsicID() == Intrinsic::alloc_token_id);533 534 uint64_t TokenID = getToken(*II, ORE);535 Value *V = ConstantInt::get(IntPtrTy, TokenID);536 II->replaceAllUsesWith(V);537 II->eraseFromParent();538}539 540} // namespace541 542AllocTokenPass::AllocTokenPass(AllocTokenOptions Opts)543 : Options(std::move(Opts)) {}544 545PreservedAnalyses AllocTokenPass::run(Module &M, ModuleAnalysisManager &MAM) {546 AllocToken Pass(Options, M, MAM);547 bool Modified = false;548 549 for (Function &F : M) {550 if (F.empty())551 continue; // declaration552 Modified |= Pass.instrumentFunction(F);553 }554 555 return Modified ? PreservedAnalyses::none().preserveSet<CFGAnalyses>()556 : PreservedAnalyses::all();557}558