1
\$\begingroup\$

https://github.com/speedrun-program/load-extender

I posted a version of these programs for review about a week ago and ended up getting good advice on how to improve them. this was the post which shows the old code: C and C++ programs which let you lengthen the time it takes to access files

Now that I've improved the code, I'm posting the new version here for review since this page says I'm allowed to.

I did everything suggested by the accepted answer in the original post except for using std::filesystem::path objects in the hash map because I couldn't figure out how to do that.

*sorry that I've been making so many small adjustments to the code occasionally. But I think at this point it should be complete.

struct_and_functions.h:

#include <stdint.h>
#include <iostream>
#include <fstream>
#include <string>
#include <atomic>
#include <vector>
#include <exception>

static bool TOO_BIG_DELAY_SEEN = false;

struct delay_sequence
{
    std::atomic<std::size_t> index;
    std::vector<std::chrono::milliseconds> delays;
    bool repeat;
    bool reset_all;

    delay_sequence() { index = 0; repeat = false; reset_all = false; }
    delay_sequence(const delay_sequence& copy_sequence)
    {
        index = copy_sequence.index.load();
        delays = copy_sequence.delays;
        repeat = copy_sequence.repeat;
        reset_all = copy_sequence.reset_all;
    }
    delay_sequence& operator=(delay_sequence copy_sequence)
    {
        index = copy_sequence.index.load();
        delays = copy_sequence.delays;
        repeat = copy_sequence.repeat;
        reset_all = copy_sequence.reset_all;
        return *this;
    }
};

static std::size_t setup_delay_sequence(delay_sequence& current_sequence, std::wstring& line, std::size_t start_position)
{
    std::size_t delay_sequence_length = 1;
    for (std::size_t line_position = start_position; line_position < line.length(); line_position++)
    {
        if (line[line_position] == L'/')
        {
            ++delay_sequence_length;
        }
        else if (line[line_position] == L'-')
        {
            --delay_sequence_length;
            if (delay_sequence_length == 0)
            {
                current_sequence.reset_all = true;
                break;
            }
            current_sequence.repeat = true;
            break;
        }
    }
    return delay_sequence_length;
}

static void remove_non_digits(std::wstring& delay_substr)
{
    std::size_t position_to_write = 0;
    for (std::size_t i = 0; i < delay_substr.length(); i++)
    {
        if (isdigit(delay_substr[i]))
        {
            delay_substr[position_to_write] = delay_substr[i];
            ++position_to_write;
        }
    }
    delay_substr.erase(position_to_write);
}

static void set_delays(std::vector<std::chrono::milliseconds>& delays, std::wstring& line, std::size_t delay_sequence_length, std::size_t slash_position)
{
    delays.reserve(delay_sequence_length);
    unsigned long long max_delay = std::chrono::milliseconds::max().count();
    while (slash_position != std::wstring::npos && delays.size() != delay_sequence_length)
    {
        std::size_t next_slash_position = line.find(L'/', slash_position + 1);
        std::wstring delay_substr = line.substr(slash_position + 1, next_slash_position - slash_position - 1);
        remove_non_digits(delay_substr);
        if (delay_substr.length() == 0)
        {
            delays.push_back(std::chrono::milliseconds(0));
            slash_position = next_slash_position;
            continue;
        }
        try
        {
            unsigned long long delay_time = std::stoull(delay_substr);
            if (delay_time > max_delay)
            {
                throw std::out_of_range(NULL);
            }
            delays.push_back(std::chrono::milliseconds(delay_time));
        }
        catch (const std::out_of_range& oor)
        {
            if (!TOO_BIG_DELAY_SEEN)
            {
                printf("delay time can only be %llu milliseconds\n", max_delay);
                TOO_BIG_DELAY_SEEN = true; // TOO_BIG_DELAY_SEEN only changed here
            }
            delays.push_back(std::chrono::milliseconds(max_delay));
        }
        slash_position = next_slash_position;
    }
}

static void set_key_and_value(rh::unordered_flat_map<std::wstring, delay_sequence>& my_rh_map, std::wstring& line)
{
    std::size_t slash_position = line.find(L'/');
    if (slash_position == std::wstring::npos)
    {
        return;
    }
    std::size_t position_before_whitespace = line.find_last_not_of(L" \t\f\v\r/", slash_position);
    if (position_before_whitespace == std::wstring::npos)
    {
        return;
    }
    std::wstring file_name = line.substr(0, position_before_whitespace + 1);
    delay_sequence current_sequence;
    std::size_t delay_sequence_length = setup_delay_sequence(current_sequence, line, slash_position + 1);
    set_delays(current_sequence.delays, line, delay_sequence_length, slash_position);
    my_rh_map[file_name] = current_sequence;
}

