topjohnwu/CPI
1
0
Fork 0
mirror of https://github.com/topjohnwu/CPI.git synced 2024-11-27 05:20:31 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

Rewrite runtime __sm_alloca

Browse Source 
Change safe memory alloca behavior: directly return memory address instead of returning index
This commit is contained in:
topjohnwu 2018-12-15 14:31:34 -05:00
parent 03f49c2022
commit da6e5ebedb
2 changed files with 37 additions and 48 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

58
pass/CPI.cpp
View File

@ -29,11 +29,10 @@ struct CPI : public ModulePass {
voidPPT = PointerType::get(voidPT, 0);
// Add global variables in libsafe_rt
smPool = new GlobalVariable(M, voidPPT, false, GlobalValue::ExternalLinkage, nullptr, "__sm_pool");
smSp = new GlobalVariable(M, intT, false, GlobalValue::ExternalLinkage, nullptr, "__sm_sp");
// Add function references in libsafe_rt
smAlloca = cast<Function>(M.getOrInsertFunction("__sm_alloca", intT));
smAlloca = cast<Function>(M.getOrInsertFunction("__sm_alloca", voidPPT));
// Find all sensitive structs
for (auto s : M.getIdentifiedStructTypes()) {
@ -56,7 +55,6 @@ struct CPI : public ModulePass {
private:
Function *smAlloca;
Value *smPool;
Value *smSp;
IntegerType *intT;
@ -72,24 +70,11 @@ private:
BasicBlock &entryBlock = F.getEntryBlock();
// Create a placeholder load for __sm_pool
auto poolAdress = new LoadInst(smPool, "smPoolAddress", entryBlock.getFirstNonPHI());
// Alloca only happens in the first basic block
hasInject |= swapFunctionPtrAlloca(entryBlock, poolAdress);
hasInject |= handleStructAlloca(entryBlock, poolAdress);
hasInject |= swapFunctionPtrAlloca(entryBlock);
hasInject |= handleStructAlloca(entryBlock);
if (hasInject) {
// Find the last __sm_alloca, and move __sm_pool load after it
CallInst *lastAlloca = nullptr;
for (auto &I : entryBlock) {
CallInst *c;
if ((c = dyn_cast<CallInst>(&I)) && c->getCalledFunction() == smAlloca)
lastAlloca = c;
}
assert(lastAlloca); /* hasInject == true implies lastAlloca != nullptr */
poolAdress->moveAfter(lastAlloca);
// Create checkpoint
auto spLoad = new LoadInst(smSp, "smStackCheckpoint", entryBlock.getFirstNonPHI());
for (auto &bb : F) {
@ -99,40 +84,35 @@ private:
new StoreInst(spLoad, smSp, ti);
}
}
} else {
// The additional load is not needed
poolAdress->eraseFromParent();
}
}
void swapStore(Instruction *idx, StoreInst *store, LoadInst *poolAddress) {
void swapStore(Instruction *addr, StoreInst *store) {
IRBuilder<> b(store);
auto off = b.CreateGEP(voidPT, poolAddress, idx);
auto cast = b.CreatePointerCast(store->getValueOperand(), voidPT);
b.CreateStore(cast, off);
b.CreateStore(cast, addr);
DEBUG(dbgs() << "SWAP:" << *store << "\n");
store->eraseFromParent();
}
void swapLoad(Instruction *idx, LoadInst *load, LoadInst *poolAddress) {
void swapLoad(Instruction *addr, LoadInst *load) {
IRBuilder<> b(load);
auto off = b.CreateGEP(voidPT, poolAddress, idx);
auto raw = b.CreateLoad(off);
auto raw = b.CreateLoad(addr);
auto cast = b.CreatePointerCast(raw, load->getType());
DEBUG(dbgs() << "SWAP:" << *load << "\n");
BasicBlock::iterator ii(load);
ReplaceInstWithValue(load->getParent()->getInstList(), ii, cast);
}
void swapPtr(Instruction *from, Instruction *to, LoadInst *poolAddress) {
void swapPtr(Instruction *from, Instruction *to) {
// Swap out all uses (store and load)
for (auto a : from->users()) {
StoreInst *s;
LoadInst *l;
if ((s = dyn_cast<StoreInst>(a))) {
swapStore(to, s, poolAddress);
swapStore(to, s);
} else if ((l = dyn_cast<LoadInst>(a))) {
swapLoad(to, l, poolAddress);
swapLoad(to, l);
} else {
/* TODO: Dunno when this will happen, spit logs for now */
DEBUG(dbgs() << "Unknown:" << *from << "\n");
@ -144,18 +124,20 @@ private:
}
}
bool swapFunctionPtrAlloca(BasicBlock &bb, LoadInst *poolAddress) {
bool swapFunctionPtrAlloca(BasicBlock &bb) {
auto v = getFunctionPtrAlloca(bb);
for (auto alloc : v) {
IRBuilder<> b(alloc);
auto idx = b.CreateCall(smAlloca, None, alloc->getName());
DEBUG(dbgs() << "ADD:" << *idx << "\n");
swapPtr(alloc, idx, poolAddress);
std::string name(alloc->getName());
auto addr = b.CreateCall(smAlloca);
DEBUG(dbgs() << "ADD:" << *addr << "\n");
swapPtr(alloc, addr);
addr->setName(name);
}
return !v.empty();
}
bool handleStructAlloca(BasicBlock &bb, LoadInst *poolAddress) {
bool handleStructAlloca(BasicBlock &bb) {
bool hasInject = false;
for (auto alloc: getSSAlloca(bb)) {
auto elist = ssMap.find(alloc->getAllocatedType());
@ -170,10 +152,10 @@ private:
if (!rmList.empty()) {
hasInject = true;
IRBuilder<> b(alloc->getNextNode());
auto idx = b.CreateCall(smAlloca, None, alloc->getName() + "." + std::to_string(fpentry));
DEBUG(dbgs() << "ADD:" << *idx << "\n");
auto addr = b.CreateCall(smAlloca, None, alloc->getName() + "." + std::to_string(fpentry));
DEBUG(dbgs() << "ADD:" << *addr << "\n");
for (auto u: rmList) {
swapPtr(u, idx, poolAddress);
swapPtr(u, addr);
}
}
}

27
safe_rt/rt.c
View File

@ -1,18 +1,25 @@
#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#define INIT_SZ 1
#define BLOCK_SZ 16
#define PTR_MEM_SZ(i) ((i) * sizeof(void *))
void **__sm_pool;
int __sm_sp = 0;
static int pool_sz = 0;
static void ***block_table;
static int table_sz = 0;
int __sm_alloca() {
if (__sm_sp == pool_sz) {
pool_sz = pool_sz ? pool_sz * 2 : INIT_SZ;
__sm_pool = (void **) realloc(__sm_pool, pool_sz * sizeof(void *));
void **__sm_alloca() {
int block_num = __sm_sp >> 4;
if (block_num == table_sz) {
int new_sz = table_sz ? table_sz * 2 : 1;
block_table = realloc(block_table, PTR_MEM_SZ(table_sz));
// Zero out the new space
memset(&block_table[table_sz], 0, PTR_MEM_SZ(new_sz - table_sz));
table_sz = new_sz;
}
// Debug message
fprintf(stderr, "__sm_alloca: %d, %d\n", __sm_sp, pool_sz);
return __sm_sp++;
if (!block_table[block_num])
block_table[block_num] = malloc(PTR_MEM_SZ(BLOCK_SZ));
fprintf(stderr, "__sm_alloca %d, %d\n", __sm_sp, table_sz << 4);
return &block_table[block_num][__sm_sp++ & 0xF];
}