Skip to content

Commit

Permalink
Remove dynamic cast for taint annotations and add mock for TaintSink
Browse files Browse the repository at this point in the history
  • Loading branch information
@mamaria-k
mamaria-k committed Dec 15, 2023
1 parent d99ac85 commit 02eb687
Show file tree
Hide file tree
Showing 3 changed files with 94 additions and 68 deletions.
16 changes: 9 additions & 7 deletions include/klee/Module/Annotation.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -47,6 +47,8 @@ struct Unknown {
std::string rawOffset;
std::string rawValue;

std::string taintType;

public:
std::vector<std::string> offset;

Expand All @@ -58,6 +60,10 @@ struct Unknown {

[[nodiscard]] const std::vector<std::string> &getOffset() const;
[[nodiscard]] std::string toString() const;

[[nodiscard]] virtual bool isTaint() const;
[[nodiscard]] optional<std::string> getTaintType() const;
[[nodiscard]] optional<std::string> getTaintTypeAsLower() const;
};

struct Deref final : public Unknown {
Expand Down Expand Up @@ -114,31 +120,27 @@ struct Free final : public Unknown {
[[nodiscard]] Kind getKind() const override;
};

//TODO: maybe separate class Taint

struct TaintOutput final : public Unknown {
std::string type;

explicit TaintOutput(const std::string &str = "TaintOutput");

[[nodiscard]] Kind getKind() const override;
[[nodiscard]] bool isTaint() const override;
};

struct TaintPropagation final : public Unknown {
std::string type;
size_t propagationParameter;

explicit TaintPropagation(const std::string &str = "TaintPropagation");

[[nodiscard]] Kind getKind() const override;
[[nodiscard]] bool isTaint() const override;
};

struct TaintSink final : public Unknown {
std::string type;

explicit TaintSink(const std::string &str = "TaintSink");

[[nodiscard]] Kind getKind() const override;
[[nodiscard]] bool isTaint() const override;
};

using Ptr = std::shared_ptr<Unknown>;
Expand Down
94 changes: 48 additions & 46 deletions lib/Core/MockBuilder.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -460,59 +460,61 @@ void MockBuilder::buildAnnotationForExternalFunctionArgs(
break;
}
case Statement::Kind::TaintOutput: {
auto x = std::dynamic_pointer_cast<Statement::TaintOutput>(statement);
// buildCallKleeTaintFunction("klee_add_taint", elem,);
break;
}
case Statement::Kind::TaintPropagation: {
// buildCallKleeTaintFunction("klee_add_taint", elem,);
break;
}
// case Statement::Kind::FormatString: {
// if (!elem->getType()->isPointerTy()) {
// klee_error("Annotation: FormatString arg not pointer");
// }
//
// std::string sinkCondName =
// "condition_sink_format_string_arg_" + std::to_string(i) +
// "_" + func->getName().str();
//
// llvm::BasicBlock *fromIf = builder->GetInsertBlock();
// llvm::Function *curFunc = fromIf->getParent();
//
// llvm::BasicBlock *sinkBB = llvm::BasicBlock::Create(
// mockModule->getContext(), sinkCondName, curFunc);
// llvm::BasicBlock *contBB = llvm::BasicBlock::Create(
// mockModule->getContext(), "continue_" + sinkCondName);
// auto brValueSink = buildCallKleeTaintFunction(
// "klee_check_taint_sink", elem,
// static_cast<size_t>(Statement::Kind::FormatString));
// builder->CreateCondBr(brValueSink, sinkBB, contBB);
//
// builder->SetInsertPoint(sinkBB);
// std::string sinkHitCondName =
// "condition_sink_hit_format_string_arg_" + std::to_string(i) +
// "_" + func->getName().str();
// auto intType = llvm::IntegerType::get(mockModule->getContext(), 1);
// auto *sinkHitCond = builder->CreateAlloca(intType, nullptr);
// buildCallKleeMakeSymbolic("klee_make_mock", sinkHitCond, intType,
// sinkHitCondName);
// fromIf = builder->GetInsertBlock();
// curFunc = fromIf->getParent();
// llvm::BasicBlock *sinkHitBB = llvm::BasicBlock::Create(
// mockModule->getContext(), sinkHitCondName, curFunc);
// auto brValueSinkHit = builder->CreateLoad(intType, sinkHitCond);
// builder->CreateCondBr(brValueSinkHit, sinkHitBB, contBB);
//
// builder->SetInsertPoint(sinkHitBB);
// buildCallKleeTaintSinkHit(
// static_cast<size_t>(Statement::Kind::FormatString));
// builder->CreateBr(contBB);
//
// curFunc->getBasicBlockList().push_back(contBB);
// builder->SetInsertPoint(contBB);
// break;
// }
case Statement::Kind::TaintSink: {
if (!elem->getType()->isPointerTy()) {
klee_error("Annotation: TaintSink arg is not pointer");
}

const std::string sinkTypeLower = statement->getTaintTypeAsLower().value();
const size_t sinkTypeNum = annotationsData.taintAnnotation.sinks[
statement->getTaintType().value()
];

const std::string sinkCondName =
"condition_sink_" + sinkTypeLower + "_arg_" +
std::to_string(i) + "_" + func->getName().str();

llvm::BasicBlock *fromIf = builder->GetInsertBlock();
llvm::Function *curFunc = fromIf->getParent();

llvm::BasicBlock *sinkBB = llvm::BasicBlock::Create(
mockModule->getContext(), sinkCondName, curFunc);
llvm::BasicBlock *contBB = llvm::BasicBlock::Create(
mockModule->getContext(), "continue_" + sinkCondName);
auto brValueSink = buildCallKleeTaintFunction(
"klee_check_taint_sink", elem, sinkTypeNum);
builder->CreateCondBr(brValueSink, sinkBB, contBB);

builder->SetInsertPoint(sinkBB);
std::string sinkHitCondName =
"condition_sink_hit_" + sinkTypeLower + "_arg_" +
std::to_string(i) + "_" + func->getName().str();
auto intType = llvm::IntegerType::get(mockModule->getContext(), 1);
auto *sinkHitCond = builder->CreateAlloca(intType, nullptr);
buildCallKleeMakeSymbolic("klee_make_mock", sinkHitCond, intType,
sinkHitCondName);
fromIf = builder->GetInsertBlock();
curFunc = fromIf->getParent();
llvm::BasicBlock *sinkHitBB = llvm::BasicBlock::Create(
mockModule->getContext(), sinkHitCondName, curFunc);
auto brValueSinkHit = builder->CreateLoad(intType, sinkHitCond);
builder->CreateCondBr(brValueSinkHit, sinkHitBB, contBB);

builder->SetInsertPoint(sinkHitBB);
buildCallKleeTaintSinkHit(sinkTypeNum);
builder->CreateBr(contBB);

curFunc->getBasicBlockList().push_back(contBB);
builder->SetInsertPoint(contBB);
break;
}
case Statement::Kind::Unknown:
default:
klee_message("Annotation: not implemented %s",
Expand Down
52 changes: 37 additions & 15 deletions lib/Module/Annotation.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,7 @@ Unknown::Unknown(const std::string &str) {
rawOffset = str.substr(startOffset, offsetLength);
if (secondColonPos != std::string::npos) {
rawValue = str.substr(secondColonPos + 1, std::string::npos);
taintType = rawValue.substr(0, rawValue.find(':'));
}
}

Expand Down Expand Up @@ -80,6 +81,18 @@ std::string Unknown::toString() const {
return rawAnnotation + ":" + rawOffset + ":" + rawValue;
}

bool Unknown::isTaint() const { return false; }

optional<std::string> Unknown::getTaintType() const {
if (!isTaint()) return nullopt;
return taintType;
}

optional<std::string> Unknown::getTaintTypeAsLower() const {
if (!isTaint()) return nullopt;
return toLower(taintType);
}

bool Unknown::operator==(const Unknown &other) const {
return this->getKind() == other.getKind() && toString() == other.toString();
}
Expand Down Expand Up @@ -128,56 +141,65 @@ Kind Free::getKind() const { return Kind::Free; }
*/

TaintOutput::TaintOutput(const std::string &str) : Unknown(str) {
if (rawValue.empty()) {
klee_error("Annotation TaintOutput: Incorrect value format, must be not empty");
// TODO: in the future, support typeless annotations (meaning all types)
if (taintType.empty()) {
klee_error("Annotation Taint: Incorrect value format, must has taint type");
}
type = rawValue;
}

Kind TaintOutput::getKind() const { return Kind::TaintOutput; }

bool TaintOutput::isTaint() const { return true; }

/*
* Format: TaintPropagation::{type}:{data}
*/

TaintPropagation::TaintPropagation(const std::string &str) : Unknown(str) {
// TODO: in the future, support typeless annotations (meaning all types)
if (taintType.empty()) {
klee_error("Annotation Taint: Incorrect value format, must has taint type");
}
if (!rawOffset.empty()) {
klee_error("Annotation TaintPropagation: Incorrect offset format, must be empty");
}
if (rawValue.empty()) {
klee_error("Annotation TaintPropagation: Incorrect value format, must be not empty");
}

const size_t colonPos = rawValue.find(':');
if (colonPos == std::string::npos) {
const std::string rawData = (colonPos == std::string::npos)
? std::string()
: rawValue.substr(colonPos + 1, std::string::npos);

if (rawData.empty()) {
klee_error("Annotation TaintPropagation: Incorrect value %s format, must be <type>:<index>", rawValue.c_str());
}
type = rawValue.substr(0, colonPos);

if (sscanf(rawValue.substr(colonPos + 1, std::string::npos).c_str(), "%zu", &propagationParameter) != 1) {
if (sscanf(rawData.c_str(), "%zu", &propagationParameter) != 1) {
klee_error("Annotation TaintPropagation: Incorrect value %s format, must be <type>:<index>", rawValue.c_str());
}
}

Kind TaintPropagation::getKind() const { return Kind::TaintPropagation; }

bool TaintPropagation::isTaint() const { return true; }

/*
* Format: TaintSink::{type}
*/

TaintSink::TaintSink(const std::string &str) : Unknown(str) {
// TODO: in the future, support typeless annotations (meaning all types)
if (taintType.empty()) {
klee_error("Annotation Taint: Incorrect value format, must has taint type");
}
//TODO: now this is true for cooddy, but it’s not clear how to use them.
if (!rawOffset.empty()) {
klee_error("Annotation TaintSink: Incorrect offset format, must be empty");
}
if (rawValue.empty()) {
klee_error("Annotation TaintSink: Incorrect value format, must be not empty");
}

type = rawValue;
}

Kind TaintSink::getKind() const { return Kind::TaintSink; }

bool TaintSink::isTaint() const { return true; }

const std::map<std::string, Statement::Kind> StringToKindMap = {
{"deref", Statement::Kind::Deref},
{"initnull", Statement::Kind::InitNull},
Expand Down

0 comments on commit 02eb687

Please sign in to comment.