static unsigned char rh_map_setup(rh::unordered_flat_map<std::wstring, delay_sequence>& my_rh_map)
{
    std::size_t line_count = 1;
    std::wstring line;
    std::wifstream my_file;
    my_file.open("files_and_delays.txt");
    if (my_file.fail())
    {
        printf("couldn't open files_and_delays.txt\n");
        return 1;
    }
    while (std::getline(my_file, line) && line_count < SIZE_MAX)
    {
        ++line_count;
    }
    if (line_count == SIZE_MAX)
    {
        printf("can't store %zu delays\n", SIZE_MAX);
        return 1;
    }
    my_rh_map.reserve(line_count);
    my_file.clear();
    if (!my_file.seekg(0, std::ios::beg))
    {
        printf("seekg on files_and_delays.txt failed\n");
        return 1;
    }
    line.clear();
    while (std::getline(my_file, line))
    {
        set_key_and_value(my_rh_map, line);
        line.clear();
    }
    my_file.close();
    return 0;
}

static void delay_file(rh::unordered_flat_map<std::wstring, delay_sequence>& my_rh_map, std::wstring& file_name)
{
    rh::unordered_flat_map<std::wstring, delay_sequence>::iterator rh_map_iter = my_rh_map.find(file_name);
    if (rh_map_iter != my_rh_map.end()) // key found
    {
        printf("%ls successfully found in hash map\n", file_name.c_str());
        delay_sequence& found_sequence = rh_map_iter->second;
        if (found_sequence.reset_all) // reset all delay sequences
        {
            for (auto& it : my_rh_map)
            {
                it.second.index = 0;
            }
            printf("all delay sequences reset\n");
        }
        else
        {
            if (found_sequence.repeat) // reset delay sequence
            {
                found_sequence.index = found_sequence.index.load() % found_sequence.delays.size();
                printf("%ls delay sequence reset\n", file_name.c_str());
            }
            if (found_sequence.index < found_sequence.delays.size())
            {
                // this is defined in the main file so it won't sleep if being used with load_extender_test.exe
                call_sleep_thread(found_sequence.delays[found_sequence.index]);
                found_sequence.index += 1;
            }
            else // delay sequence already finished
            {
                printf("%ls delay sequence already finished\n", file_name.c_str());
            }
        }
    }
    else // key not found
    {
        printf("%ls not found in hash map\n", file_name.c_str());
    }
}

load_extender_test.cpp

#include <chrono>
#include <thread>
#include "robin_hood.h"
// thank you for making this hash map, martinus.
// here's his github: https://github.com/martinus/robin-hood-hashing

// this is defined here so delay_file function will use the right version
static void call_sleep_thread(std::chrono::milliseconds duration)
{
    unsigned long long total_milliseconds = duration.count();
    unsigned long long total_seconds = total_milliseconds / 1000;
    short remaining_milliseconds = total_milliseconds % 1000;
    printf("sleep for %llu second(s) and %d millisecond(s)\n", total_seconds, remaining_milliseconds);
}

namespace rh = robin_hood;

#include "struct_and_functions.h" // included here so alias can be used

// these are global so the hooked function can see them
static rh::unordered_flat_map<std::wstring, delay_sequence> my_rh_map;
// SETUP_SUCCEEDED only changed here
static unsigned char SETUP_SUCCEEDED = rh_map_setup(my_rh_map);

static void print_rh_map()
{
    printf("\n---------- hash map current state ----------\n");
    for (auto& it : my_rh_map)
    {
        printf("%ls / ", it.first.c_str());
        delay_sequence& sequence_to_print = it.second;
        for (std::size_t i = 0; i < sequence_to_print.delays.size(); i++)
        {
            // casted to unsigned long long to get rid of warning message
            // visual studio on windows wants %lu, but gcc on linux wants %llu
            printf("%llu ", (unsigned long long)sequence_to_print.delays[i].count());
        }
        if (sequence_to_print.reset_all)
        {
            printf("RESET_ALL ");
        }
        else if (sequence_to_print.repeat)
        {
            printf("REPEAT ");
        }
        printf(": INDEX %zu", sequence_to_print.index.load());
        printf("\n");
    }
    printf("---------------------------------------------\n\n");
}

static void test_all_inputs()
{
    std::wstring test_path;
    std::wifstream input_test_file;
    input_test_file.open("test_input.txt");
    if (input_test_file.fail())
    {
        printf("couldn't open test_input.txt\n");
        return;
    }
    if (SETUP_SUCCEEDED == 0)
    {
        print_rh_map();
        while (std::getline(input_test_file, test_path))
        {
            printf("testing input: %ls\n", test_path.c_str());
            delay_file(my_rh_map, test_path);
            print_rh_map();
            test_path.clear();
        }
    }
    else
    {
        printf("setup failed\n");
    }
    input_test_file.close();
}

