Skip to content

Commit bc922f5

Browse files
mamaria-kmamaria-k
committed
Add taint rules
1 parent 6a872a2 commit bc922f5

14 files changed

Lines changed: 135 additions & 84 deletions

‎configs/taint-annotations.json‎

Lines changed: 12 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,20 @@
11
{
22
"SensitiveDataLeak": [
3-
"SensitiveDataSource"
3+
{
4+
"source": "SensitiveDataSource",
5+
"rule": "TAINT.SDE"
6+
}
47
],
58
"FormatString": [
6-
"UntrustedSource"
9+
{
10+
"source": "UntrustedSource",
11+
"rule": "SV.STR.FMT.TAINT"
12+
}
713
],
814
"Execute": [
9-
"UntrustedSource"
15+
{
16+
"source": "UntrustedSource",
17+
"rule": "TAINT.STRING.CLI"
18+
}
1019
]
1120
}

‎include/klee/Core/MockBuilder.h‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -50,8 +50,8 @@ class MockBuilder {
5050
const std::string &symbolicName);
5151
llvm::CallInst *buildCallKleeTaintFunction(const std::string &functionName,
5252
llvm::Value *source, size_t taint,
53-
bool returnBool);
54-
void buildCallKleeTaintSinkHit(size_t taintSink);
53+
llvm::Type *returnType);
54+
void buildCallKleeTaintHit(llvm::Value *taintRule);
5555

5656
void buildAnnotationTaintOutput(llvm::Value *elem,
5757
const Statement::Ptr &statement);

‎include/klee/Module/SarifReport.h‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -51,6 +51,10 @@ enum ReachWithError {
5151
NullCheckAfterDerefException,
5252
Reachable,
5353
None,
54+
55+
TaintFormatString,
56+
TaintSensitiveData,
57+
TaintExecute,
5458
};
5559

5660
const char *getErrorString(ReachWithError error);

‎include/klee/Module/TaintAnnotation.h‎

Lines changed: 21 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -10,18 +10,30 @@ using json = nlohmann::json;
1010

