Skip to content

Commit

Permalink
Clean up code
Browse files Browse the repository at this point in the history
  • Loading branch information
@crosire
crosire committed Dec 15, 2018
1 parent 6d2df47 commit 2469d01
Show file tree
Hide file tree
Showing 2 changed files with 25 additions and 31 deletions.
39 changes: 16 additions & 23 deletions source/blink.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -86,44 +86,37 @@ void blink::application::run()

{ print("Reading PE debug info directory ...");

guid pdb_guid;
std::string pdb_path;
struct RSDS_DEBUG_FORMAT
{
uint32_t signature;
guid guid;
uint32_t age;
char path[1];
} const *debug_data = nullptr;

// Search debug directory for program debug database file name
const IMAGE_DATA_DIRECTORY debug_directory = headers->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_DEBUG];
const auto debug_directory_entries = reinterpret_cast<const IMAGE_DEBUG_DIRECTORY *>(_image_base + debug_directory.VirtualAddress);

for (unsigned int i = 0; i < debug_directory.Size / sizeof(IMAGE_DEBUG_DIRECTORY); i++)
{
if (debug_directory_entries[i].Type != IMAGE_DEBUG_TYPE_CODEVIEW)
continue;

struct RSDS_DEBUG_FORMAT
if (debug_directory_entries[i].Type == IMAGE_DEBUG_TYPE_CODEVIEW)
{
uint32_t signature;
guid guid;
uint32_t age;
char path[1];
};
debug_data = reinterpret_cast<const RSDS_DEBUG_FORMAT *>(reinterpret_cast<const BYTE *>(_image_base + debug_directory_entries[i].AddressOfRawData));

const auto data = reinterpret_cast<const RSDS_DEBUG_FORMAT *>(reinterpret_cast<const BYTE *>(_image_base + debug_directory_entries[i].AddressOfRawData));

if (data->signature == 0x53445352) // RSDS
{
pdb_guid = data->guid;
pdb_path = data->path;
break;
if (debug_data->signature == 0x53445352) // RSDS
break;
}
}

if (!pdb_path.empty())
if (debug_data != nullptr)
{
print(" Found program debug database: " + pdb_path);
print(" Found program debug database: " + std::string(debug_data->path));

pdb_reader pdb(pdb_path);
pdb_reader pdb(debug_data->path);

// Check if the debug information actually matches the executable
if (pdb.guid() == pdb_guid)
if (pdb.guid() == debug_data->guid)
{
const auto pdb_symbols = pdb.symbols(_image_base);
const auto pdb_source_files = pdb.sourcefiles();
Expand Down Expand Up @@ -247,7 +240,7 @@ void blink::application::run()

BYTE watcher_buffer[4096];

// Check for modifications to any of the source code files (need not monitor renames as well because some editors modify temporary files before renaming them to the actual one)
// Check for modifications to any of the source code files
while (ReadDirectoryChangesW(watcher_handle, watcher_buffer, sizeof(watcher_buffer), TRUE, FILE_NOTIFY_CHANGE_LAST_WRITE | FILE_NOTIFY_CHANGE_FILE_NAME, &size, nullptr, nullptr))
{
bool first_notification = true;
Expand Down
17 changes: 9 additions & 8 deletions source/main.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -37,11 +37,12 @@ DWORD CALLBACK remote_main(BYTE *image_base)
const auto headers = reinterpret_cast<const IMAGE_NT_HEADERS *>(image_base + reinterpret_cast<const IMAGE_DOS_HEADER *>(image_base)->e_lfanew);

// Apply base relocations
auto relocation = reinterpret_cast<const IMAGE_BASE_RELOCATION *>(image_base + headers->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_BASERELOC].VirtualAddress);
const auto relocation_delta = image_base - reinterpret_cast<const BYTE *>(headers->OptionalHeader.ImageBase);

if (relocation_delta != 0) // No need to relocate anything if the delta is zero
{
auto relocation = reinterpret_cast<const IMAGE_BASE_RELOCATION *>(image_base + headers->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_BASERELOC].VirtualAddress);

while (relocation->VirtualAddress != 0)
{
const auto field_count = (relocation->SizeOfBlock - sizeof(IMAGE_BASE_RELOCATION)) / sizeof(WORD);
Expand All @@ -61,7 +62,7 @@ DWORD CALLBACK remote_main(BYTE *image_base)
*reinterpret_cast<UINT64 *>(image_base + relocation->VirtualAddress + (field & 0xFFF)) += static_cast<INT64>(relocation_delta);
break;
default:
return 1; // Exit when encountering an unknown relocation type
return ERROR_IMAGE_AT_DIFFERENT_BASE; // Exit when encountering an unknown relocation type
}
}

Expand All @@ -88,28 +89,29 @@ DWORD CALLBACK remote_main(BYTE *image_base)
{
if (IMAGE_SNAP_BY_ORDINAL(import_name_table[k].u1.Ordinal))
{
// Import by ordinal
const auto import = IMAGE_ORDINAL(import_name_table[k].u1.Ordinal);

import_address_table[k].u1.AddressOfData = reinterpret_cast<DWORD_PTR>(GetProcAddress(module, reinterpret_cast<LPCSTR>(import)));
}
else
{
// Import by function name
const auto import = reinterpret_cast<const IMAGE_IMPORT_BY_NAME *>(image_base + import_name_table[k].u1.AddressOfData);

import_address_table[k].u1.AddressOfData = reinterpret_cast<DWORD_PTR>(GetProcAddress(module, import->Name));
}
}
}

