Add an API to perform one step of DiscoverFunction

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,16 +91,16 @@ 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
     if (IsIndirectJumpThunkInstruction(Instruction)) {
         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};
-}
+}

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);
+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);