Ejemplo n.º 1
0
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;
}