// Call constructors
// Call global C/C++ constructors
_initterm(__xi_a, __xi_z);
_initterm(__xc_a, __xc_z);
#pragma endregion

// Run main loop
blink::application().run();

// Clean up handles
CloseHandle(console);

return 0;
Expand All @@ -132,7 +134,7 @@ int main(int argc, char *argv[])
for (int i = 1; i < argc; ++i, command_line += ' ')
command_line += argv[i];

if (!CreateProcessA(nullptr, const_cast<char *>(command_line.data()), nullptr, nullptr, FALSE, CREATE_NEW_CONSOLE, nullptr, nullptr, &startup_info, &process_info))
if (!CreateProcessA(nullptr, command_line.data(), nullptr, nullptr, FALSE, CREATE_NEW_CONSOLE, nullptr, nullptr, &startup_info, &process_info))
{
std::cout << "Failed to start target application process!" << std::endl;
return GetLastError();
Expand All @@ -152,9 +154,8 @@ int main(int argc, char *argv[])
}

// Open target application process
const DWORD access = PROCESS_CREATE_THREAD | PROCESS_VM_OPERATION | PROCESS_VM_READ | PROCESS_VM_WRITE | PROCESS_DUP_HANDLE | PROCESS_QUERY_INFORMATION;
const HANDLE local_process = GetCurrentProcess();
const HANDLE remote_process = OpenProcess(access, FALSE, pid);
const HANDLE remote_process = OpenProcess(PROCESS_CREATE_THREAD | PROCESS_VM_OPERATION | PROCESS_VM_READ | PROCESS_VM_WRITE | PROCESS_DUP_HANDLE | PROCESS_QUERY_INFORMATION, FALSE, pid);

if (remote_process == nullptr)
{
Expand All @@ -171,7 +172,7 @@ int main(int argc, char *argv[])
CloseHandle(remote_process);

std::cout << "Machine architecture mismatch between target application and this application!" << std::endl;
return ERROR_BAD_ENVIRONMENT;
return ERROR_IMAGE_MACHINE_TYPE_MISMATCH;
}

std::cout << "Launching in target application ..." << std::endl;
Expand Down

0 comments on commit 2469d01

Please sign in to comment.