diff --git a/injector/src/inject.c b/injector/src/inject.c
index 9305b64..ea85020 100644
--- a/injector/src/inject.c
+++ b/injector/src/inject.c
@@ -1,5 +1,22 @@
 #include <inject.h>
 
+#define JUMP_SIZE (6 + sizeof(void*))
+
+// Original values to recover after the injection
+// Recovery is performed by the assembly payload
+#pragma pack(push, 1)
+struct recovery_data {
+    void *entryPointAddress;
+    char entryPointData[JUMP_SIZE];
+
+    void *importDescriptorAddress;
+    IMAGE_IMPORT_DESCRIPTOR importDescriptorData;
+
+    void *sizeFieldAddress;
+    DWORD sizeFieldData;
+};
+#pragma pack(pop)
+
 static inline void write_protected_process_memory(HANDLE process, void *address, const void *buf, size_t size) {
     DWORD oldProtect;
     VirtualProtectEx(process, address, size, PAGE_EXECUTE_READWRITE, &oldProtect);
@@ -13,13 +30,6 @@ static inline void write_protected_process_memory(HANDLE process, void *address,
 void inject(HANDLE process, const void *payload, size_t payloadSize, const wchar_t *dllPath) {
     size_t _; // Contrary to the docs, {Write,Read}ProcessMemory likes to crash if the last arg is NULL
 
-    // Inject the loader into the module
-    size_t dllPathLen = (wcslen(dllPath) + 1) * sizeof(wchar_t);
-
-    char *remoteAlloc = VirtualAllocEx(process, NULL, payloadSize + dllPathLen, MEM_COMMIT, PAGE_EXECUTE_READWRITE);
-    WriteProcessMemory(process, remoteAlloc, payload, payloadSize, &_);
-    WriteProcessMemory(process, remoteAlloc + payloadSize, dllPath, dllPathLen, &_);
-
     // Find the EXE header in the process
     char exeHeader[1024];
     IMAGE_DOS_HEADER *dosHeader = NULL;
@@ -64,25 +74,60 @@ void inject(HANDLE process, const void *payload, size_t payloadSize, const wchar
 
     char *exe = (char*)memoryInfo.BaseAddress;
 
+
+    // Inject the loader into the process
+    const unsigned char JUMP_INST[] = { 0xFF, 0x25, 0x00, 0x00, 0x00, 0x00 };
+   
+    size_t dllPathSize = (wcslen(dllPath) + 1) * sizeof(wchar_t);
+
+    size_t allocSize = payloadSize + dllPathSize + sizeof(struct recovery_data);
+    char *remoteAlloc = VirtualAllocEx(process, NULL, allocSize, MEM_COMMIT, PAGE_EXECUTE_READWRITE);
+
+    // Write the assembly payload and dll path
+    WriteProcessMemory(process, remoteAlloc, payload, payloadSize, &_);
+    WriteProcessMemory(process, remoteAlloc + payloadSize, dllPath, dllPathSize, &_);
+
+
+    // Modify the executable to run the assembly payload
+    // Recovery data structure
+    struct recovery_data rd;
+
     // Replace the entry point with a jump to the loader
     char *entryPoint = exe + ntHeaders->OptionalHeader.AddressOfEntryPoint;
 
-    const unsigned char JUMP_INST[] = { 0xFF, 0x25, 0x00, 0x00, 0x00, 0x00 };
+    // Save the original entry point address and bytes
+    rd.entryPointAddress = entryPoint;
+    ReadProcessMemory(process, rd.entryPointAddress, rd.entryPointData, sizeof(rd.entryPointData), &_);
 
+    // Replace the entry point with a jump to the assembly payload
     write_protected_process_memory(process, entryPoint, JUMP_INST, sizeof(JUMP_INST));
     write_protected_process_memory(process, entryPoint + sizeof(JUMP_INST), &remoteAlloc, sizeof(remoteAlloc));
 
+
     // Break the import table to prevent any dlls from being loaded
     // Step 1: break the first import descriptor
     char *importDescriptors = exe + ntHeaders->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT].VirtualAddress;
+
+    // Save the original descriptor address and bytes
+    rd.importDescriptorAddress = importDescriptors;
+    ReadProcessMemory(process, rd.importDescriptorAddress, &rd.importDescriptorData, sizeof(rd.importDescriptorData), &_);
+
+    // Overwrite with zeroes    
     IMAGE_IMPORT_DESCRIPTOR firstDescriptor;
     ZeroMemory(&firstDescriptor, sizeof(firstDescriptor));
-
     write_protected_process_memory(process, importDescriptors, &firstDescriptor, sizeof(firstDescriptor));
 
     // Step 2: break the image data directory entry
-    size_t ddOffset = ((char*)&(ntHeaders->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT].Size)) - exeHeader;
-    DWORD newSize = 0;
+    char* ddAddr = ((char*)&(ntHeaders->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT].Size)) - exeHeader + exe;
+    
+    // Save the original value
+    rd.sizeFieldAddress = ddAddr;
+    ReadProcessMemory(process, rd.sizeFieldAddress, &rd.sizeFieldData, sizeof(rd.sizeFieldData), &_);
 
-    write_protected_process_memory(process, exe + ddOffset, &newSize, sizeof(newSize));
+    // Set to 0
+    DWORD newSize = 0;
+    write_protected_process_memory(process, ddAddr, &newSize, sizeof(newSize));
+
+    // Write recovery data to the allocation
+    WriteProcessMemory(process, remoteAlloc + payloadSize + dllPathSize, &rd, sizeof(rd), &_);
 }