1 files changed, 5 insertions, 13 deletions

diff --git a/injector.c b/injector.c
index ea53f61..9dbfe40 100644
--- a/injector.c
+++ b/injector.c
@@ -246,18 +246,10 @@ static int prepare_last_section(pe_buffer *pe, uint8_t **ptrToShellcode)
 }
 
 static uint8_t *
-RvaToPtr(uint8_t *buf, off_t rva)
+RvaToPtr(pe_buffer *pe, off_t rva)
 {
-    PIMAGE_DOS_HEADER dosHdr = (PIMAGE_DOS_HEADER) buf;
-    PIMAGE_FILE_HEADER fileHdr =
-        (PIMAGE_FILE_HEADER)(buf +
-                             dosHdr->e_lfanew +
-                             sizeof(DWORD));
-    PIMAGE_SECTION_HEADER sections =
-        (PIMAGE_SECTION_HEADER)(buf +
-        dosHdr->e_lfanew +
-        sizeof(IMAGE_NT_HEADERS));
-    size_t nSections = fileHdr->NumberOfSections;
+    PIMAGE_SECTION_HEADER sections = pe->secHdr;
+    size_t nSections = pe->fileHdr->NumberOfSections;
     off_t pos = 0;
 
     for (SIZE_T i = 0; i < nSections; ++i) {
@@ -267,7 +259,7 @@ RvaToPtr(uint8_t *buf, off_t rva)
         }
         if (rva < pos) {
             rva -= sections[i].VirtualAddress;
-            return buf + rva + sections[i].PointerToRawData;
+            return pe->buf + rva + sections[i].PointerToRawData;
         }
     }
     return NULL;
@@ -313,7 +305,7 @@ static int search_patch_near_entry(pe_buffer *pe, patch_ctx *pa)
 {
     const size_t maxbytes = 30;
     size_t i;
-    uint8_t *insts8 = RvaToPtr(pe->buf,
+    uint8_t *insts8 = RvaToPtr(pe,
                                pe->optHdr->AddressOfEntryPoint);
     uint16_t *insts16;
     uint32_t relAddr;