int main(int argc, char* argv[])
{
    printf("\ntest start\n");
    test_all_inputs();
    printf("test finished, press Enter to exit\n");
    std::wstring input;
    std::getline(std::wcin, input);
    return 0;
}

load_extender.cpp

#include <chrono>
#include <thread>
#include <stdio.h>
#include <dlfcn.h>
#include "robin_hood.h"
// thank you for making this hash map, martinus.
// here's his github: https://github.com/martinus/robin-hood-hashing

#define DISABLE_PRINTF

#ifdef DISABLE_PRINTF
    #define printf(fmt, ...) (0)
#endif

// this is defined here so delay_file function will use the right version
static void call_sleep_thread(std::chrono::milliseconds duration)
{
    std::this_thread::sleep_for(duration);
}

namespace rh = robin_hood;

#include "struct_and_functions.h" // included here so alias can be used

// these are global so the hooked function can see them
static rh::unordered_flat_map<std::wstring, delay_sequence> my_rh_map;
// SETUP_SUCCEEDED only changed here
static unsigned char SETUP_SUCCEEDED = rh_map_setup(my_rh_map);

static auto original_fopen = reinterpret_cast<FILE *(*)(const char *path, const char *mode)>(dlsym(RTLD_NEXT, "fopen"));

FILE *fopen(const char *path, const char *mode)
{
    if (SETUP_SUCCEEDED == 0)
    {
        // finding part of path which shows file name
        std::size_t file_name_index = SIZE_MAX; // overflows to 0 and checks entire path if no slash is found
        std::size_t end_index = 0;
        for (; path[end_index] != '\0'; end_index++)
        {
            if (path[end_index] == '/')
            {
                file_name_index = end_index;
            }
        }
        std::wstring file_name(&path[file_name_index + 1], &path[end_index]);
        delay_file(my_rh_map, file_name);
    }
    return original_fopen(path, mode);
}

load_extender_dll.cpp

// allows std::chrono::milliseconds::max().count();
#define NOMINMAX

#include <chrono>
#include <thread>
#include <Windows.h>
#include <easyhook.h>
#include "robin_hood.h"
// thank you for making this hash map, martinus.
// here's his github: https://github.com/martinus/robin-hood-hashing

#define DISABLE_PRINTF

#ifdef DISABLE_PRINTF
    #define printf(fmt, ...) (0)
#endif

// this is defined here so delay_file function will use the right version
static void call_sleep_thread(std::chrono::milliseconds duration)
{
    std::this_thread::sleep_for(duration);
}

namespace rh = robin_hood;

// included here so alias can be used
#include "struct_and_functions.h"

// these are global so the hooked function can see them
static rh::unordered_flat_map<std::wstring, delay_sequence> my_rh_map;
static unsigned char SETUP_SUCCEEDED = rh_map_setup(my_rh_map);

static NTSTATUS WINAPI NtOpenFileHook(
    PHANDLE           FileHandle,
    ACCESS_MASK        DesiredAccess,
    POBJECT_ATTRIBUTES ObjectAttributes,
    PIO_STATUS_BLOCK  IoStatusBlock,
    ULONG              ShareAccess,
    ULONG              OpenOptions)
{
    if (SETUP_SUCCEEDED == 0)
    {
        std::wstring file_path = ObjectAttributes->ObjectName->Buffer;
        // +1 so '\' isn't included. If '\' isn't found, the whole wstring is checked because npos is -1
        std::wstring file_name = file_path.substr(file_path.rfind(L"\\") + 1);
        delay_file(my_rh_map, file_name);
    }
    return NtOpenFile(FileHandle, DesiredAccess, ObjectAttributes, IoStatusBlock, ShareAccess, OpenOptions);
}

extern "C" void __declspec(dllexport) __stdcall NativeInjectionEntryPoint(REMOTE_ENTRY_INFO * inRemoteInfo);

void __stdcall NativeInjectionEntryPoint(REMOTE_ENTRY_INFO* inRemoteInfo) {
    HOOK_TRACE_INFO hHook1 = { NULL };
    LhInstallHook(
        GetProcAddress(GetModuleHandle(TEXT("ntdll")), "NtOpenFile"),
        NtOpenFileHook,
        NULL,
        &hHook1);

    ULONG ACLEntries[1] = { 0 };
    LhSetExclusiveACL(ACLEntries, 1, &hHook1);
    return;
}

load_extender_exe.cpp

#include <tchar.h>
#include <iostream>
#include <string>
#include <Windows.h>
#include <easyhook.h>

