From d3a863be7dcf898c8119b9d438b1bc2cf220cd59 Mon Sep 17 00:00:00 2001
From: Alexandre Tullot <alexandre.tullot@protonmail.com>
Date: Fri, 16 Jun 2023 16:27:15 +0200
Subject: [PATCH] Add BuildIndex endpoint to generate the index file (#14)

* Add BuildIndex endpoint to ttd_types

* Generate idx file during engine init

* Remove useless parameters
---
 libyarattd/include/libyarattd_ttd.h       |  2 ++
 libyarattd/include/libyarattd_ttd_types.h | 10 +++++-----
 libyarattd/src/libyarattd_ttd.c           | 19 +++++++++++++++----
 libyarattd/src/libyarattd_virtual_alloc.c | 15 ---------------
 4 files changed, 22 insertions(+), 24 deletions(-)

diff --git a/libyarattd/include/libyarattd_ttd.h b/libyarattd/include/libyarattd_ttd.h
index a020869..f188ec5 100644
--- a/libyarattd/include/libyarattd_ttd.h
+++ b/libyarattd/include/libyarattd_ttd.h
@@ -18,6 +18,8 @@ char* base64_encode(
     size_t input_length,
     size_t* output_length);
 
+void build_index_from_engine(TTD_Replay_ReplayEngine* engine);
+
 static unsigned char VERSION_GUID[] = {
     0xA5, 0x20, 0x34, 0x4D, 0xEF, 0x37, 0x14, 0x41, 0xAE, 0x91, 0x63, 0xD0,
     0x37, 0x8C, 0x84, 0xA9, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00};
diff --git a/libyarattd/include/libyarattd_ttd_types.h b/libyarattd/include/libyarattd_ttd_types.h
index 0d9f0b9..3b3fe04 100644
--- a/libyarattd/include/libyarattd_ttd_types.h
+++ b/libyarattd/include/libyarattd_ttd_types.h
@@ -97,6 +97,9 @@ typedef struct TTD_Replay_ExceptionEvent
   TTD_Replay_Exception* info;
 } TTD_Replay_ExceptionEvent;
 
+typedef enum TTD_Replay_IndexStatus TTD_Replay_IndexStatus;
+typedef enum TTD_Replay_IndexBuildFlags TTD_Replay_IndexBuildFlags;
+
 typedef struct TTD_Replay_IReplayEngine_vftable
 {
   // const void* (__fastcall*
@@ -227,11 +230,8 @@ typedef struct TTD_Replay_IReplayEngine_vftable
   struct TTD_Replay_ICursor*(__fastcall* NewCursor)(
       TTD_Replay_ReplayEngine* self,
       const unsigned char* guid);
-  //	enum TTD::Replay::IndexStatus(__stdcall __high*
-  //_BuildIndex_ReplayEngine_Replay_TTD__UEAA_AW4IndexStatus_23_P6AXPEBXPEBUIndexBuildProgressType_23__Z0W4IndexBuildFlags_23__Z)(void(__stdcall
-  //__high*)(const void*, const struct TTD::Replay::IndexBuildProgressType*),
-  // const void*, enum TTD::Replay::IndexBuildFlags);
-  void* unk44;
+  TTD_Replay_IndexStatus(
+      __stdcall* BuildIndex)(TTD_Replay_ReplayEngine* self, void* callback);
   //	enum TTD::Replay::IndexStatus(__high*
   //_GetIndexStatus_ReplayEngine_Replay_TTD__UEBA_AW4IndexStatus_23_XZ)(void);
   void* unk45;
diff --git a/libyarattd/src/libyarattd_ttd.c b/libyarattd/src/libyarattd_ttd.c
index 018ead7..e3f273b 100644
--- a/libyarattd/src/libyarattd_ttd.c
+++ b/libyarattd/src/libyarattd_ttd.c
@@ -57,13 +57,18 @@ int init_ttd_engine(TTD_Replay_ReplayEngine** engine, wchar_t* filename)
   result = CreateReplayEngineWithHandshake(tmp, &instance, VERSION_GUID);
   *engine = (TTD_Replay_ReplayEngine*) instance;
 
-  if ((*engine)->IReplayEngine->Initialize((*engine), filename) == TRUE)
+  if ((*engine)->IReplayEngine->Initialize((*engine), filename) != TRUE)
   {
-    return ERROR_SUCCESS;
+    fwprintf(stdout, L"Failed to initialize ReplayEngine\n");
+    return ERROR_INTERNAL_FATAL_ERROR;
   }
-  else
+
+  // Generate if needed the idx file of the trace file. This file is needed by
+  // TTDReplay.dll to call some API endpoints like GetCrossPlatformContext
+  build_index_from_engine(*engine);
+  if (check_idx_file(filename) != ERROR_SUCCESS)
   {
-    fwprintf(stdout, L"Failed to initialize ReplayEngine\n");
+    fwprintf(stderr, L"Failed to generate index file\n");
     return ERROR_INTERNAL_FATAL_ERROR;
   }
 }
@@ -99,3 +104,9 @@ char* base64_encode(
 
   return encoded_data;
 }
+
+void dummy_callback() {}
+void build_index_from_engine(TTD_Replay_ReplayEngine* engine)
+{
+  engine->IReplayEngine->BuildIndex(engine, &dummy_callback);
+}
diff --git a/libyarattd/src/libyarattd_virtual_alloc.c b/libyarattd/src/libyarattd_virtual_alloc.c
index 556472d..e3bd41b 100644
--- a/libyarattd/src/libyarattd_virtual_alloc.c
+++ b/libyarattd/src/libyarattd_virtual_alloc.c
@@ -191,21 +191,6 @@ int build_virtual_alloc_map_from_cache(
 
 int build_virtual_alloc_map(YR_TTD_SCHEDULER* scheduler)
 {
-  // Check that the idx file exists in the same path as the trace file
-  // This file is needed by TTDReplay.dll to use
-  // ICursor.GetCrossPlatformContext
-  if (check_idx_file(scheduler->path) != ERROR_SUCCESS)
-  {
-    fwprintf(
-        stderr,
-        L"Error: idx file not found\nTo use the Virtual Alloc mode, you need "
-        L"to have the idx file alongside the run file recored by TTD.\nIf "
-        L"you "
-        L"don't have this file, you can generate it automatically by "
-        L"openning the .run file with WinDbg.\n");
-    return ERROR_INTERNAL_FATAL_ERROR;
-  }
-
   // Save current cursor position
   Position* last = scheduler->engine->IReplayEngine->GetLastPosition(
       scheduler->engine);