BasicBlock* LinearPass::CreateBlockWorker(std::vector<BasicBlock>* Blocks, duint Start, duint End, bool Call, bool Jmp, bool Ret, bool Pad)
{
BasicBlock block { Start, End - 1, 0, 0, 0 };
// Check for calls
if(Call)
block.SetFlag(BASIC_BLOCK_FLAG_CALL);
// Check for returns
if(Ret)
block.SetFlag(BASIC_BLOCK_FLAG_RET);
// Check for interrupts
if(Pad)
block.SetFlag(BASIC_BLOCK_FLAG_PAD);
Blocks->push_back(block);
// std::vector::back() incurs a very large performance overhead (30% slower) when
// used in debug mode. This code eliminates it from showing up in the profiler.
#ifdef _DEBUG
return &Blocks->data()[Blocks->size() - 1];
#else
return &Blocks->back();
#endif // _DEBUG
}
bool FunctionPass::ResolveKnownFunctionEnd(FunctionDef* Function)
{
// Helper to link final blocks to function
auto startBlock = FindBBlockInRange(Function->VirtualStart);
auto endBlock = FindBBlockInRange(Function->VirtualEnd);
if(!startBlock || !endBlock)
return false;
// Find block start/end indices
Function->BBlockStart = FindBBlockIndex(startBlock);
Function->BBlockEnd = FindBBlockIndex(endBlock);
// Set the flag for blocks that have been scanned
for(BasicBlock* block = startBlock; (duint)block <= (duint)endBlock; block++)
{
// Block now in use
block->SetFlag(BASIC_BLOCK_FLAG_FUNCTION);
// Counter
Function->InstrCount += block->InstrCount;
}
return true;
}
bool FunctionPass::ResolveFunctionEnd(FunctionDef* Function, BasicBlock* LastBlock)
{
ASSERT_TRUE(Function->VirtualStart != 0);
// Find the first basic block of the function
BasicBlock* block = FindBBlockInRange(Function->VirtualStart);
if(!block)
{
ASSERT_ALWAYS("Block should exist at this point");
return false;
}
// The maximum address is determined by any jump that extends past
// a RET or other terminating basic block. A function may have multiple
// return statements.
duint maximumAddr = 0;
// Loop forever until the end is found
for(; (duint)block <= (duint)LastBlock; block++)
{
if(block->GetFlag(BASIC_BLOCK_FLAG_CALL_TARGET) && block->VirtualStart != Function->VirtualStart)
{
block--;
break;
}
// Block is now in use
block->SetFlag(BASIC_BLOCK_FLAG_FUNCTION);
// Increment instruction count
Function->InstrCount += block->InstrCount;
// Calculate max from just linear instructions
maximumAddr = max(maximumAddr, block->VirtualEnd);
// Find maximum jump target
if(!block->GetFlag(BASIC_BLOCK_FLAG_CALL) && !block->GetFlag(BASIC_BLOCK_FLAG_INDIRECT))
{
if(block->Target != 0 && block->Target >= maximumAddr)
{
// Here's a problem: Compilers add tail-call elimination with a jump.
// Solve this by creating a maximum jump limit.
auto targetBlock = FindBBlockInRange(block->Target);
// If (target block found) and (target block is not called)
if(targetBlock && !targetBlock->GetFlag(BASIC_BLOCK_FLAG_CALL_TARGET))
{
duint blockEnd = targetBlock->VirtualEnd;
//
// Edge case when a compiler emits:
//
// pop ebp
// jmp some_func
// int3
// int3
// some_func:
// push ebp
//
// Where INT3 will align "some_func" to 4, 8, 12, or 16.
// INT3 padding is also optional (if the jump fits perfectly).
//
if(true/*block->GetFlag(BASIC_BLOCK_FLAG_ABSJMP)*/)
{
{
// Check if padding is aligned to 4
auto nextBlock = block + 1;
if((duint)nextBlock <= (duint)LastBlock)
{
if(nextBlock->GetFlag(BASIC_BLOCK_FLAG_PAD))
{
// If this block is aligned to 4 bytes at the end
if((nextBlock->VirtualEnd + 1) % 4 == 0)
blockEnd = block->VirtualEnd;
}
}
}
}
// Now calculate the maximum end address, taking into account the jump destination
maximumAddr = max(maximumAddr, blockEnd);
}
}
}
// Sanity check
ASSERT_TRUE(maximumAddr >= block->VirtualStart);
// Does this node contain the maximum address?
if(maximumAddr >= block->VirtualStart && maximumAddr <= block->VirtualEnd)
{
// It does! There's 4 possibilities next:
//
// 1. Return
// 2. Tail-call elimination
// 3. Optimized loop
// 4. Function continues to next block
//
// 1.
if(block->GetFlag(BASIC_BLOCK_FLAG_RET))
break;
if(block->Target != 0)
{
// NOTE: Both must be an absolute jump
if(block->GetFlag(BASIC_BLOCK_FLAG_ABSJMP))
{
// 2.
if(block->VirtualEnd == maximumAddr)
break;
// 3.
if(block->Target >= Function->VirtualStart && block->Target < block->VirtualEnd)
break;
}
}
// 4. Continue
}
}
// Loop is done. Set the information in the function structure.
Function->VirtualEnd = block->VirtualEnd;
Function->BBlockEnd = FindBBlockIndex(block);
return true;
}
void LinearPass::AnalysisOverlapWorker(duint Start, duint End, BBlockArray* Insertions)
{
// Comparison function to see if two blocks overlap
auto BlockOverlapsRemove = [](BasicBlock * A, BasicBlock * B) -> BasicBlock*
{
// Do the blocks overlap?
if(max(A->VirtualStart, B->VirtualStart) <= min((A->VirtualEnd - 1), (B->VirtualEnd - 1)))
{
// Return the block that should be removed
if(A->Size() > B->Size())
return B;
return A;
}
return nullptr;
};
// Get a pointer to pure data
const auto blocks = m_MainBlocks.data();
for(duint i = Start; i < End; i++)
{
const auto curr = &blocks[i];
const auto next = &blocks[i + 1];
// Current versus next (overlap -> delete)
BasicBlock* removal = BlockOverlapsRemove(curr, next);
if(removal)
removal->SetFlag(BASIC_BLOCK_FLAG_DELETE);
// Find blocks that need to be split in two because
// of CALL/JMP targets
//
// Find targets in the virtual range
if(ValidateAddress(curr->Target))
{
removal = FindBBlockInRange(curr->Target);
if(removal)
{
if(curr->GetFlag(BASIC_BLOCK_FLAG_CALL))
removal->SetFlag(BASIC_BLOCK_FLAG_CALL_TARGET);
// If the target does not equal the block start...
if(curr->Target != removal->VirtualStart)
{
// Mark for deletion
removal->SetFlag(BASIC_BLOCK_FLAG_DELETE);
// Block part 1
BasicBlock block1;
block1.VirtualStart = removal->VirtualStart;
block1.VirtualEnd = curr->Target;
block1.Target = 0;
block1.Flags = BASIC_BLOCK_FLAG_CUTOFF; // Attributes of the top half
block1.InstrCount = removal->InstrCount;
// Block part 2
BasicBlock block2;
block2.VirtualStart = curr->Target;
block2.VirtualEnd = removal->VirtualEnd;
block2.Target = removal->Target;
block2.Flags = removal->Flags; // Attributes of the bottom half (taken from original block)
block2.InstrCount = removal->InstrCount;
Insertions->push_back(block1);
Insertions->push_back(block2);
}
}
}
}
}
void LinearPass::AnalysisWorker(duint Start, duint End, BBlockArray* Blocks)
{
Capstone disasm;
duint blockBegin = Start; // BBlock starting virtual address
duint blockEnd = 0; // BBlock ending virtual address
bool blockPrevPad = false; // Indicator if the last instruction was padding
BasicBlock* lastBlock = nullptr; // Avoid an expensive call to std::vector::back()
int insnCount = 0; // Temporary number of instructions counted for a block
for(duint i = Start; i < End;)
{
if(!disasm.Disassemble(i, TranslateAddress(i), int(End - i)))
{
// Skip instructions that can't be determined
i++;
continue;
}
// Increment counters
i += disasm.Size();
blockEnd = i;
insnCount++;
// The basic block ends here if it is a branch
bool call = disasm.InGroup(CS_GRP_CALL); // CALL
bool jmp = disasm.InGroup(CS_GRP_JUMP); // JUMP
bool ret = disasm.InGroup(CS_GRP_RET); // RETURN
bool padding = disasm.IsFilling(); // INSTRUCTION PADDING
if(padding)
{
// PADDING is treated differently. They are all created as their
// own separate block for more analysis later.
duint realBlockEnd = blockEnd - disasm.Size();
if((realBlockEnd - blockBegin) > 0)
{
// The next line terminates the BBlock before the INT instruction.
// Early termination, faked as an indirect JMP. Rare case.
lastBlock = CreateBlockWorker(Blocks, blockBegin, realBlockEnd, false, false, false, false);
lastBlock->SetFlag(BASIC_BLOCK_FLAG_PREPAD);
blockBegin = realBlockEnd;
lastBlock->InstrCount = insnCount;
insnCount = 0;
}
}
if(call || jmp || ret || padding)
{
// Was this a padding instruction?
if(padding && blockPrevPad)
{
// Append it to the previous block
lastBlock->VirtualEnd = blockEnd;
}
else
{
// Otherwise use the default route: create a new entry
auto block = lastBlock = CreateBlockWorker(Blocks, blockBegin, blockEnd, call, jmp, ret, padding);
// Counters
lastBlock->InstrCount = insnCount;
insnCount = 0;
if(!padding)
{
// Check if absolute jump, regardless of operand
if(disasm.GetId() == X86_INS_JMP)
block->SetFlag(BASIC_BLOCK_FLAG_ABSJMP);
// Figure out the operand type(s)
const auto & operand = disasm.x86().operands[0];
if(operand.type == X86_OP_IMM)
{
// Branch target immediate
block->Target = (duint)operand.imm;
}
else
{
// Indirects (no operand, register, or memory)
block->SetFlag(BASIC_BLOCK_FLAG_INDIRECT);
if(operand.type == X86_OP_MEM &&
operand.mem.base == X86_REG_RIP &&
operand.mem.index == X86_REG_INVALID &&
operand.mem.scale == 1)
{
/*
block->SetFlag(BASIC_BLOCK_FLAG_INDIRPTR);
block->Target = (duint)operand.mem.disp;
*/
}
}
}
}
// Reset the loop variables
blockBegin = i;
blockPrevPad = padding;
}
}
}