void get_exit_input()
{
    std::wcout << "Press Enter to exit";
    std::wstring input;
    std::getline(std::wcin, input);
    std::getline(std::wcin, input);
}

int _tmain(int argc, _TCHAR* argv[])
{
    WCHAR* dllToInject32 = NULL;
    WCHAR* dllToInject64 = NULL;
    LPCWSTR lpApplicationName = argv[0];
    DWORD lpBinaryType;
    if (GetBinaryType(lpApplicationName, &lpBinaryType) == 0 || (lpBinaryType != 0 && lpBinaryType != 6))
    {
        std::wcout << "ERROR: This exe wasn't identified as 32-bit or as 64-bit";
        get_exit_input();
        return 1;
    }
    else if (lpBinaryType == 0)
    {
        dllToInject32 = (WCHAR*)L"load_extender_32.dll";
    }
    else
    {
        dllToInject64 = (WCHAR*)L"load_extender_64.dll";
    }
    DWORD processId;
    std::wcout << "Enter the target process Id: ";
    std::cin >> processId;

    wprintf(L"Attempting to inject dll\n\n");

    // Inject dllToInject into the target process Id, passing 
    // freqOffset as the pass through data.
    NTSTATUS nt = RhInjectLibrary(
        processId,   // The process to inject into
        0,           // ThreadId to wake up upon injection
        EASYHOOK_INJECT_DEFAULT,
        dllToInject32, // 32-bit
        dllToInject64, // 64-bit
        NULL, // data to send to injected DLL entry point
        0 // size of data to send
    );

    if (nt != 0)
    {
        printf("RhInjectLibrary failed with error code = %d\n", nt);
        PWCHAR err = RtlGetLastErrorString();
        std::wcout << err << "\n";
        get_exit_input();
        return 1;
    }
    else
    {
        std::wcout << L"Library injected successfully.\n";
    }

    get_exit_input();
    return 0;
}
```
\$\endgroup\$
1

1 Answer 1

4
\$\begingroup\$

General Observations

Interesting project. The code looks a little too much like C and not enough like C++. You might want to convert the structs to classes to take advatage of data encapsulation.

I suggest that you follow more of the advice given in the answer on your first question:

Make your code as platform-independent as possible

You should not need two completely different implementations. A lot of code can be shared. The only things that are different between Linux and Windows are the function calls that need to be intercepted, and how you intercept them.

Rather than have 2 separate code bases, one for Linux and one for Windows, the code that is platform specific should be within ifdefs.

The major problem with having 2 separate sets of files for a program is that it doubles the cost of the maintenance. If there is a bug in the code it needs to be fixed in 2 places rather than one.

Follow Conventions

A convention that started in the C program is that macros or symbolic constants should be in all caps and regular variables should follow the naming convention used throughout the rest of the program whether that is snake_case or camelCase. The following statement is a violation of that convention:

static bool TOO_BIG_DELAY_SEEN = false;

since TOO_BIG_DELAY_SEEN is not a const declaration.

It is generally a bad idea to declare a variable in a header file since that can lead to linking problems. In this case the static declaration localizes the variable but it does exist in every file that includes the header.

Function Bodies Should be in .cpp Files

The code defines a structure with a number of functions in the file struct_and_functions.h. For a number of reasons it is better to define these as pointers to functions in the struct and define the body of the function in a .cpp file. This means that the code will only be compiled once rather than multiple times and can prevent linking problems. It also decreases build times since the functions are supplied by a .o file (Linux) or .obj file (Windows). This will decrease the size of the binary is that header file is included in multiple .cpp files.

\$\endgroup\$
5
  • \$\begingroup\$ Thank you for helping with this, I’ll try to change it using your advice. For TOO_BIG_DELAY_SEEN would it be better if I made it in rh_map_setup and passed it to set_key_and_value and set_delays? \$\endgroup\$ Jan 12, 2021 at 21:51
  • 1
    \$\begingroup\$ I can't find the declaration/body of rh_map_setup() in your code, only the usage, so I can't say. \$\endgroup\$
    – pacmaninbw
    Jan 12, 2021 at 22:03
  • \$\begingroup\$ They’re in struct_and_functions.h, the 2nd and 3rd lowest functions. I ended up making it global because I thought it seemed clumsy to pass it up multiple functions just to use it in a single area to prevent a message from repeating. But I maybe that’s the wrong way to deal with it. \$\endgroup\$ Jan 12, 2021 at 22:11
  • \$\begingroup\$ I tried looking up mutable default arguments for c++ but I didn’t see anything for it. \$\endgroup\$ Jan 12, 2021 at 22:30
  • \$\begingroup\$ Nevermind, I think I finally found something. I can make it in set_delays as a static local variable. \$\endgroup\$ Jan 12, 2021 at 22:46

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.