Skip to content

Commit 82252b5

Browse files
mamaria-kmamaria-k
committed
Add handlers and fix
1 parent 40d4a65 commit 82252b5

File tree

9 files changed

+190
-71
lines changed

9 files changed

+190
-71
lines changed

‎include/klee/Module/TaintAnnotation.h

+3-12
Original file line numberDiff line numberDiff line change
@@ -14,19 +14,10 @@ using source_ty = size_t;
1414
using sink_ty = size_t;
1515
using rule_ty = size_t;
1616

17-
struct TaintHitInfo final {
18-
source_ty source;
19-
sink_ty sink;
20-
21-
explicit TaintHitInfo(source_ty source, sink_ty sink);
22-
23-
bool operator<(const TaintHitInfo &other) const;
24-
bool operator==(const TaintHitInfo &other) const;
25-
};
26-
27-
using TaintHitsMap = std::map<TaintHitInfo, rule_ty>;
28-
2917
struct TaintAnnotation final {
18+
using TaintHitsSink = std::map<source_ty, rule_ty>;
19+
using TaintHitsMap = std::map<sink_ty, TaintHitsSink>;
20+
3021
TaintHitsMap hits;
3122

3223
std::unordered_map<std::string, source_ty> sources;

‎include/klee/klee.h

+1-1
Original file line numberDiff line numberDiff line change
@@ -67,7 +67,7 @@ bool klee_check_taint_source(void *array, size_t taint_source);
6767
* \arg addr - The start of the object.
6868
* \arg taint_sink - Taint sink.
6969
*/
70-
size_t klee_get_taint_rule(void *array, size_t taint_sink);
70+
uint64_t klee_get_taint_rule(void *array, size_t taint_sink);
7171

7272
/* klee_taint_hit - Execute taint hit with rule.
7373
*

‎lib/Core/Executor.cpp

+128-1
Original file line numberDiff line numberDiff line change
@@ -5028,7 +5028,7 @@ void Executor::terminateStateOnTargetError(ExecutionState &state,
50285028
}
50295029

50305030
void Executor::terminateStateOnTargetTaintError(ExecutionState &state,
5031-
size_t rule) {
5031+
uint64_t rule) {
50325032
if (rule >= annotationsData.taintAnnotation.rules.size()) {
50335033
terminateStateOnUserError(state, "Incorrect rule id");
50345034
}
@@ -5541,6 +5541,133 @@ void Executor::executeFree(ExecutionState &state, ref<PointerExpr> address,
55415541
}
55425542
}
55435543

5544+
void Executor::executeChangeTaintSource(ExecutionState &state,
5545+
klee::KInstruction *target,
5546+
ref<PointerExpr> address,
5547+
uint64_t source, bool isAdd) {
5548+
address = optimizer.optimizeExpr(address, true);
5549+
ref<Expr> isNullPointer = Expr::createIsZero(address->getValue());
5550+
StatePair zeroPointer =
5551+
forkInternal(state, isNullPointer, BranchType::ResolvePointer);
5552+
if (zeroPointer.first) {
5553+
auto error =
5554+
(isReadFromSymbolicArray(address->getBase()) && zeroPointer.second)
5555+
? ReachWithErrorType::MayBeNullPointerException
5556+
: ReachWithErrorType::MustBeNullPointerException;
5557+
terminateStateOnTargetError(*zeroPointer.first, ReachWithError(error));
5558+
}
5559+
if (zeroPointer.second) { // address != 0
5560+
ExactResolutionList rl;
5561+
resolveExact(*zeroPointer.second, address,
5562+
typeSystemManager->getUnknownType(), rl, "сhangeTaintSource");
5563+
for (Executor::ExactResolutionList::iterator it = rl.begin(), ie = rl.end();
5564+
it != ie; ++it) {
5565+
const MemoryObject *mo = it->first;
5566+
RefObjectPair op =
5567+
it->second->addressSpace.findOrLazyInitializeObject(mo);
5568+
ref<const ObjectState> os = op.second;
5569+
5570+
ObjectState *wos = it->second->addressSpace.getWriteable(mo, os.get());
5571+
if (wos->readOnly) {
5572+
terminateStateOnProgramError(*(it->second),
5573+
"memory error: object read only",
5574+
StateTerminationType::ReadOnly);
5575+
} else {
5576+
wos->updateTaint(Expr::createTaintBySource(source), isAdd);
5577+
}
5578+
}
5579+
}
5580+
}
5581+
5582+
void Executor::executeCheckTaintSource(ExecutionState &state,
5583+
klee::KInstruction *target,
5584+
ref<PointerExpr> address,
5585+
uint64_t source) {
5586+
address = optimizer.optimizeExpr(address, true);
5587+
ref<Expr> isNullPointer = Expr::createIsZero(address->getValue());
5588+
StatePair zeroPointer =
5589+
forkInternal(state, isNullPointer, BranchType::ResolvePointer);
5590+
if (zeroPointer.first) {
5591+
auto error =
5592+
(isReadFromSymbolicArray(address->getBase()) && zeroPointer.second)
5593+
? ReachWithErrorType::MayBeNullPointerException
5594+
: ReachWithErrorType::MustBeNullPointerException;
5595+
terminateStateOnTargetError(*zeroPointer.first, ReachWithError(error));
5596+
}
5597+
if (zeroPointer.second) {
5598+
ExactResolutionList rl;
5599+
resolveExact(*zeroPointer.second, address,
5600+
typeSystemManager->getUnknownType(), rl, "checkTaintSource");
5601+
5602+
for (Executor::ExactResolutionList::iterator it = rl.begin(), ie = rl.end();
5603+
it != ie; ++it) {
5604+
const MemoryObject *mo = it->first;
5605+
RefObjectPair op =
5606+
it->second->addressSpace.findOrLazyInitializeObject(mo);
5607+
ref<const ObjectState> os = op.second;
5608+
5609+
ref<Expr> taintSource =
5610+
ExtractExpr::create(os->readTaint(), source, Expr::Bool);
5611+
bindLocal(target, *it->second, taintSource);
5612+
}
5613+
}
5614+
}
5615+
5616+
void Executor::executeGetTaintRule(ExecutionState &state,
5617+
klee::KInstruction *target,
5618+
ref<PointerExpr> address, uint64_t sink) {
5619+
const auto &hitsBySink = annotationsData.taintAnnotation.hits[sink];
5620+
ref<Expr> hitsBySinkTaint = Expr::createEmptyTaint();
5621+
for (const auto [source, rule] : hitsBySink) {
5622+
hitsBySinkTaint =
5623+
OrExpr::create(hitsBySinkTaint, Expr::createTaintBySource(source));
5624+
}
5625+
5626+
address = optimizer.optimizeExpr(address, true);
5627+
ref<Expr> isNullPointer = Expr::createIsZero(address->getValue());
5628+
StatePair zeroPointer =
5629+
forkInternal(state, isNullPointer, BranchType::ResolvePointer);
5630+
if (zeroPointer.first) {
5631+
auto error =
5632+
(isReadFromSymbolicArray(address->getBase()) && zeroPointer.second)
5633+
? ReachWithErrorType::MayBeNullPointerException
5634+
: ReachWithErrorType::MustBeNullPointerException;
5635+
terminateStateOnTargetError(*zeroPointer.first, ReachWithError(error));
5636+
}
5637+
if (zeroPointer.second) {
5638+
ExactResolutionList rl;
5639+
resolveExact(*zeroPointer.second, address,
5640+
typeSystemManager->getUnknownType(), rl, "getTaintRule");
5641+
5642+
for (Executor::ExactResolutionList::iterator it = rl.begin(), ie = rl.end();
5643+
it != ie; ++it) {
5644+
const MemoryObject *mo = it->first;
5645+
RefObjectPair op =
5646+
it->second->addressSpace.findOrLazyInitializeObject(mo);
5647+
ref<const ObjectState> os = op.second;
5648+
5649+
ref<Expr> hits = AndExpr::create(os->readTaint(), hitsBySinkTaint);
5650+
5651+
auto curState = it->second;
5652+
for (size_t source = 0;
5653+
source < annotationsData.taintAnnotation.sources.size(); source++) {
5654+
ref<Expr> taintSource = ExtractExpr::create(hits, source, Expr::Bool);
5655+
StatePair taintSourceStates =
5656+
forkInternal(*curState, taintSource, BranchType::Taint);
5657+
if (taintSourceStates.first) {
5658+
bindLocal(target, *taintSourceStates.first,
5659+
ConstantExpr::create(hitsBySink.at(source) + 1,
5660+
Expr::Int64)); // return (rule + 1)
5661+
}
5662+
if (taintSourceStates.second) {
5663+
curState = taintSourceStates.second;
5664+
}
5665+
}
5666+
bindLocal(target, *curState, ConstantExpr::create(0, Expr::Int64));
5667+
}
5668+
}
5669+
}
5670+
55445671
bool Executor::resolveExact(ExecutionState &estate, ref<Expr> address,
55455672
KType *type, ExactResolutionList &results,
55465673
const std::string &name) {

‎lib/Core/Executor.h

+13-1
Original file line numberDiff line numberDiff line change
@@ -361,6 +361,18 @@ class Executor : public Interpreter {
361361
void executeFree(ExecutionState &state, ref<PointerExpr> address,
362362
KInstruction *target = 0);
363363

364+
void executeChangeTaintSource(ExecutionState &state,
365+
klee::KInstruction *target,
366+
ref<PointerExpr> address, uint64_t source,
367+
bool isAdd);
368+
369+
void executeCheckTaintSource(ExecutionState &state,
370+
klee::KInstruction *target,
371+
ref<PointerExpr> address, uint64_t source);
372+
373+
void executeGetTaintRule(ExecutionState &state, klee::KInstruction *target,
374+
ref<PointerExpr> address, uint64_t sink);
375+
364376
/// Serialize a landingpad instruction so it can be handled by the
365377
/// libcxxabi-runtime
366378
MemoryObject *serializeLandingpad(ExecutionState &state,
@@ -634,7 +646,7 @@ class Executor : public Interpreter {
634646
/// Then just call `terminateStateOnError`
635647
void terminateStateOnTargetError(ExecutionState &state, ReachWithError error);
636648

637-
void terminateStateOnTargetTaintError(ExecutionState &state, size_t rule);
649+
void terminateStateOnTargetTaintError(ExecutionState &state, uint64_t rule);
638650

639651
/// Call error handler and terminate state in case of program errors
640652
/// (e.g. free()ing globals, out-of-bound accesses)

‎lib/Core/Memory.cpp

+2-4
Original file line numberDiff line numberDiff line change
@@ -476,8 +476,7 @@ ref<Expr> ObjectState::readTaint(ref<Expr> offset, Expr::Width width) const {
476476
// Otherwise, follow the slow general case.
477477
unsigned NumBytes = width / 8;
478478
assert(width == NumBytes * 8 && "Invalid read size!");
479-
ref<Expr> Res(0);
480-
ref<Expr> null = Expr::createPointer(0);
479+
ref<Expr> Res = Expr::createEmptyTaint();
481480
for (unsigned i = 0; i != NumBytes; ++i) {
482481
unsigned idx = Context::get().isLittleEndian() ? i : (NumBytes - i - 1);
483482
ref<Expr> Byte = readTaint8(
@@ -496,8 +495,7 @@ ref<Expr> ObjectState::readTaint(unsigned offset, Expr::Width width) const {
496495
// Otherwise, follow the slow general case.
497496
unsigned NumBytes = width / 8;
498497
assert(width == NumBytes * 8 && "Invalid width for read size!");
499-
ref<Expr> Res(0);
500-
ref<Expr> null = Expr::createPointer(0);
498+
ref<Expr> Res = Expr::createEmptyTaint();
501499
for (unsigned i = 0; i != NumBytes; ++i) {
502500
unsigned idx = Context::get().isLittleEndian() ? i : (NumBytes - i - 1);
503501
ref<Expr> Byte = readTaint8(offset + idx);

‎lib/Core/MockBuilder.cpp

+9-9
Original file line numberDiff line numberDiff line change
@@ -611,7 +611,7 @@ void MockBuilder::buildAnnotationForExternalFunctionArgs(
611611
Statement::InitNull *initNullPtr = nullptr;
612612

613613
bool isMocked = false;
614-
std::string mockName = "klee_mock" + func->getName().str() + "_arg_" +
614+
std::string mockName = "klee_mock_" + func->getName().str() + "_arg_" +
615615
std::to_string(i) + "_" +
616616
std::to_string(offsetIndex);
617617

@@ -681,10 +681,10 @@ void MockBuilder::buildAnnotationForExternalFunctionArgs(
681681
klee_error("Annotation: TaintOutput arg is not pointer");
682682
}
683683

684-
if (!isMocked) {
685-
buildCallKleeMakeMockAll(elem, mockName);
686-
isMocked = true;
687-
}
684+
// if (!isMocked) {
685+
// buildCallKleeMakeMockAll(elem, mockName);
686+
// isMocked = true;
687+
// }
688688
buildAnnotationTaintOutput(elem, statement);
689689
break;
690690
}
@@ -693,10 +693,10 @@ void MockBuilder::buildAnnotationForExternalFunctionArgs(
693693
klee_error("Annotation: TaintPropagation arg is not pointer");
694694
}
695695

696-
if (!isMocked) {
697-
buildCallKleeMakeMockAll(elem, mockName);
698-
isMocked = true;
699-
}
696+
// if (!isMocked) {
697+
// buildCallKleeMakeMockAll(elem, mockName);
698+
// isMocked = true;
699+
// }
700700
buildAnnotationTaintPropagation(elem, statement, func,
701701
"_arg_" + std::to_string(i) + "_");
702702
break;

‎lib/Core/SpecialFunctionHandler.cpp

+19-19
Original file line numberDiff line numberDiff line change
@@ -1298,10 +1298,10 @@ void SpecialFunctionHandler::handleAddTaint(klee::ExecutionState &state,
12981298
return;
12991299
}
13001300

1301-
// uint64_t taintSource = dyn_cast<ConstantExpr>(arguments[1])->getZExtValue();
1302-
// printf("klee_add_taint source: %zu\n", taintSource);
1303-
// executor.executeChangeTaintSource(
1304-
// state, target, executor.makePointer(arguments[0]), taintSource, true);
1301+
uint64_t taintSource = dyn_cast<ConstantExpr>(arguments[1])->getZExtValue();
1302+
printf("klee_add_taint source: %zu\n", taintSource);
1303+
executor.executeChangeTaintSource(
1304+
state, target, executor.makePointer(arguments[0]), taintSource, true);
13051305
}
13061306

13071307
void SpecialFunctionHandler::handleClearTaint(
@@ -1314,10 +1314,10 @@ void SpecialFunctionHandler::handleClearTaint(
13141314
return;
13151315
}
13161316

1317-
// uint64_t taintSource = dyn_cast<ConstantExpr>(arguments[1])->getZExtValue();
1318-
// printf("klee_clear_taint source: %zu\n", taintSource);
1319-
// executor.executeChangeTaintSource(
1320-
// state, target, executor.makePointer(arguments[0]), taintSource, false);
1317+
uint64_t taintSource = dyn_cast<ConstantExpr>(arguments[1])->getZExtValue();
1318+
printf("klee_clear_taint source: %zu\n", taintSource);
1319+
executor.executeChangeTaintSource(
1320+
state, target, executor.makePointer(arguments[0]), taintSource, false);
13211321
}
13221322

13231323
void SpecialFunctionHandler::handleCheckTaintSource(
@@ -1330,10 +1330,10 @@ void SpecialFunctionHandler::handleCheckTaintSource(
13301330
return;
13311331
}
13321332

1333-
// uint64_t taintSource = dyn_cast<ConstantExpr>(arguments[1])->getZExtValue();
1334-
// printf("klee_check_taint_source source: %zu\n", taintSource);
1335-
// executor.executeCheckTaintSource(
1336-
// state, target, executor.makePointer(arguments[0]), taintSource);
1333+
uint64_t taintSource = dyn_cast<ConstantExpr>(arguments[1])->getZExtValue();
1334+
printf("klee_check_taint_source source: %zu\n", taintSource);
1335+
executor.executeCheckTaintSource(
1336+
state, target, executor.makePointer(arguments[0]), taintSource);
13371337
}
13381338

13391339
void SpecialFunctionHandler::handleGetTaintRule(
@@ -1346,14 +1346,14 @@ void SpecialFunctionHandler::handleGetTaintRule(
13461346
return;
13471347
}
13481348

1349-
// // TODO: now mock
1350-
ref<Expr> result = ConstantExpr::create(1, Expr::Int64);
1351-
executor.bindLocal(target, state, result);
1349+
// // // TODO: now mock
1350+
// ref<Expr> result = ConstantExpr::create(1, Expr::Int64);
1351+
// executor.bindLocal(target, state, result);
13521352

1353-
// uint64_t taintSink = dyn_cast<ConstantExpr>(arguments[1])->getZExtValue();
1354-
// printf("klee_get_taint_rule source: %zu\n", taintSink);
1355-
// executor.executeGetTaintRule(state, target,
1356-
// executor.makePointer(arguments[0]), taintSink);
1353+
uint64_t taintSink = dyn_cast<ConstantExpr>(arguments[1])->getZExtValue();
1354+
printf("klee_get_taint_rule sink: %zu\n", taintSink);
1355+
executor.executeGetTaintRule(state, target,
1356+
executor.makePointer(arguments[0]), taintSink);
13571357
}
13581358

13591359
void SpecialFunctionHandler::handleTaintHit(klee::ExecutionState &state,

‎lib/Expr/Expr.cpp

+12-12
Original file line numberDiff line numberDiff line change
@@ -559,18 +559,18 @@ ref<ConstantExpr> Expr::createTaintBySource(uint64_t source) {
559559

560560
ref<Expr> Expr::combineTaints(const ref<Expr> &taintL,
561561
const ref<Expr> &taintR) {
562-
if (SelectExpr *sel = dyn_cast<SelectExpr>(taintL)) {
563-
taintL->dump();
564-
}
565-
if (PointerExpr *sel = dyn_cast<PointerExpr>(taintL)) {
566-
taintR->dump();
567-
}
568-
if (ConstantExpr *sel = dyn_cast<ConstantExpr>(taintL)) {
569-
if (ConstantExpr *ser = dyn_cast<ConstantExpr>(taintR)) {
570-
sel->getAPValue().dump();
571-
ser->getAPValue().dump();
572-
}
573-
}
562+
// if (SelectExpr *sel = dyn_cast<SelectExpr>(taintL)) {
563+
// taintL->dump();
564+
// }
565+
// if (PointerExpr *sel = dyn_cast<PointerExpr>(taintL)) {
566+
// taintR->dump();
567+
// }
568+
// if (ConstantExpr *sel = dyn_cast<ConstantExpr>(taintL)) {
569+
// if (ConstantExpr *ser = dyn_cast<ConstantExpr>(taintR)) {
570+
// sel->getAPValue().dump();
571+
// ser->getAPValue().dump();
572+
// }
573+
// }
574574
return OrExpr::create(taintL, taintR);
575575
}
576576

‎lib/Module/TaintAnnotation.cpp

+3-12
Original file line numberDiff line numberDiff line change
@@ -5,16 +5,6 @@
55

66
namespace klee {
77

8-
TaintHitInfo::TaintHitInfo(source_ty source, sink_ty sink)
9-
: source(source), sink(sink) {}
10-
11-
bool TaintHitInfo::operator<(const TaintHitInfo &other) const {
12-
return (sink < other.sink) || (source < other.source);
13-
}
14-
bool TaintHitInfo::operator==(const TaintHitInfo &other) const {
15-
return (source == other.source) && (sink == other.sink);
16-
}
17-
188
TaintAnnotation::TaintAnnotation(const std::string &path) {
199
if (path.empty()) {
2010
return;
@@ -60,11 +50,12 @@ TaintAnnotation::TaintAnnotation(const std::string &path) {
6050
klee_error("Taint annotations: Incorrect file format");
6151
}
6252

53+
std::map<source_ty, rule_ty> hitsForSink;
6354
for (auto &taintHitJson : item.value()) {
64-
hits[TaintHitInfo(sources[taintHitJson["source"]], sinksCounter)] =
55+
hitsForSink[sources[taintHitJson["source"]]] =
6556
rulesMap[taintHitJson["rule"]];
6657
}
67-
58+
hits[sinksCounter] = hitsForSink;
6859
sinksCounter++;
6960
}
7061
}

0 commit comments

Comments
 (0)