diff --git a/src/AssemblyAnalyzer.cpp b/src/AssemblyAnalyzer.cpp index fdef7c2..859a3fd 100644 --- a/src/AssemblyAnalyzer.cpp +++ b/src/AssemblyAnalyzer.cpp @@ -51,9 +51,20 @@ bool IsVirtualTableJumpThunkInstruction(const ZydisDecodedInstruction& Instructi FirstOp.mem.index == ZydisRegister::ZYDIS_REGISTER_NONE; } -//Performs code discovery for the function located at the given offset, determining it's kind and real address +//Performs code discovery for the function located at the given offset, determining its kind and real address //by going through thunks if it's necessary and returns info about the final function -FunctionInfo DiscoverFunction(uint8_t* FunctionPtr) { +FunctionInfo DiscoverFunction(unsigned char* FunctionPtr) { + for (;;) { + const FunctionInfoStep Step = DiscoverFunctionStep(FunctionPtr); + if (Step.bIsDone) { + return Step.FinalFunctionInfo; + } + FunctionPtr = Step.IntermediateAddress; + } +} + +//Iterative step +FunctionInfoStep DiscoverFunctionStep(unsigned char* FunctionPtr) { ZydisDecoder decoder; ZydisDecoderInit(&decoder, ZYDIS_MACHINE_MODE_LONG_64, ZYDIS_ADDRESS_WIDTH_64); @@ -80,7 +91,7 @@ FunctionInfo DiscoverFunction(uint8_t* FunctionPtr) { ZyanU64 ResultJumpAddress; const bool bSuccessCalculation = ZYAN_SUCCESS(ZydisCalcAbsoluteAddress(&Instruction, &Instruction.operands[0], (ZyanU64) FunctionPtr, &ResultJumpAddress)); assert(bSuccessCalculation); - return DiscoverFunction((uint8_t*) ResultJumpAddress); + return (unsigned char*)ResultJumpAddress; } //test for indirect jump thunks, usually encountered in import table functions wrappers @@ -88,8 +99,8 @@ FunctionInfo DiscoverFunction(uint8_t* FunctionPtr) { ZyanU64 ResultMemoryLocation; const bool bSuccessCalculation = ZYAN_SUCCESS(ZydisCalcAbsoluteAddress(&Instruction, &Instruction.operands[0], (ZyanU64) FunctionPtr, &ResultMemoryLocation)); assert(bSuccessCalculation); - uint8_t* ResultJumpAddress = *(uint8_t**) ResultMemoryLocation; - return DiscoverFunction(ResultJumpAddress); + unsigned char* ResultJumpAddress = *(unsigned char**)ResultMemoryLocation; + return ResultJumpAddress; } //test for virtual table call thunk @@ -113,4 +124,4 @@ FunctionInfo DiscoverFunction(uint8_t* FunctionPtr) { //We can assume this is correct function pointer now return FunctionInfo{true, false, FunctionPtr}; -} \ No newline at end of file +} diff --git a/src/AssemblyAnalyzer.h b/src/AssemblyAnalyzer.h index 58b834a..3d87c6e 100644 --- a/src/AssemblyAnalyzer.h +++ b/src/AssemblyAnalyzer.h @@ -20,6 +20,17 @@ struct FunctionInfo { unsigned int VirtualTableFunctionOffset; }; +//Union that either contains the fully-discovered function info, or the next address to analyze. +struct FunctionInfoStep { + bool bIsDone; + union { + FunctionInfo FinalFunctionInfo; // if bIsDone + unsigned char* IntermediateAddress; // if !bIsDone + }; + FunctionInfoStep(FunctionInfo FinalFunctionInfo) : bIsDone(true), FinalFunctionInfo(FinalFunctionInfo) {} + FunctionInfoStep(unsigned char* IntermediateAddress) : bIsDone(false), IntermediateAddress(IntermediateAddress) {} +}; + typedef void(*LogDebugMessageFunctionType)(const char*); /** Sets up function used for debugging logging, if it's enabled by build configuration */ @@ -33,4 +44,10 @@ AA_API void SetDebugLoggingHook(LogDebugMessageFunctionType LogDebugMessage); * @param FunctionPtr pointer to the function's code to analyze * @return information about the function */ -AA_API FunctionInfo DiscoverFunction(unsigned char* FunctionPtr); \ No newline at end of file +AA_API FunctionInfo DiscoverFunction(unsigned char* FunctionPtr); + +/** + * Performs a single step of function analysis, stopping after following a single address. + * This can be used to find the next thunk without following the entire chain. + */ +AA_API FunctionInfoStep DiscoverFunctionStep(unsigned char* FunctionPtr);