Skip to content

Commit

Permalink
chore: refactor
Browse files Browse the repository at this point in the history
  • Loading branch information
hwk2077 committed Sep 11, 2023
1 parent ba40e4d commit 00143d7
Show file tree
Hide file tree
Showing 2 changed files with 111 additions and 140 deletions.
207 changes: 103 additions & 104 deletions libs2eplugins/src/s2e/Plugins/Searchers/ValidPathSearcher.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -19,17 +19,8 @@ S2E_DEFINE_PLUGIN(ValidPathSearcher, "ValidPathSearcher S2E plugin", "ValidPathS

void ValidPathSearcher::initialize() {
s2e()->getExecutor()->setSearcher(this);
m_onStateForkConn =
s2e()->getCorePlugin()->onStateFork.connect(sigc::mem_fun(*this, &ValidPathSearcher::onStateFork));
m_onTranslateBlockEndConn = s2e()->getCorePlugin()->onTranslateBlockEnd.connect(
sigc::mem_fun(*this, &ValidPathSearcher::onTranslateBlockEnd));

onARMFunctionConn = s2e()->getPlugin<ARMFunctionMonitor>();
onARMFunctionConn->onARMFunctionCallEvent.connect(sigc::mem_fun(*this, &ValidPathSearcher::onARMFunctionCall));
onARMFunctionConn->onARMFunctionReturnEvent.connect(sigc::mem_fun(*this, &ValidPathSearcher::onARMFunctionReturn));

initConnection();
m_curState = nullptr;
m_lstState = nullptr;
m_selfSwitch = false;
m_searcherActive = false;
}
Expand All @@ -39,61 +30,14 @@ klee::ExecutionState &ValidPathSearcher::selectState() {
<< "[current stateID: " << m_curState->getID() << "] "
<< "\n";

if (m_selfSwitch) {
m_selfSwitch = false;
uint32_t callerPC = m_functionCalRetMap[m_returnPC];
uint32_t forkPC = m_idStateMap[m_curState->getID()].forkPC;
getDebugStream() << "self switch "
<< "[callerPC: " << hexval(callerPC) << "] "
<< "[forkPC: " << hexval(forkPC) << "] "
<< "\n";
for (auto &it : m_forkStateItems[callerPC][forkPC]) {
if (it->getID() != m_curState->getID() && m_idStateMap[it->getID()].flag != INVALID) {
m_idStateMap[it->getID()].flag = VALID;
m_lstState = m_curState;
m_curState = m_idStateMap[it->getID()].state;
if (m_curState == nullptr) {
getDebugStream() << "m_curState is nullptr"
<< "\n";
}
getDebugStream() << "[new self-switched StateID: " << m_curState->getID() << "] "
<< "\n";
return *m_curState;
}
}
getDebugStream() << "no option for self switch"
<< "\n";
if (m_selfSwitch && selectSelfSwitchedState()) {
return *m_curState;
}

for (auto it = m_idStateMap.rbegin(); it != m_idStateMap.rend(); ++it) {
if (it->second.state->getID() != m_curState->getID() && it->second.state != m_lstState &&
it->second.flag == UNVISITED) {
it->second.flag = VALID;
m_curState = it->second.state;
if (m_curState == nullptr) {
getDebugStream() << "m_curState is nullptr"
<< "\n";
}
getDebugStream() << "[new selected unvisited stateID: " << m_curState->getID() << "] "
<< "\n";
m_lstState = m_curState;
return *m_curState;
}
if (selectUnvisitedState()) {
return *m_curState;
}
for (auto it = m_idStateMap.rbegin(); it != m_idStateMap.rend(); ++it) {
if (it->second.state->getID() != m_curState->getID() && it->second.state != m_lstState &&
it->second.flag == VALID) {
m_curState = it->second.state;
if (m_curState == nullptr) {
getDebugStream() << "m_curState is nullptr"
<< "\n";
}
m_lstState = m_curState;
getDebugStream() << "[new selected valid stateID: " << m_curState->getID() << "] "
<< "\n";
return *m_curState;
}
if (selectValidState()) {
return *m_curState;
}

getDebugStream() << "no option"
Expand All @@ -103,14 +47,6 @@ klee::ExecutionState &ValidPathSearcher::selectState() {

void ValidPathSearcher::update(klee::ExecutionState *current, const klee::StateSet &addedStates,
const klee::StateSet &removedStates) {
// if (m_curState) {
// getDebugStream() << "update "
// << "[current stateID: " << m_curState->getID() << "] "
// << "\n";
// } else {
// getDebugStream() << "update "
// << "\n";
// }
for (auto it : removedStates) {
S2EExecutionState *removedState = dynamic_cast<S2EExecutionState *>(it);
auto found = m_idStateMap.find(removedState->getID());
Expand All @@ -120,11 +56,6 @@ void ValidPathSearcher::update(klee::ExecutionState *current, const klee::StateS
<< "\n";
}
}
// for (auto it : m_idStateMap) {
// getDebugStream() << "[traversal stateID: " << it.second.state->getID() << "] "
// << "[flag: " << it.second.flag << "] "
// << "\n";
// }
}

bool ValidPathSearcher::empty() {
Expand All @@ -138,21 +69,6 @@ bool ValidPathSearcher::empty() {
return false;
}

bool ValidPathSearcher::insertToForkStates(uint32_t callerPC, uint32_t forkPC, S2EExecutionState *state) {
const auto &foundSameCallerPC = m_forkStateItems.find(callerPC);
if (foundSameCallerPC != m_forkStateItems.end()) {
const auto &foundSameForkPC = foundSameCallerPC->second.find(forkPC);
if (foundSameForkPC != foundSameCallerPC->second.end()) {
if (foundSameForkPC->second.size() >= 2) {
return false;
}
}
}
m_forkStateItems[callerPC][forkPC].push_back(state);
m_idStateMap[state->getID()] = {state, forkPC, UNVISITED};
return true;
}

void ValidPathSearcher::onStateFork(S2EExecutionState *state, const std::vector<S2EExecutionState *> &newStates,
const std::vector<klee::ref<klee::Expr>> &newConditions) {
getDebugStream() << "onStateFork" << '\n';
Expand Down Expand Up @@ -202,21 +118,104 @@ void ValidPathSearcher::onARMFunctionReturn(S2EExecutionState *state, uint32_t p
return;
}

// m_selfSwitch = true;
// // s2e()->getExecutor()->selectNextState(m_curState);
// s2e()->getExecutor()->suspendState(m_curState);

// uint32_t callerPC = m_functionCalRetMap[m_returnPC];
// uint32_t forkPC = m_idStateMap[state->getID()].forkPC;
// for (auto &it : m_forkStateItems[callerPC][forkPC]) {
// if (m_idStateMap[it->getID()].state->getID() != m_curState->getID() &&
// m_idStateMap[it->getID()].flag != INVALID) {
// m_selfSwitch = true;
// m_curState = m_idStateMap[it->getID()].state;
// s2e()->getExecutor()->selectNextState(m_curState);
// }
// }
m_selfSwitch = true;
// s2e()->getExecutor()->selectNextState(m_curState);
s2e()->getExecutor()->suspendState(m_curState);
// s2e()->getExecutor()->setCpuExitRequest();
}

void ValidPathSearcher::initConnection() {
m_onStateForkConn =
s2e()->getCorePlugin()->onStateFork.connect(sigc::mem_fun(*this, &ValidPathSearcher::onStateFork));
m_onTranslateBlockEndConn = s2e()->getCorePlugin()->onTranslateBlockEnd.connect(
sigc::mem_fun(*this, &ValidPathSearcher::onTranslateBlockEnd));

onARMFunctionConn = s2e()->getPlugin<ARMFunctionMonitor>();
onARMFunctionConn->onARMFunctionCallEvent.connect(sigc::mem_fun(*this, &ValidPathSearcher::onARMFunctionCall));
onARMFunctionConn->onARMFunctionReturnEvent.connect(sigc::mem_fun(*this, &ValidPathSearcher::onARMFunctionReturn));
}

bool ValidPathSearcher::insertToForkStates(uint32_t callerPC, uint32_t forkPC, S2EExecutionState *state) {
const auto &foundSameCallerPC = m_forkStateItems.find(callerPC);
if (foundSameCallerPC != m_forkStateItems.end()) {
const auto &foundSameForkPC = foundSameCallerPC->second.find(forkPC);
if (foundSameForkPC != foundSameCallerPC->second.end()) {
if (foundSameForkPC->second.size() >= 2) {
return false;
}
}
}
m_forkStateItems[callerPC][forkPC].push_back(state);
m_idStateMap[state->getID()] = {state, forkPC, UNVISITED, 0};
return true;
}

bool ValidPathSearcher::selectSelfSwitchedState() {
m_selfSwitch = false;
uint32_t callerPC = m_functionCalRetMap[m_returnPC];
uint32_t forkPC = m_idStateMap[m_curState->getID()].forkPC;
getDebugStream() << "self switch "
<< "[callerPC: " << hexval(callerPC) << "] "
<< "[forkPC: " << hexval(forkPC) << "] "
<< "\n";

for (auto &it : m_forkStateItems[callerPC][forkPC]) {
const int id = it->getID();
if (m_curState->getID() != id && m_idStateMap[id].flag != INVALID) {
if (m_idStateMap[id].switchCnt >= MAX_SWITCH_CNT) {
m_idStateMap[id].flag = INVALID;
continue;
}
m_idStateMap[id].switchCnt++;
m_idStateMap[id].flag = VALID;
m_curState = m_idStateMap[id].state;
getDebugStream() << "[new self-switched StateID: " << m_curState->getID() << "] "
<< "\n";
return true;
}
}

getDebugStream() << "no option for self switch"
<< "\n";
return false;
}

bool ValidPathSearcher::selectUnvisitedState() {
for (auto it = m_idStateMap.rbegin(); it != m_idStateMap.rend(); ++it) {
if (it->second.state->getID() != m_curState->getID() && it->second.flag == UNVISITED) {
if (it->second.switchCnt >= MAX_SWITCH_CNT) {
it->second.flag = INVALID;
continue;
}
it->second.switchCnt++;
it->second.flag = VALID;
m_curState = it->second.state;
getDebugStream() << "[new selected unvisited stateID: " << m_curState->getID() << "] "
<< "\n";
return true;
}
}
return false;
}

bool ValidPathSearcher::selectValidState() {
for (auto it = m_idStateMap.rbegin(); it != m_idStateMap.rend(); ++it) {
if (it->second.state->getID() != m_curState->getID() && it->second.flag == VALID) {
if (it->second.switchCnt >= MAX_SWITCH_CNT) {
it->second.flag = INVALID;
continue;
}
it->second.switchCnt++;
m_curState = it->second.state;
getDebugStream() << "[new selected valid stateID: " << m_curState->getID() << "] "
<< "\n";
return true;
}
}
return false;
}



} // namespace plugins
} // namespace s2e
44 changes: 8 additions & 36 deletions libs2eplugins/src/s2e/Plugins/Searchers/ValidPathSearcher.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -35,9 +35,11 @@ struct StateItem {
S2EExecutionState *state;
uint32_t forkPC;
uint32_t flag;
uint32_t switchCnt;
};

enum StatePhase { UNVISITED, INVALID, VALID };
const int MAX_SWITCH_CNT = 20;

class ValidPathSearcher : public Plugin, public klee::Searcher {
S2E_PLUGIN
Expand All @@ -57,52 +59,25 @@ class ValidPathSearcher : public Plugin, public klee::Searcher {
uint32_t m_callerPC;
uint32_t m_returnPC;
S2EExecutionState *m_curState;
S2EExecutionState *m_lstState;
bool m_selfSwitch;
bool m_searcherActive;

std::map<uint32_t, std::map<uint32_t, std::vector<S2EExecutionState *>>> m_forkStateItems;
std::map<uint32_t, StateItem> m_idStateMap;
std::map<uint32_t, uint32_t> m_functionCalRetMap;

void initConnection();
bool insertToForkStates(uint32_t callerPC, uint32_t forkPC, S2EExecutionState *state);
bool selectSelfSwitchedState();
bool selectUnvisitedState();
bool selectValidState();

public:
/**
* callback when state fork
* @param state
* @param newStates
* @param newConditions
*/
void onStateFork(S2EExecutionState *state, const std::vector<S2EExecutionState *> &newStates,
const std::vector<klee::ref<klee::Expr>> &newConditions);

/**
*
* @param signal
* @param state
* @param tb
* @param pc
* @param isStatic
* @param staticTargetPc
*/
void onTranslateBlockEnd(ExecutionSignal *signal, S2EExecutionState *state, TranslationBlock *tb, uint64_t pc,
bool isStatic, uint64_t staticTargetPc);

/**
* callback when function call.
*
* @param state s2e state
* @param pcCaller pc of caller
* @param pcCtxHashVal pc and context hash value
* @param pcReturn pc of return
*/
void onARMFunctionCall(S2EExecutionState *state, uint32_t pcCaller, uint64_t pcCtxHashVal, uint32_t pcReturn);

/**
* callback when function return.
*
* @param state s2e state
* @param pcReturn pc of return
*/
void onARMFunctionReturn(S2EExecutionState *state, uint32_t pcReturn);

/**
Expand All @@ -128,9 +103,6 @@ class ValidPathSearcher : public Plugin, public klee::Searcher {
* @return true, if
*/
virtual bool empty();

private:
bool insertToForkStates(uint32_t callerPC, uint32_t forkPC, S2EExecutionState *state);
};

} // namespace plugins
Expand Down

0 comments on commit 00143d7

Please sign in to comment.