1111
namespace klee {
1212

13-
// TODO: set<size_t> to set<pair<size_t, size)t>>, where pair is sink and target
14-
using TaintSinksSourcesMap = std::map<size_t, std::set<size_t>>;
13+
using source_ty = size_t;
14+
using sink_ty = size_t;
15+
using rule_ty = size_t;
1516

16-
class TaintAnnotation final {
17-
public:
18-
TaintSinksSourcesMap sinksToSources;
19-
std::map<std::string, size_t> sinks;
20-
std::map<std::string, size_t> sources;
21-
// TODO: add map from size_t target to string
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+
29+
struct TaintAnnotation final {
30+
TaintHitsMap hits;
31+
32+
std::unordered_map<std::string, source_ty> sources;
33+
std::unordered_map<std::string, sink_ty> sinks;
34+
std::vector<std::string> rules;
2235

2336
explicit TaintAnnotation(const std::string &path);
24-
virtual ~TaintAnnotation();
2537
};
2638

2739
} // namespace klee

‎include/klee/klee.h‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -61,19 +61,19 @@ void klee_clear_taint(void *array, size_t taint_source);
6161
*/
6262
bool klee_check_taint_source(void *array, size_t taint_source);
6363

64-
/* klee_check_taint_sink - Check that the contents of the object pointer
65-
* to \arg addr causes a taint sink \arg taint_sink hit.
64+
/* klee_get_taint_rule - Return taint rule id +1 if contents of the object pointer
65+
* to \arg addr causes a taint sink \arg taint_sink hit, else return 0.
6666
*
6767
* \arg addr - The start of the object.
6868
* \arg taint_sink - Taint sink.
6969
*/
70-
bool klee_check_taint_sink(void *array, size_t taint_sink);
70+
size_t klee_get_taint_rule(void *array, size_t taint_sink);
7171

72-
/* klee_taint_sink_hit - Execute taint sink \arg taint_sink hit.
72+
/* klee_taint_hit - Execute taint hit with rule.
7373
*
74-
* \arg taint_sink - Taint sink.
74+
* \arg rule - Taint rule id.
7575
*/
76-
void klee_taint_sink_hit(size_t taint_sink);
76+
void klee_taint_hit(size_t rule);
7777

7878
/* klee_range - Construct a symbolic value in the signed interval
7979
* [begin,end).

‎lib/Core/Executor.cpp‎

Lines changed: 6 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -5026,18 +5026,12 @@ void Executor::terminateStateOnTargetError(ExecutionState &state,
50265026
}
50275027

50285028
// TODO: add taint target errors to taint-annotations.json and change function
5029-
void Executor::terminateStateOnTaintError(ExecutionState &state, size_t sink) {
5030-
// reportStateOnTargetError(state, error);
5031-
5032-
const auto &sinks = annotationsData.taintAnnotation.sinks;
5033-
auto taintSink = std::find_if(sinks.begin(), sinks.end(),
5034-
[&](const std::pair<std::string, size_t> &i) {
5035-
return i.second == sink;
5036-
});
5037-
if (taintSink == std::end(sinks)) {
5038-
klee_error("Unknown sink index");
5039-
}
5040-
terminateStateOnProgramError(state, taintSink->first + " taint error",
5029+
void Executor::terminateStateOnTargetTaintError(ExecutionState &state, size_t rule) {
5030+
const std::string &ruleStr = annotationsData.taintAnnotation.rules[rule];
5031+
5032+
// reportStateOnTargetError(state, rule);
5033+
5034+
terminateStateOnProgramError(state, ruleStr + " taint error",
50415035
StateTerminationType::Taint);
50425036
}
50435037

‎lib/Core/Executor.h‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -634,7 +634,7 @@ class Executor : public Interpreter {
634634
/// Then just call `terminateStateOnError`
635635
void terminateStateOnTargetError(ExecutionState &state, ReachWithError error);
636636

637-
void terminateStateOnTaintError(ExecutionState &state, size_t sink);
637+
void terminateStateOnTargetTaintError(ExecutionState &state, size_t rule);
638638

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

‎lib/Core/MockBuilder.cpp‎

Lines changed: 28 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -421,7 +421,8 @@ void MockBuilder::buildAnnotationTaintOutput(llvm::Value *elem,
421421
taintOutputPtr->getTaintType().c_str());
422422
return;
423423
}
424-
buildCallKleeTaintFunction("klee_add_taint", elem, source->second, false);
424+
buildCallKleeTaintFunction("klee_add_taint", elem, source->second,
425+
llvm::Type::getVoidTy(mockModule->getContext()));
425426
}
426427

427428
void MockBuilder::buildAnnotationTaintPropagation(
@@ -460,11 +461,13 @@ void MockBuilder::buildAnnotationTaintPropagation(
460461
llvm::Value *propagationValue =
461462
func->getArg(taintPropagationPtr->propagationParameterIndex);
462463
auto brValuePropagate = buildCallKleeTaintFunction(
463-
"klee_check_taint_source", propagationValue, source->second, true);
464+
"klee_check_taint_source", propagationValue, source->second,
465+
llvm::Type::getInt1Ty(mockModule->getContext()));
464466
builder->CreateCondBr(brValuePropagate, propagateBB, contBB);
465467

466468
builder->SetInsertPoint(propagateBB);
467-
buildCallKleeTaintFunction("klee_add_taint", elem, source->second, false);
469+
buildCallKleeTaintFunction("klee_add_taint", elem, source->second,
470+
llvm::Type::getVoidTy(mockModule->getContext()));
468471
builder->CreateBr(contBB);
469472

470473
curFunc->getBasicBlockList().push_back(contBB);
@@ -495,8 +498,13 @@ void MockBuilder::buildAnnotationTaintSink(llvm::Value *elem,
495498
llvm::BasicBlock::Create(mockModule->getContext(), sinkCondName, curFunc);
496499
llvm::BasicBlock *contBB = llvm::BasicBlock::Create(
497500
mockModule->getContext(), "continue_" + sinkCondName);
498-
auto brValueSink = buildCallKleeTaintFunction("klee_check_taint_sink", elem,
499-
sink->second, true);
501+
auto taintRule = buildCallKleeTaintFunction(
502+
"klee_get_taint_rule", elem, sink->second,
503+
llvm::Type::getInt64Ty(mockModule->getContext()));
504+
const auto brValueSink =
505+
builder->CreateCmp(llvm::CmpInst::Predicate::ICMP_NE, taintRule,
506+
llvm::ConstantInt::get(mockModule->getContext(),
507+
llvm::APInt(64, 0, false)));
500508
builder->CreateCondBr(brValueSink, sinkBB, contBB);
501509

502510
builder->SetInsertPoint(sinkBB);
@@ -509,13 +517,16 @@ void MockBuilder::buildAnnotationTaintSink(llvm::Value *elem,
509517
sinkHitCondName);
510518
fromIf = builder->GetInsertBlock();
511519
curFunc = fromIf->getParent();
512-
llvm::BasicBlock *sinkHitBB = llvm::BasicBlock::Create(
520+
llvm::BasicBlock *taintHitBB = llvm::BasicBlock::Create(
513521
mockModule->getContext(), sinkHitCondName, curFunc);
514-
auto brValueSinkHit = builder->CreateLoad(intType, sinkHitCond);
515-
builder->CreateCondBr(brValueSinkHit, sinkHitBB, contBB);
516-
517-
builder->SetInsertPoint(sinkHitBB);
518-
buildCallKleeTaintSinkHit(sink->second);
522+
auto brValueTaintHit = builder->CreateLoad(intType, sinkHitCond);
523+
builder->CreateCondBr(brValueTaintHit, taintHitBB, contBB);
524+
525+
builder->SetInsertPoint(taintHitBB);
526+
const auto taintRuleId = builder->CreateSub(
527+
taintRule, llvm::ConstantInt::get(mockModule->getContext(),
528+
llvm::APInt(64, 1, false)));
529+
buildCallKleeTaintHit(taintRuleId);
519530
builder->CreateBr(contBB);
520531

521532
curFunc->getBasicBlockList().push_back(contBB);
@@ -525,11 +536,7 @@ void MockBuilder::buildAnnotationTaintSink(llvm::Value *elem,
525536
llvm::CallInst *
526537
MockBuilder::buildCallKleeTaintFunction(const std::string &functionName,
527538
llvm::Value *source, size_t taint,
528-
bool returnBool) {
529-
const auto returnType = returnBool
530-
? llvm::Type::getInt1Ty(mockModule->getContext())
531-
: llvm::Type::getVoidTy(mockModule->getContext());
532-
539+
llvm::Type *returnType) {
533540
auto *kleeTaintFunctionType = llvm::FunctionType::get(
534541
returnType,
535542
{llvm::Type::getInt8PtrTy(mockModule->getContext()),
@@ -592,17 +599,13 @@ MockBuilder::buildCallKleeTaintFunction(const std::string &functionName,
592599
llvm::APInt(64, taint, false))});
593600
}
594601

595-
void MockBuilder::buildCallKleeTaintSinkHit(size_t taintSink) {
596-
auto *kleeTaintSinkHitType = llvm::FunctionType::get(
602+
void MockBuilder::buildCallKleeTaintHit(llvm::Value *taintRule) {
603+
auto *kleeTaintHitType = llvm::FunctionType::get(
597604
llvm::Type::getVoidTy(mockModule->getContext()),
598605
{llvm::Type::getInt64Ty(mockModule->getContext())}, false);
599-
auto kleeTaintSinkHitCallee = mockModule->getOrInsertFunction(
600-
"klee_taint_sink_hit", kleeTaintSinkHitType);
601-
602-
builder->CreateCall(
603-
kleeTaintSinkHitCallee,
604-
{llvm::ConstantInt::get(mockModule->getContext(),
605-
llvm::APInt(64, taintSink, false))});
606+
auto kleeTaintSinkHitCallee =
607+
mockModule->getOrInsertFunction("klee_taint_hit", kleeTaintHitType);
608+
builder->CreateCall(kleeTaintSinkHitCallee, {taintRule});
606609
}
607610

608611
void MockBuilder::buildAnnotationForExternalFunctionArgs(

‎lib/Core/SpecialFunctionHandler.cpp‎

Lines changed: 13 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -133,8 +133,8 @@ static SpecialFunctionHandler::HandlerInfo handlerInfo[] = {
133133
add("klee_add_taint", handleAddTaint, false),
134134
add("klee_clear_taint", handleClearTaint, false),
135135
add("klee_check_taint_source", handleCheckTaintSource, true),
136-
add("klee_check_taint_sink", handleCheckTaintSink, true),
137-
add("klee_taint_sink_hit", handleTaintSinkHit, false),
136+
add("klee_get_taint_rule", handleGetTaintRule, true),
137+
add("klee_taint_hit", handleTaintHit, false),
138138

139139
#ifdef SUPPORT_KLEE_EH_CXX
140140
add("_klee_eh_Unwind_RaiseException_impl", handleEhUnwindRaiseExceptionImpl,
@@ -1254,35 +1254,37 @@ void SpecialFunctionHandler::handleCheckTaintSource(
12541254
}
12551255
}
12561256

1257-
void SpecialFunctionHandler::handleCheckTaintSink(
1257+
void SpecialFunctionHandler::handleGetTaintRule(
12581258
klee::ExecutionState &state, klee::KInstruction *target,
12591259
std::vector<ref<Expr>> &arguments) {
12601260
if (arguments.size() != 2) {
12611261
executor.terminateStateOnUserError(state,
12621262
"Incorrect number of arguments to "
1263-
"klee_check_taint_sink(void*, size_t)");
1263+
"klee_get_taint_rule(void*, size_t)");
12641264
return;
12651265
}
12661266

1267-
// ref<Expr> result = ConstantExpr::create(true, Expr::Bool);
1268-
// executor.bindLocal(target, state, result);
1267+
// ref<Expr> result = ConstantExpr::create(4, Expr::Int64);
1268+
// executor.bindLocal(target, state, result);
12691269
}
12701270

1271-
void SpecialFunctionHandler::handleTaintSinkHit(
1271+
void SpecialFunctionHandler::handleTaintHit(
12721272
klee::ExecutionState &state, klee::KInstruction *target,
12731273
std::vector<ref<Expr>> &arguments) {
12741274
if (arguments.size() != 1) {
12751275
executor.terminateStateOnUserError(
1276-
state, "Incorrect number of arguments to klee_taint_sink_hit(size_t)");
1276+
state, "Incorrect number of arguments to klee_taint_hit(size_t)");
12771277
return;
12781278
}
12791279

12801280
char *end = nullptr;
1281-
size_t sink = strtoul(arguments[0]->toString().c_str(), &end, 10);
1281+
size_t rule = strtoul(arguments[0]->toString().c_str(), &end, 10);
12821282
if (*end != '\0' || errno == ERANGE) {
12831283
executor.terminateStateOnUserError(
1284-
state, "Incorrect argument 0 to klee_taint_sink_hit(size_t)");
1284+
state, "Incorrect argument 0 to klee_taint_hit(size_t)");
12851285
}
12861286

1287-
executor.terminateStateOnTaintError(state, sink);
1287+
klee_warning("!!!: %s\n", arguments[0]->toString().c_str());
1288+
1289+
executor.terminateStateOnTargetTaintError(state, rule);
12881290
}

‎lib/Core/SpecialFunctionHandler.h‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -184,8 +184,8 @@ class SpecialFunctionHandler {
184184
HANDLER(handleAddTaint);
185185
HANDLER(handleClearTaint);
186186
HANDLER(handleCheckTaintSource);
187-
HANDLER(handleCheckTaintSink);
188-
HANDLER(handleTaintSinkHit);
187+
HANDLER(handleGetTaintRule);
188+
HANDLER(handleTaintHit);
189189
#undef HANDLER
190190
};
191191
} // namespace klee

0 commit comments

Comments
 (0)