From 5c28287cd86cb1e394190f6a308dcc4aae5d95bd Mon Sep 17 00:00:00 2001 From: seele Date: Mon, 23 Mar 2026 16:22:31 +0800 Subject: [PATCH 1/4] feat: implement path resolution and environment variable retrieval in win payload --- src/catter-hook/win/payload/exports.def | 2 + src/catter-hook/win/payload/main.cc | 93 +----- src/catter-hook/win/payload/resolver.h | 131 +++++++++ src/catter-hook/win/payload/util.cc | 325 +++++++++++++++++++++ src/catter-hook/win/payload/util.h | 17 ++ src/catter-hook/win/payload/winapi.h | 76 +++++ tests/unit/catter-hook/win/payload/main.cc | 2 - tests/unit/catter-hook/win/payload/util.cc | 155 ++++++++++ xmake.lua | 9 +- 9 files changed, 720 insertions(+), 90 deletions(-) create mode 100644 src/catter-hook/win/payload/exports.def create mode 100644 src/catter-hook/win/payload/resolver.h create mode 100644 src/catter-hook/win/payload/util.cc create mode 100644 src/catter-hook/win/payload/util.h create mode 100644 src/catter-hook/win/payload/winapi.h delete mode 100644 tests/unit/catter-hook/win/payload/main.cc create mode 100644 tests/unit/catter-hook/win/payload/util.cc diff --git a/src/catter-hook/win/payload/exports.def b/src/catter-hook/win/payload/exports.def new file mode 100644 index 00000000..47a75ae8 --- /dev/null +++ b/src/catter-hook/win/payload/exports.def @@ -0,0 +1,2 @@ +LIBRARY catter-hook-win64 +EXPORTS diff --git a/src/catter-hook/win/payload/main.cc b/src/catter-hook/win/payload/main.cc index 16c552e1..39f65ae1 100644 --- a/src/catter-hook/win/payload/main.cc +++ b/src/catter-hook/win/payload/main.cc @@ -1,91 +1,12 @@ #include -#include #include -#include #include #include -#include -#include -#include #include #include -#include "win/env.h" - -namespace catter::win { -namespace { -template -std::basic_string get_app_name(const char_t* application_name, const char_t* command_line) { - if(application_name == nullptr) { - if(command_line == nullptr) { - return {}; - } - auto view = std::basic_string_view(command_line); - auto first_space = view.find_first_of(char_t(' ')); - return std::basic_string(view.substr(0, first_space)); - } else { - return std::basic_string(application_name); - } -} - -template -DWORD FixGetEnvironmentVariable(const char_t* name, char_t* buffer, DWORD size) { - if constexpr(std::is_same_v) { - return GetEnvironmentVariableA(name, buffer, size); - } else { - return GetEnvironmentVariableW(name, buffer, size); - } -} - -template -std::basic_string get_proxy_path() { - constexpr size_t buffer_size = 256; - char_t buffer[buffer_size]; - - auto len = FixGetEnvironmentVariable(catter::win::ENV_VAR_PROXY_PATH, - buffer, - buffer_size); - if(len == 0) { - return {}; - } - - if(len < buffer_size) { - return std::basic_string(buffer, len); - } - - std::basic_string path; - path.resize(len); - FixGetEnvironmentVariable(catter::win::ENV_VAR_PROXY_PATH, path.data(), len); - path.pop_back(); - return path; -} - -template -std::basic_string get_ipc_id() { - constexpr size_t buffer_size = 64; - char_t buffer[buffer_size]; - - auto len = - FixGetEnvironmentVariable(catter::win::ENV_VAR_IPC_ID, buffer, buffer_size); - if(len == 0) { - return {}; - } - - if(len < buffer_size) { - return std::basic_string(buffer, len); - } - - std::basic_string id; - id.resize(len); - FixGetEnvironmentVariable(catter::win::ENV_VAR_IPC_ID, id.data(), len); - id.pop_back(); - return id; -} - -} // namespace - -} // namespace catter::win +#include "win/payload/util.h" // Use anonymous namespace to avoid exporting symbols namespace { @@ -110,9 +31,9 @@ struct CreateProcessA { auto converted_cmdline = std::format("{} -p {} --exec {} -- {}", - catter::win::get_proxy_path(), - catter::win::get_ipc_id(), - catter::win::get_app_name(lpApplicationName, lpCommandLine), + catter::win::payload::get_proxy_path(), + catter::win::payload::get_ipc_id(), + catter::win::payload::resolve_abspath(lpApplicationName, lpCommandLine), std::string_view(lpCommandLine ? lpCommandLine : "")); return original(nullptr, @@ -146,9 +67,9 @@ struct CreateProcessW { auto converted_cmdline = std::format(L"{} -p {} --exec {} -- {}", - catter::win::get_proxy_path(), - catter::win::get_ipc_id(), - catter::win::get_app_name(lpApplicationName, lpCommandLine), + catter::win::payload::get_proxy_path_wide(), + catter::win::payload::get_ipc_id_wide(), + catter::win::payload::resolve_abspath(lpApplicationName, lpCommandLine), std::wstring_view(lpCommandLine ? lpCommandLine : L"")); return original(nullptr, diff --git a/src/catter-hook/win/payload/resolver.h b/src/catter-hook/win/payload/resolver.h new file mode 100644 index 00000000..fc168684 --- /dev/null +++ b/src/catter-hook/win/payload/resolver.h @@ -0,0 +1,131 @@ +#pragma once + +#include +#include +#include +#include + +#include "winapi.h" + +namespace catter::win::payload { + +template +class Resolver { +public: + explicit Resolver(std::vector> search_paths) + : m_search_paths(std::move(search_paths)) {} + + std::basic_string resolve(std::basic_string_view app_name) const { + if(app_name.empty()) { + return {}; + } + + auto original_name = std::basic_string(app_name); + auto fixed_name = fix_app_name(app_name); + if(contains_path(fixed_name)) { + auto direct = to_absolute(fixed_name); + if(!direct.empty() && is_file(direct)) { + return direct; + } + // Keep original input when the path cannot be resolved to an existing file. + return original_name; + } + + for(const auto& search_path: m_search_paths) { + auto candidate = join_path(search_path, fixed_name); + auto absolute_candidate = to_absolute(candidate); + if(absolute_candidate.empty()) { + continue; + } + if(is_file(absolute_candidate)) { + return absolute_candidate; + } + } + + // Keep original input when search paths do not resolve an existing file. + return original_name; + } + +private: + static constexpr char_t path_sep = char_t('\\'); + + static constexpr bool is_path_sep(char_t c) { + return c == char_t('\\') || c == char_t('/'); + } + + static bool contains_path(std::basic_string_view value) { + return std::find_if(value.begin(), value.end(), [](char_t c) { + return is_path_sep(c) || c == char_t(':'); + }) + != value.end(); + } + + static bool has_extension(std::basic_string_view file_name) { + return file_name.find(char_t('.')) != std::basic_string_view::npos; + } + + static std::basic_string_view exe_suffix() { + if constexpr(std::is_same_v) { + return ".exe"; + } else { + return L".exe"; + } + } + + static std::basic_string fix_app_name(std::basic_string_view app_name) { + if(contains_path(app_name) || app_name.back() == char_t('.') || has_extension(app_name)) { + return std::basic_string(app_name); + } + + std::basic_string fixed(app_name); + fixed.append(exe_suffix()); + return fixed; + } + + static std::basic_string join_path(std::basic_string_view directory, + std::basic_string_view file_name) { + if(directory.empty()) { + return std::basic_string(file_name); + } + + std::basic_string out(directory); + if(!is_path_sep(out.back())) { + out.push_back(path_sep); + } + out.append(file_name); + return out; + } + + static std::basic_string to_absolute(std::basic_string_view path) { + std::basic_string input(path); + auto required_size = FixGetFullPathName(input.c_str(), 0, nullptr, nullptr); + if(required_size == 0) { + return {}; + } + + std::basic_string absolute(required_size, char_t('\0')); + auto written = + FixGetFullPathName(input.c_str(), required_size, absolute.data(), nullptr); + if(written == 0 || written >= required_size) { + return {}; + } + + absolute.resize(written); + return absolute; + } + + static bool is_file(std::basic_string_view path) { + auto attrs = FixGetFileAttributes(std::basic_string(path).c_str()); + return attrs != INVALID_FILE_ATTRIBUTES && (attrs & FILE_ATTRIBUTE_DIRECTORY) == 0; + } + + std::vector> m_search_paths; +}; + +template +Resolver create_app_name_resolver(); + +template +Resolver create_command_line_resolver(); + +} \ No newline at end of file diff --git a/src/catter-hook/win/payload/util.cc b/src/catter-hook/win/payload/util.cc new file mode 100644 index 00000000..f3895b10 --- /dev/null +++ b/src/catter-hook/win/payload/util.cc @@ -0,0 +1,325 @@ +#include +#include +#include +#include + +#include + +#include "win/env.h" + +#include "resolver.h" +#include "util.h" + + +namespace catter::win::payload { + +namespace { + +template +constexpr char_t k_path_delimiter = char_t(';'); + +template +constexpr char_t k_path_sep = char_t('\\'); + +template +std::vector> split_path_var(std::basic_string_view value) { + std::vector> segments; + size_t start = 0; + while(start <= value.size()) { + auto split_pos = value.find(k_path_delimiter, start); + auto token = value.substr(start, split_pos == std::basic_string_view::npos + ? value.size() - start + : split_pos - start); + if(!token.empty()) { + segments.emplace_back(token); + } + + if(split_pos == std::basic_string_view::npos) { + break; + } + start = split_pos + 1; + } + return segments; +} + +template +void push_if_non_empty(std::vector>& out, + std::basic_string value) { + if(!value.empty()) { + out.push_back(std::move(value)); + } +} + +template +std::basic_string get_env_var(const char_t* name, size_t stack_buffer_size) { + std::basic_string value; + value.resize(stack_buffer_size); + + auto len = FixGetEnvironmentVariable(name, value.data(), static_cast(value.size())); + if(len == 0) { + return {}; + } + + if(len < value.size()) { + value.resize(len); + return value; + } + + value.resize(len); + len = FixGetEnvironmentVariable(name, value.data(), static_cast(value.size())); + if(len == 0 || len > value.size()) { + return {}; + } + + if(len > 0 && value[len - 1] == char_t('\0')) { + value.resize(len - 1); + } else { + value.resize(len); + } + + return value; +} + +template +std::basic_string get_current_directory() { + auto required_size = FixGetCurrentDirectory(0, nullptr); + if(required_size == 0) { + return {}; + } + + std::basic_string value(required_size, char_t('\0')); + auto written = FixGetCurrentDirectory(required_size, value.data()); + if(written == 0 || written >= required_size) { + return {}; + } + + value.resize(written); + return value; +} + +template +std::basic_string get_module_directory() { + std::basic_string module_path(MAX_PATH, char_t('\0')); + while(true) { + auto written = FixGetModuleFileName(nullptr, + module_path.data(), + static_cast(module_path.size())); + if(written == 0) { + return {}; + } + if(written < module_path.size() - 1) { + module_path.resize(written); + break; + } + module_path.resize(module_path.size() * 2); + } + + auto pos = module_path.find_last_of(k_path_sep); + if(pos == std::basic_string::npos) { + return {}; + } + + return module_path.substr(0, pos); +} + +template +std::basic_string get_system_directory() { + std::basic_string value(MAX_PATH, char_t('\0')); + while(true) { + auto written = + FixGetSystemDirectory(value.data(), static_cast(value.size())); + if(written == 0) { + return {}; + } + if(written < value.size()) { + value.resize(written); + return value; + } + value.resize(written + 1); + } +} + +template +std::basic_string get_windows_directory() { + std::basic_string value(MAX_PATH, char_t('\0')); + while(true) { + auto written = + FixGetWindowsDirectory(value.data(), static_cast(value.size())); + if(written == 0) { + return {}; + } + if(written < value.size()) { + value.resize(written); + return value; + } + value.resize(written + 1); + } +} + +template +std::basic_string get_current_drive_root() { + auto cwd = get_current_directory(); + if(cwd.size() < 2 || cwd[1] != char_t(':')) { + return {}; + } + + std::basic_string root; + root.push_back(cwd[0]); + root.push_back(char_t(':')); + root.push_back(k_path_sep); + return root; +} + +template +constexpr std::basic_string_view k_path_env_name = {}; + +template <> +constexpr std::basic_string_view k_path_env_name = "PATH"; + +template <> +constexpr std::basic_string_view k_path_env_name = L"PATH"; + +template +constexpr std::basic_string_view k_system16_name = {}; + +template <> +constexpr std::basic_string_view k_system16_name = "System"; + +template <> +constexpr std::basic_string_view k_system16_name = L"System"; + +} // namespace + +template +std::basic_string extract_first_token(std::basic_string_view command_line) { + constexpr char_t quote_char = char_t('"'); + constexpr char_t space_char = char_t(' '); + + auto trimmed = [](std::basic_string_view value) { + auto first_non_space = std::find_if_not(value.begin(), value.end(), [](char_t c) { + return c == space_char; + }); + return std::basic_string(first_non_space, value.end()); + }(command_line); + + std::basic_string_view view(trimmed); + if(view.empty()) { + return {}; + } + + if(view.front() == quote_char) { + auto closing_quote = view.find(quote_char, 1); + if(closing_quote == std::basic_string_view::npos) { + return std::basic_string(view.substr(1)); + } + return std::basic_string(view.substr(1, closing_quote - 1)); + } + + auto first_space = std::find_if(view.begin(), view.end(), [](char_t c) { + return c == space_char; + }); + return std::basic_string(view.begin(), first_space); +} + +template +Resolver create_app_name_resolver() { + std::vector> search_paths; + // Search order for explicit application_name: + // 1) current drive root, 2) current directory. + push_if_non_empty(search_paths, get_current_drive_root()); + push_if_non_empty(search_paths, get_current_directory()); + return Resolver(std::move(search_paths)); +} + +template +Resolver create_command_line_resolver() { + std::vector> search_paths; + + // Search order for command line token: + // 1) directory of current process image. + // 2) current directory. + // 3) 32-bit system directory. + // 4) 16-bit system directory named "System" under Windows directory. + // 5) Windows directory. + // 6) directories listed in PATH. + auto module_dir = get_module_directory(); + auto current_dir = get_current_directory(); + auto system_dir = get_system_directory(); + auto windows_dir = get_windows_directory(); + + push_if_non_empty(search_paths, std::move(module_dir)); + push_if_non_empty(search_paths, std::move(current_dir)); + push_if_non_empty(search_paths, std::move(system_dir)); + if(!windows_dir.empty()) { + auto system16_dir = windows_dir; + if(system16_dir.back() != k_path_sep && system16_dir.back() != char_t('/')) { + system16_dir.push_back(k_path_sep); + } + system16_dir.append(k_system16_name); + push_if_non_empty(search_paths, std::move(system16_dir)); + } + push_if_non_empty(search_paths, windows_dir); + + auto path_value = get_env_var(k_path_env_name.data(), 4096); + auto path_segments = split_path_var(path_value); + for(auto& segment: path_segments) { + push_if_non_empty(search_paths, std::move(segment)); + } + + return Resolver(std::move(search_paths)); +} + +template +std::basic_string resolve_abspath_impl(const char_t* application_name, + const char_t* command_line) { + std::basic_string raw_app_name; + + // CreateProcess takes lpApplicationName first; when absent, the executable + // is parsed from the first token in lpCommandLine. + if(application_name != nullptr && application_name[0] != char_t('\0')) { + raw_app_name.assign(application_name); + return create_app_name_resolver().resolve(raw_app_name); + } + + raw_app_name = extract_first_token(command_line == nullptr + ? std::basic_string_view{} + : std::basic_string_view{command_line}); + if(raw_app_name.empty()) { + return {}; + } + + return create_command_line_resolver().resolve(raw_app_name); +} + +template Resolver create_app_name_resolver(); +template Resolver create_app_name_resolver(); +template Resolver create_command_line_resolver(); +template Resolver create_command_line_resolver(); + + + +std::string resolve_abspath(const char* application_name, const char* command_line) { + return resolve_abspath_impl(application_name, command_line); +} + +std::wstring resolve_abspath(const wchar_t* application_name, const wchar_t* command_line) { + return resolve_abspath_impl(application_name, command_line); +} + +std::string get_proxy_path() { + return get_env_var(catter::win::ENV_VAR_PROXY_PATH, 256); +} + +std::wstring get_proxy_path_wide() { + return get_env_var(catter::win::ENV_VAR_PROXY_PATH, 256); +} + +std::string get_ipc_id() { + return get_env_var(catter::win::ENV_VAR_IPC_ID, 64); +} + +std::wstring get_ipc_id_wide() { + return get_env_var(catter::win::ENV_VAR_IPC_ID, 64); +} + +} // namespace catter::win::payload diff --git a/src/catter-hook/win/payload/util.h b/src/catter-hook/win/payload/util.h new file mode 100644 index 00000000..748e16d0 --- /dev/null +++ b/src/catter-hook/win/payload/util.h @@ -0,0 +1,17 @@ +#pragma once + +#include +#include + +namespace catter::win::payload { + +std::string resolve_abspath(const char* application_name, const char* command_line); +std::wstring resolve_abspath(const wchar_t* application_name, const wchar_t* command_line); + +std::string get_proxy_path(); +std::wstring get_proxy_path_wide(); + +std::string get_ipc_id(); +std::wstring get_ipc_id_wide(); + +} // namespace catter::win::payload diff --git a/src/catter-hook/win/payload/winapi.h b/src/catter-hook/win/payload/winapi.h new file mode 100644 index 00000000..e44cdcbf --- /dev/null +++ b/src/catter-hook/win/payload/winapi.h @@ -0,0 +1,76 @@ +#include + +#include + + +namespace catter::win::payload { + +template +concept CharT = std::is_same_v || std::is_same_v; + +template +DWORD FixGetEnvironmentVariable(const char_t* name, char_t* buffer, DWORD size) { + if constexpr(std::is_same_v) { + return GetEnvironmentVariableA(name, buffer, size); + } else { + return GetEnvironmentVariableW(name, buffer, size); + } +} + +template +DWORD FixGetFullPathName(const char_t* file_name, + DWORD buffer_size, + char_t* buffer, + char_t** file_part) { + if constexpr(std::is_same_v) { + return GetFullPathNameA(file_name, buffer_size, buffer, file_part); + } else { + return GetFullPathNameW(file_name, buffer_size, buffer, file_part); + } +} + +template +DWORD FixGetFileAttributes(const char_t* path) { + if constexpr(std::is_same_v) { + return GetFileAttributesA(path); + } else { + return GetFileAttributesW(path); + } +} + +template +DWORD FixGetCurrentDirectory(DWORD size, char_t* buffer) { + if constexpr(std::is_same_v) { + return GetCurrentDirectoryA(size, buffer); + } else { + return GetCurrentDirectoryW(size, buffer); + } +} + +template +DWORD FixGetModuleFileName(HMODULE module, char_t* buffer, DWORD size) { + if constexpr(std::is_same_v) { + return GetModuleFileNameA(module, buffer, size); + } else { + return GetModuleFileNameW(module, buffer, size); + } +} + +template +UINT FixGetSystemDirectory(char_t* buffer, UINT size) { + if constexpr(std::is_same_v) { + return GetSystemDirectoryA(buffer, size); + } else { + return GetSystemDirectoryW(buffer, size); + } +} + +template +UINT FixGetWindowsDirectory(char_t* buffer, UINT size) { + if constexpr(std::is_same_v) { + return GetWindowsDirectoryA(buffer, size); + } else { + return GetWindowsDirectoryW(buffer, size); + } +} +} \ No newline at end of file diff --git a/tests/unit/catter-hook/win/payload/main.cc b/tests/unit/catter-hook/win/payload/main.cc deleted file mode 100644 index 0ade447e..00000000 --- a/tests/unit/catter-hook/win/payload/main.cc +++ /dev/null @@ -1,2 +0,0 @@ - -// TODO diff --git a/tests/unit/catter-hook/win/payload/util.cc b/tests/unit/catter-hook/win/payload/util.cc new file mode 100644 index 00000000..9d66abd4 --- /dev/null +++ b/tests/unit/catter-hook/win/payload/util.cc @@ -0,0 +1,155 @@ + +#include "win/payload/util.h" + +#include + +#include +#include +#include + +#include + +namespace ct = catter; +namespace fs = std::filesystem; + +namespace { + +struct TempSandbox { + fs::path root; + + TempSandbox() { + root = fs::temp_directory_path() / + (L"catter_win_payload_ut_" + std::to_wstring(GetCurrentProcessId()) + L"_" + + std::to_wstring(GetTickCount64())); + fs::create_directories(root); + } + + ~TempSandbox() { + std::error_code ec; + fs::remove_all(root, ec); + } +}; + +void touch_file(const fs::path& file) { + fs::create_directories(file.parent_path()); + std::ofstream out(file, std::ios::binary); + out << "test"; +} + +bool same_existing_file(const fs::path& left, const fs::path& right) { + std::error_code ec; + if(fs::equivalent(left, right, ec)) { + return true; + } + return false; +} + +struct ScopedCurrentDirectory { + std::wstring original; + + ScopedCurrentDirectory(const fs::path& path) { + wchar_t buffer[MAX_PATH]; + auto len = GetCurrentDirectoryW(MAX_PATH, buffer); + if(len > 0) { + original.assign(buffer, len); + } + SetCurrentDirectoryW(path.wstring().c_str()); + } + + ~ScopedCurrentDirectory() { + if(!original.empty()) { + SetCurrentDirectoryW(original.c_str()); + } + } +}; + +struct ScopedEnvVar { + std::wstring name; + std::wstring original; + bool had_original = false; + + ScopedEnvVar(std::wstring name_value, const std::wstring& value) : name(std::move(name_value)) { + wchar_t buffer[32768]; + auto len = GetEnvironmentVariableW(name.c_str(), buffer, 32768); + if(len > 0 && len < 32768) { + had_original = true; + original.assign(buffer, len); + } + SetEnvironmentVariableW(name.c_str(), value.c_str()); + } + + ~ScopedEnvVar() { + if(had_original) { + SetEnvironmentVariableW(name.c_str(), original.c_str()); + } else { + SetEnvironmentVariableW(name.c_str(), nullptr); + } + } +}; + +TEST_SUITE(win_payload_util) { + TEST_CASE(resolve_abspath_from_application_name_appends_exe_and_searches_current_directory) { + TempSandbox sandbox; + ScopedCurrentDirectory scope(sandbox.root); + + auto app_path = sandbox.root / "clang.exe"; + touch_file(app_path); + + auto resolved = ct::win::payload::resolve_abspath("clang", "clang -c main.cc"); + EXPECT_TRUE(same_existing_file(fs::path(resolved), app_path)); + }; + + TEST_CASE(resolve_abspath_from_application_name_with_path_does_not_append_exe) { + TempSandbox sandbox; + ScopedCurrentDirectory scope(sandbox.root); + + auto app_path = sandbox.root / "bin" / "lld"; + touch_file(app_path); + + auto resolved = ct::win::payload::resolve_abspath("bin\\lld", "bin\\lld /v"); + EXPECT_TRUE(same_existing_file(fs::path(resolved), app_path)); + EXPECT_TRUE(fs::path(resolved).filename() == "lld"); + }; + + TEST_CASE(resolve_abspath_from_command_line_searches_path_variable) { + TempSandbox sandbox; + auto cwd = sandbox.root / "cwd"; + fs::create_directories(cwd); + ScopedCurrentDirectory scope(cwd); + + auto app_dir = sandbox.root / "pathbin"; + auto app_path = app_dir / "runner.exe"; + touch_file(app_path); + + ScopedEnvVar path_scope(L"PATH", app_dir.wstring()); + + auto resolved = ct::win::payload::resolve_abspath(nullptr, "runner --version"); + EXPECT_TRUE(same_existing_file(fs::path(resolved), app_path)); + }; + + TEST_CASE(resolve_abspath_from_command_line_supports_quoted_token) { + TempSandbox sandbox; + ScopedCurrentDirectory scope(sandbox.root); + + auto app_path = sandbox.root / "quoted.exe"; + touch_file(app_path); + + auto resolved = ct::win::payload::resolve_abspath(nullptr, "\"quoted\" --help"); + EXPECT_TRUE(same_existing_file(fs::path(resolved), app_path)); + }; + + TEST_CASE(get_proxy_path_reads_environment_variable) { + ScopedEnvVar scope(L"CATTER_PROXY_PATH", L"C:\\tmp\\proxy.exe"); + EXPECT_TRUE(ct::win::payload::get_proxy_path() == "C:\\tmp\\proxy.exe"); + EXPECT_TRUE(ct::win::payload::get_proxy_path_wide() == L"C:\\tmp\\proxy.exe"); + }; + + TEST_CASE(get_ipc_id_reads_environment_variable) { + ScopedEnvVar scope(L"CATTER_IPC_ID", L"12345"); + EXPECT_TRUE(ct::win::payload::get_ipc_id() == "12345"); + EXPECT_TRUE(ct::win::payload::get_ipc_id_wide() == L"12345"); + }; + +}; + +} // namespace diff --git a/xmake.lua b/xmake.lua index 27fd32cb..04009230 100644 --- a/xmake.lua +++ b/xmake.lua @@ -140,13 +140,16 @@ target("catter-hook-win64") set_default(is_plat("windows")) set_kind("shared") add_includedirs("src/catter-hook/") - add_files("src/catter-hook/win/payload/main.cc") + add_files("src/catter-hook/win/payload/*.cc") add_syslinks("user32", "advapi32") add_packages("minhook") - if get_config("toolchain") == "msvc" or get_config("toolchain") == "clang-cl" then + local toolchain = get_config("toolchain") + if toolchain == "msvc" or toolchain == "clang-cl" then add_cxxflags("/EHs-c-", "/GR-") + add_shflags("/DEF:src/catter-hook/win/payload/exports.def") else add_cxxflags("-fno-exceptions", "-fno-rtti") + add_shflags("-Wl,/DEF:src/catter-hook/win/payload/exports.def") end @@ -265,6 +268,8 @@ target("ut-catter-hook-win64") set_kind("binary") add_rules("ut-base") + add_includedirs("src/catter-hook/") + add_files("src/catter-hook/win/payload/util.cc") add_files("tests/unit/catter-hook/win/**.cc") add_deps("common") From db15ba4c1bdec6af38120a9520a9777aa6ffc7ed Mon Sep 17 00:00:00 2001 From: seele Date: Mon, 23 Mar 2026 17:03:30 +0800 Subject: [PATCH 2/4] Refactor catter-hook payload: Enhance resolver and utility functions - Updated `CreateProcessA` and `CreateProcessW` to use template functions for proxy path and IPC ID retrieval. - Refactored `Resolver` class to improve path resolution logic and added new methods for dynamic environment variable retrieval. - Introduced `CallWinApiWithGrowingBuffer` to handle Windows API calls with dynamic buffer sizes. - Created `resolver.cc` to implement the `Resolver` class methods. - Added unit tests for resolver functionality, including application name resolution and command line parsing. - Removed redundant code from `util.cc` and `util.h`, consolidating functionality into the `Resolver` class. --- src/catter-hook/win/payload/main.cc | 8 +- src/catter-hook/win/payload/resolver.cc | 99 ++++++ src/catter-hook/win/payload/resolver.h | 221 ++++++++------ src/catter-hook/win/payload/util.cc | 285 ++---------------- src/catter-hook/win/payload/util.h | 16 +- src/catter-hook/win/payload/winapi.h | 164 +++++++++- .../unit/catter-hook/win/payload/resolver.cc | 139 +++++++++ tests/unit/catter-hook/win/payload/util.cc | 167 ++-------- xmake.lua | 1 + 9 files changed, 603 insertions(+), 497 deletions(-) create mode 100644 src/catter-hook/win/payload/resolver.cc create mode 100644 tests/unit/catter-hook/win/payload/resolver.cc diff --git a/src/catter-hook/win/payload/main.cc b/src/catter-hook/win/payload/main.cc index 39f65ae1..c50c6aea 100644 --- a/src/catter-hook/win/payload/main.cc +++ b/src/catter-hook/win/payload/main.cc @@ -31,8 +31,8 @@ struct CreateProcessA { auto converted_cmdline = std::format("{} -p {} --exec {} -- {}", - catter::win::payload::get_proxy_path(), - catter::win::payload::get_ipc_id(), + catter::win::payload::get_proxy_path(), + catter::win::payload::get_ipc_id(), catter::win::payload::resolve_abspath(lpApplicationName, lpCommandLine), std::string_view(lpCommandLine ? lpCommandLine : "")); @@ -67,8 +67,8 @@ struct CreateProcessW { auto converted_cmdline = std::format(L"{} -p {} --exec {} -- {}", - catter::win::payload::get_proxy_path_wide(), - catter::win::payload::get_ipc_id_wide(), + catter::win::payload::get_proxy_path(), + catter::win::payload::get_ipc_id(), catter::win::payload::resolve_abspath(lpApplicationName, lpCommandLine), std::wstring_view(lpCommandLine ? lpCommandLine : L"")); diff --git a/src/catter-hook/win/payload/resolver.cc b/src/catter-hook/win/payload/resolver.cc new file mode 100644 index 00000000..a0689cba --- /dev/null +++ b/src/catter-hook/win/payload/resolver.cc @@ -0,0 +1,99 @@ +#include "resolver.h" + +namespace catter::win::payload { + +template +std::basic_string Resolver::resolve(std::basic_string_view app_name) const { + if(app_name.empty()) { + return {}; + } + + auto original_name = std::basic_string(app_name); + auto fixed_name = fix_app_name(app_name); + if(contains_path(fixed_name)) { + auto direct = to_absolute(fixed_name); + if(!direct.empty() && is_file(direct)) { + return direct; + } + // Keep original input when the path cannot be resolved to an existing file. + return original_name; + } + + for(const auto& search_path: m_search_paths) { + auto candidate = join_path(search_path, fixed_name); + auto absolute_candidate = to_absolute(candidate); + if(absolute_candidate.empty()) { + continue; + } + if(is_file(absolute_candidate)) { + return absolute_candidate; + } + } + + // Keep original input when search paths do not resolve an existing file. + return original_name; +} + +template +bool Resolver::contains_path(std::basic_string_view value) { + return std::find_if(value.begin(), value.end(), [](char_t c) { + return is_path_sep(c) || c == char_t(':'); + }) + != value.end(); +} + +template +bool Resolver::has_extension(std::basic_string_view file_name) { + return file_name.find(char_t('.')) != std::basic_string_view::npos; +} + +template +std::basic_string_view Resolver::exe_suffix() { + if constexpr(std::is_same_v) { + return ".exe"; + } else { + return L".exe"; + } +} + +template +std::basic_string Resolver::fix_app_name(std::basic_string_view app_name) { + if(contains_path(app_name) || app_name.back() == char_t('.') || has_extension(app_name)) { + return std::basic_string(app_name); + } + + std::basic_string fixed(app_name); + fixed.append(exe_suffix()); + return fixed; +} + +template +std::basic_string Resolver::join_path(std::basic_string_view directory, + std::basic_string_view file_name) { + if(directory.empty()) { + return std::basic_string(file_name); + } + + std::basic_string out(directory); + if(!is_path_sep(out.back())) { + out.push_back(path_sep); + } + out.append(file_name); + return out; +} + +template +std::basic_string Resolver::to_absolute(std::basic_string_view path) { + return GetFullPathNameDynamic(path); +} + +template +bool Resolver::is_file(std::basic_string_view path) { + auto attrs = FixGetFileAttributes(std::basic_string(path).c_str()); + return attrs != INVALID_FILE_ATTRIBUTES && (attrs & FILE_ATTRIBUTE_DIRECTORY) == 0; +} + +template class Resolver; +template class Resolver; + +} // namespace catter::win::payload diff --git a/src/catter-hook/win/payload/resolver.h b/src/catter-hook/win/payload/resolver.h index fc168684..85fd02e2 100644 --- a/src/catter-hook/win/payload/resolver.h +++ b/src/catter-hook/win/payload/resolver.h @@ -3,6 +3,8 @@ #include #include #include +#include +#include #include #include "winapi.h" @@ -12,120 +14,159 @@ namespace catter::win::payload { template class Resolver { public: - explicit Resolver(std::vector> search_paths) - : m_search_paths(std::move(search_paths)) {} + explicit Resolver(std::vector> search_paths) : + m_search_paths(std::move(search_paths)) {} - std::basic_string resolve(std::basic_string_view app_name) const { - if(app_name.empty()) { - return {}; - } - - auto original_name = std::basic_string(app_name); - auto fixed_name = fix_app_name(app_name); - if(contains_path(fixed_name)) { - auto direct = to_absolute(fixed_name); - if(!direct.empty() && is_file(direct)) { - return direct; - } - // Keep original input when the path cannot be resolved to an existing file. - return original_name; - } - - for(const auto& search_path: m_search_paths) { - auto candidate = join_path(search_path, fixed_name); - auto absolute_candidate = to_absolute(candidate); - if(absolute_candidate.empty()) { - continue; - } - if(is_file(absolute_candidate)) { - return absolute_candidate; - } - } - - // Keep original input when search paths do not resolve an existing file. - return original_name; - } + std::basic_string resolve(std::basic_string_view app_name) const; private: - static constexpr char_t path_sep = char_t('\\'); + constexpr static char_t path_sep = char_t('\\'); - static constexpr bool is_path_sep(char_t c) { + constexpr static bool is_path_sep(char_t c) { return c == char_t('\\') || c == char_t('/'); } - static bool contains_path(std::basic_string_view value) { - return std::find_if(value.begin(), value.end(), [](char_t c) { - return is_path_sep(c) || c == char_t(':'); - }) - != value.end(); - } + static bool contains_path(std::basic_string_view value); + static bool has_extension(std::basic_string_view file_name); + static std::basic_string_view exe_suffix(); + static std::basic_string fix_app_name(std::basic_string_view app_name); + static std::basic_string join_path(std::basic_string_view directory, + std::basic_string_view file_name); + static std::basic_string to_absolute(std::basic_string_view path); + static bool is_file(std::basic_string_view path); - static bool has_extension(std::basic_string_view file_name) { - return file_name.find(char_t('.')) != std::basic_string_view::npos; - } + std::vector> m_search_paths; +}; - static std::basic_string_view exe_suffix() { - if constexpr(std::is_same_v) { - return ".exe"; - } else { - return L".exe"; - } - } +template +Resolver create_app_name_resolver(); - static std::basic_string fix_app_name(std::basic_string_view app_name) { - if(contains_path(app_name) || app_name.back() == char_t('.') || has_extension(app_name)) { - return std::basic_string(app_name); - } +template +Resolver create_command_line_resolver(); - std::basic_string fixed(app_name); - fixed.append(exe_suffix()); - return fixed; - } +extern template class Resolver; +extern template class Resolver; - static std::basic_string join_path(std::basic_string_view directory, - std::basic_string_view file_name) { - if(directory.empty()) { - return std::basic_string(file_name); - } +namespace detail { - std::basic_string out(directory); - if(!is_path_sep(out.back())) { - out.push_back(path_sep); - } - out.append(file_name); - return out; - } +template +constexpr char_t k_path_delimiter = char_t(';'); - static std::basic_string to_absolute(std::basic_string_view path) { - std::basic_string input(path); - auto required_size = FixGetFullPathName(input.c_str(), 0, nullptr, nullptr); - if(required_size == 0) { - return {}; +template +constexpr char_t k_path_sep = char_t('\\'); + +template +std::vector> split_path_var(std::basic_string_view value) { + std::vector> segments; + size_t start = 0; + while(start <= value.size()) { + auto split_pos = value.find(k_path_delimiter, start); + auto token = + value.substr(start, + split_pos == std::basic_string_view::npos ? value.size() - start + : split_pos - start); + if(!token.empty()) { + segments.emplace_back(token); } - std::basic_string absolute(required_size, char_t('\0')); - auto written = - FixGetFullPathName(input.c_str(), required_size, absolute.data(), nullptr); - if(written == 0 || written >= required_size) { - return {}; + if(split_pos == std::basic_string_view::npos) { + break; } + start = split_pos + 1; + } + return segments; +} - absolute.resize(written); - return absolute; +template +void push_if_non_empty(std::vector>& out, + std::basic_string value) { + if(!value.empty()) { + out.push_back(std::move(value)); } +} - static bool is_file(std::basic_string_view path) { - auto attrs = FixGetFileAttributes(std::basic_string(path).c_str()); - return attrs != INVALID_FILE_ATTRIBUTES && (attrs & FILE_ATTRIBUTE_DIRECTORY) == 0; +template +std::basic_string get_current_drive_root() { + auto cwd = GetCurrentDirectoryDynamic(); + if(cwd.size() < 2 || cwd[1] != char_t(':')) { + return {}; } - std::vector> m_search_paths; -}; + std::basic_string root; + root.push_back(cwd[0]); + root.push_back(char_t(':')); + root.push_back(k_path_sep); + return root; +} template -Resolver create_app_name_resolver(); +constexpr std::basic_string_view k_path_env_name = {}; + +template <> +constexpr inline std::basic_string_view k_path_env_name = "PATH"; + +template <> +constexpr inline std::basic_string_view k_path_env_name = L"PATH"; template -Resolver create_command_line_resolver(); +constexpr std::basic_string_view k_system16_name = {}; + +template <> +constexpr inline std::basic_string_view k_system16_name = "System"; + +template <> +constexpr inline std::basic_string_view k_system16_name = L"System"; + +} // namespace detail + +template +Resolver create_app_name_resolver() { + std::vector> search_paths; + // Search order for explicit application_name: + // 1) current drive root, 2) current directory. + detail::push_if_non_empty(search_paths, detail::get_current_drive_root()); + detail::push_if_non_empty(search_paths, GetCurrentDirectoryDynamic()); + return Resolver(std::move(search_paths)); +} + +template +Resolver create_command_line_resolver() { + std::vector> search_paths; + + // Search order for command line token: + // 1) directory of current process image. + // 2) current directory. + // 3) 32-bit system directory. + // 4) 16-bit system directory named "System" under Windows directory. + // 5) Windows directory. + // 6) directories listed in PATH. + auto module_dir = GetModuleDirectory(nullptr); + auto current_dir = GetCurrentDirectoryDynamic(); + auto system_dir = GetSystemDirectoryDynamic(); + auto windows_dir = GetWindowsDirectoryDynamic(); + + detail::push_if_non_empty(search_paths, std::move(module_dir)); + detail::push_if_non_empty(search_paths, std::move(current_dir)); + detail::push_if_non_empty(search_paths, std::move(system_dir)); + if(!windows_dir.empty()) { + auto system16_dir = windows_dir; + if(system16_dir.back() != detail::k_path_sep && + system16_dir.back() != char_t('/')) { + system16_dir.push_back(detail::k_path_sep); + } + system16_dir.append(detail::k_system16_name); + detail::push_if_non_empty(search_paths, std::move(system16_dir)); + } + detail::push_if_non_empty(search_paths, windows_dir); + + auto path_value = + GetEnvironmentVariableDynamic(detail::k_path_env_name.data(), 4096); + auto path_segments = detail::split_path_var(path_value); + for(auto& segment: path_segments) { + detail::push_if_non_empty(search_paths, std::move(segment)); + } + + return Resolver(std::move(search_paths)); +} -} \ No newline at end of file +} // namespace catter::win::payload diff --git a/src/catter-hook/win/payload/util.cc b/src/catter-hook/win/payload/util.cc index f3895b10..2fc39aeb 100644 --- a/src/catter-hook/win/payload/util.cc +++ b/src/catter-hook/win/payload/util.cc @@ -1,204 +1,24 @@ #include #include #include -#include - -#include #include "win/env.h" #include "resolver.h" #include "util.h" - namespace catter::win::payload { namespace { -template -constexpr char_t k_path_delimiter = char_t(';'); - -template -constexpr char_t k_path_sep = char_t('\\'); - -template -std::vector> split_path_var(std::basic_string_view value) { - std::vector> segments; - size_t start = 0; - while(start <= value.size()) { - auto split_pos = value.find(k_path_delimiter, start); - auto token = value.substr(start, split_pos == std::basic_string_view::npos - ? value.size() - start - : split_pos - start); - if(!token.empty()) { - segments.emplace_back(token); - } - - if(split_pos == std::basic_string_view::npos) { - break; - } - start = split_pos + 1; - } - return segments; -} - -template -void push_if_non_empty(std::vector>& out, - std::basic_string value) { - if(!value.empty()) { - out.push_back(std::move(value)); - } -} - -template -std::basic_string get_env_var(const char_t* name, size_t stack_buffer_size) { - std::basic_string value; - value.resize(stack_buffer_size); - - auto len = FixGetEnvironmentVariable(name, value.data(), static_cast(value.size())); - if(len == 0) { - return {}; - } - - if(len < value.size()) { - value.resize(len); - return value; - } - - value.resize(len); - len = FixGetEnvironmentVariable(name, value.data(), static_cast(value.size())); - if(len == 0 || len > value.size()) { - return {}; - } - - if(len > 0 && value[len - 1] == char_t('\0')) { - value.resize(len - 1); - } else { - value.resize(len); - } - - return value; -} - -template -std::basic_string get_current_directory() { - auto required_size = FixGetCurrentDirectory(0, nullptr); - if(required_size == 0) { - return {}; - } - - std::basic_string value(required_size, char_t('\0')); - auto written = FixGetCurrentDirectory(required_size, value.data()); - if(written == 0 || written >= required_size) { - return {}; - } - - value.resize(written); - return value; -} - -template -std::basic_string get_module_directory() { - std::basic_string module_path(MAX_PATH, char_t('\0')); - while(true) { - auto written = FixGetModuleFileName(nullptr, - module_path.data(), - static_cast(module_path.size())); - if(written == 0) { - return {}; - } - if(written < module_path.size() - 1) { - module_path.resize(written); - break; - } - module_path.resize(module_path.size() * 2); - } - - auto pos = module_path.find_last_of(k_path_sep); - if(pos == std::basic_string::npos) { - return {}; - } - - return module_path.substr(0, pos); -} - -template -std::basic_string get_system_directory() { - std::basic_string value(MAX_PATH, char_t('\0')); - while(true) { - auto written = - FixGetSystemDirectory(value.data(), static_cast(value.size())); - if(written == 0) { - return {}; - } - if(written < value.size()) { - value.resize(written); - return value; - } - value.resize(written + 1); - } -} - -template -std::basic_string get_windows_directory() { - std::basic_string value(MAX_PATH, char_t('\0')); - while(true) { - auto written = - FixGetWindowsDirectory(value.data(), static_cast(value.size())); - if(written == 0) { - return {}; - } - if(written < value.size()) { - value.resize(written); - return value; - } - value.resize(written + 1); - } -} - -template -std::basic_string get_current_drive_root() { - auto cwd = get_current_directory(); - if(cwd.size() < 2 || cwd[1] != char_t(':')) { - return {}; - } - - std::basic_string root; - root.push_back(cwd[0]); - root.push_back(char_t(':')); - root.push_back(k_path_sep); - return root; -} - -template -constexpr std::basic_string_view k_path_env_name = {}; - -template <> -constexpr std::basic_string_view k_path_env_name = "PATH"; - -template <> -constexpr std::basic_string_view k_path_env_name = L"PATH"; - -template -constexpr std::basic_string_view k_system16_name = {}; - -template <> -constexpr std::basic_string_view k_system16_name = "System"; - -template <> -constexpr std::basic_string_view k_system16_name = L"System"; - -} // namespace - template std::basic_string extract_first_token(std::basic_string_view command_line) { constexpr char_t quote_char = char_t('"'); constexpr char_t space_char = char_t(' '); auto trimmed = [](std::basic_string_view value) { - auto first_non_space = std::find_if_not(value.begin(), value.end(), [](char_t c) { - return c == space_char; - }); + auto first_non_space = + std::find_if_not(value.begin(), value.end(), [](char_t c) { return c == space_char; }); return std::basic_string(first_non_space, value.end()); }(command_line); @@ -215,61 +35,12 @@ std::basic_string extract_first_token(std::basic_string_view com return std::basic_string(view.substr(1, closing_quote - 1)); } - auto first_space = std::find_if(view.begin(), view.end(), [](char_t c) { - return c == space_char; - }); + auto first_space = + std::find_if(view.begin(), view.end(), [](char_t c) { return c == space_char; }); return std::basic_string(view.begin(), first_space); } -template -Resolver create_app_name_resolver() { - std::vector> search_paths; - // Search order for explicit application_name: - // 1) current drive root, 2) current directory. - push_if_non_empty(search_paths, get_current_drive_root()); - push_if_non_empty(search_paths, get_current_directory()); - return Resolver(std::move(search_paths)); -} - -template -Resolver create_command_line_resolver() { - std::vector> search_paths; - - // Search order for command line token: - // 1) directory of current process image. - // 2) current directory. - // 3) 32-bit system directory. - // 4) 16-bit system directory named "System" under Windows directory. - // 5) Windows directory. - // 6) directories listed in PATH. - auto module_dir = get_module_directory(); - auto current_dir = get_current_directory(); - auto system_dir = get_system_directory(); - auto windows_dir = get_windows_directory(); - - push_if_non_empty(search_paths, std::move(module_dir)); - push_if_non_empty(search_paths, std::move(current_dir)); - push_if_non_empty(search_paths, std::move(system_dir)); - if(!windows_dir.empty()) { - auto system16_dir = windows_dir; - if(system16_dir.back() != k_path_sep && system16_dir.back() != char_t('/')) { - system16_dir.push_back(k_path_sep); - } - system16_dir.append(k_system16_name); - push_if_non_empty(search_paths, std::move(system16_dir)); - } - push_if_non_empty(search_paths, windows_dir); - - auto path_value = get_env_var(k_path_env_name.data(), 4096); - auto path_segments = split_path_var(path_value); - for(auto& segment: path_segments) { - push_if_non_empty(search_paths, std::move(segment)); - } - - return Resolver(std::move(search_paths)); -} - -template +template std::basic_string resolve_abspath_impl(const char_t* application_name, const char_t* command_line) { std::basic_string raw_app_name; @@ -291,35 +62,31 @@ std::basic_string resolve_abspath_impl(const char_t* application_name, return create_command_line_resolver().resolve(raw_app_name); } -template Resolver create_app_name_resolver(); -template Resolver create_app_name_resolver(); -template Resolver create_command_line_resolver(); -template Resolver create_command_line_resolver(); - - - -std::string resolve_abspath(const char* application_name, const char* command_line) { - return resolve_abspath_impl(application_name, command_line); -} - -std::wstring resolve_abspath(const wchar_t* application_name, const wchar_t* command_line) { - return resolve_abspath_impl(application_name, command_line); -} - -std::string get_proxy_path() { - return get_env_var(catter::win::ENV_VAR_PROXY_PATH, 256); -} +} // namespace -std::wstring get_proxy_path_wide() { - return get_env_var(catter::win::ENV_VAR_PROXY_PATH, 256); +template +std::basic_string resolve_abspath(const char_t* application_name, + const char_t* command_line) { + return resolve_abspath_impl(application_name, command_line); } -std::string get_ipc_id() { - return get_env_var(catter::win::ENV_VAR_IPC_ID, 64); +template +std::basic_string get_proxy_path() { + return GetEnvironmentVariableDynamic(catter::win::ENV_VAR_PROXY_PATH, 256); } -std::wstring get_ipc_id_wide() { - return get_env_var(catter::win::ENV_VAR_IPC_ID, 64); -} +template +std::basic_string get_ipc_id() { + return GetEnvironmentVariableDynamic(catter::win::ENV_VAR_IPC_ID, 64); +} + +template std::basic_string resolve_abspath(const char* application_name, + const char* command_line); +template std::basic_string resolve_abspath(const wchar_t* application_name, + const wchar_t* command_line); +template std::basic_string get_proxy_path(); +template std::basic_string get_proxy_path(); +template std::basic_string get_ipc_id(); +template std::basic_string get_ipc_id(); } // namespace catter::win::payload diff --git a/src/catter-hook/win/payload/util.h b/src/catter-hook/win/payload/util.h index 748e16d0..01d8b78f 100644 --- a/src/catter-hook/win/payload/util.h +++ b/src/catter-hook/win/payload/util.h @@ -3,15 +3,17 @@ #include #include -namespace catter::win::payload { +#include "winapi.h" -std::string resolve_abspath(const char* application_name, const char* command_line); -std::wstring resolve_abspath(const wchar_t* application_name, const wchar_t* command_line); +namespace catter::win::payload { -std::string get_proxy_path(); -std::wstring get_proxy_path_wide(); +template +std::basic_string resolve_abspath(const char_t* application_name, + const char_t* command_line); +template +std::basic_string get_proxy_path(); -std::string get_ipc_id(); -std::wstring get_ipc_id_wide(); +template +std::basic_string get_ipc_id(); } // namespace catter::win::payload diff --git a/src/catter-hook/win/payload/winapi.h b/src/catter-hook/win/payload/winapi.h index e44cdcbf..b601cc86 100644 --- a/src/catter-hook/win/payload/winapi.h +++ b/src/catter-hook/win/payload/winapi.h @@ -1,13 +1,63 @@ +#pragma once + #include -#include +#include +#include +#include +#ifndef NOMINMAX +#define NOMINMAX +#endif + +#include namespace catter::win::payload { template concept CharT = std::is_same_v || std::is_same_v; +struct WinApiBufferDecision { + bool done = false; + size_t output_size = 0; + size_t next_size = 0; +}; + +template +std::basic_string CallWinApiWithGrowingBuffer(size_t initial_size, + api_call_t&& call, + decision_t&& decide) { + constexpr auto max_size = static_cast((std::numeric_limits::max)()); + size_t buffer_size = initial_size < 1 ? 1 : initial_size; + if(buffer_size > max_size) { + return {}; + } + + std::basic_string buffer(buffer_size, char_t('\0')); + while(true) { + auto result = call(buffer.data(), static_cast(buffer.size())); + if(result == 0) { + return {}; + } + + auto decision = decide(result, buffer.size()); + if(decision.done) { + if(decision.output_size > buffer.size()) { + return {}; + } + buffer.resize(decision.output_size); + return buffer; + } + + auto next_size = + decision.next_size > (buffer.size() + 1) ? decision.next_size : (buffer.size() + 1); + if(next_size > max_size) { + return {}; + } + buffer.resize(next_size); + } +} + template DWORD FixGetEnvironmentVariable(const char_t* name, char_t* buffer, DWORD size) { if constexpr(std::is_same_v) { @@ -73,4 +123,114 @@ UINT FixGetWindowsDirectory(char_t* buffer, UINT size) { return GetWindowsDirectoryW(buffer, size); } } -} \ No newline at end of file + +template +std::basic_string GetEnvironmentVariableDynamic(const char_t* name, + size_t initial_size = 256) { + return CallWinApiWithGrowingBuffer( + initial_size, + [&](char_t* buffer, DWORD size) { + return FixGetEnvironmentVariable(name, buffer, size); + }, + [](DWORD result, size_t buffer_size) -> WinApiBufferDecision { + if(static_cast(result) < buffer_size) { + return WinApiBufferDecision{.done = true, + .output_size = static_cast(result)}; + } + return WinApiBufferDecision{.next_size = static_cast(result)}; + }); +} + +template +std::basic_string GetCurrentDirectoryDynamic(size_t initial_size = MAX_PATH) { + return CallWinApiWithGrowingBuffer( + initial_size, + [&](char_t* buffer, DWORD size) { return FixGetCurrentDirectory(size, buffer); }, + [](DWORD result, size_t buffer_size) -> WinApiBufferDecision { + if(static_cast(result) < buffer_size) { + return WinApiBufferDecision{.done = true, + .output_size = static_cast(result)}; + } + return WinApiBufferDecision{.next_size = static_cast(result) + 1}; + }); +} + +template +std::basic_string GetModulePathDynamic(HMODULE module, size_t initial_size = MAX_PATH) { + return CallWinApiWithGrowingBuffer( + initial_size, + [&](char_t* buffer, DWORD size) { + return FixGetModuleFileName(module, buffer, size); + }, + [](DWORD result, size_t buffer_size) -> WinApiBufferDecision { + auto written = static_cast(result); + if(written < buffer_size) { + return WinApiBufferDecision{.done = true, .output_size = written}; + } + return WinApiBufferDecision{.next_size = buffer_size * 2}; + }); +} + +template +std::basic_string GetModuleDirectory(HMODULE module, size_t initial_size = MAX_PATH) { + auto module_path = GetModulePathDynamic(module, initial_size); + if(module_path.empty()) { + return {}; + } + + auto pos = module_path.find_last_of(char_t('\\')); + if(pos == std::basic_string::npos) { + pos = module_path.find_last_of(char_t('/')); + } + if(pos == std::basic_string::npos) { + return {}; + } + return module_path.substr(0, pos); +} + +template +std::basic_string GetSystemDirectoryDynamic(size_t initial_size = MAX_PATH) { + return CallWinApiWithGrowingBuffer( + initial_size, + [&](char_t* buffer, UINT size) { return FixGetSystemDirectory(buffer, size); }, + [](UINT result, size_t buffer_size) -> WinApiBufferDecision { + if(static_cast(result) < buffer_size) { + return WinApiBufferDecision{.done = true, + .output_size = static_cast(result)}; + } + return WinApiBufferDecision{.next_size = static_cast(result) + 1}; + }); +} + +template +std::basic_string GetWindowsDirectoryDynamic(size_t initial_size = MAX_PATH) { + return CallWinApiWithGrowingBuffer( + initial_size, + [&](char_t* buffer, UINT size) { return FixGetWindowsDirectory(buffer, size); }, + [](UINT result, size_t buffer_size) -> WinApiBufferDecision { + if(static_cast(result) < buffer_size) { + return WinApiBufferDecision{.done = true, + .output_size = static_cast(result)}; + } + return WinApiBufferDecision{.next_size = static_cast(result) + 1}; + }); +} + +template +std::basic_string GetFullPathNameDynamic(std::basic_string_view path, + size_t initial_size = MAX_PATH) { + auto input = std::basic_string(path); + return CallWinApiWithGrowingBuffer( + initial_size, + [&](char_t* buffer, DWORD size) { + return FixGetFullPathName(input.c_str(), size, buffer, nullptr); + }, + [](DWORD result, size_t buffer_size) -> WinApiBufferDecision { + if(static_cast(result) < buffer_size) { + return WinApiBufferDecision{.done = true, + .output_size = static_cast(result)}; + } + return WinApiBufferDecision{.next_size = static_cast(result) + 1}; + }); +} +} // namespace catter::win::payload diff --git a/tests/unit/catter-hook/win/payload/resolver.cc b/tests/unit/catter-hook/win/payload/resolver.cc new file mode 100644 index 00000000..016ae69a --- /dev/null +++ b/tests/unit/catter-hook/win/payload/resolver.cc @@ -0,0 +1,139 @@ +#include "win/payload/resolver.h" +#include "win/payload/util.h" + +#include + +#include +#include +#include + +#include + +namespace ct = catter; +namespace fs = std::filesystem; + +namespace { + +struct TempSandbox { + fs::path root; + + TempSandbox() { + root = fs::temp_directory_path() / + (L"catter_win_payload_resolver_ut_" + std::to_wstring(GetCurrentProcessId()) + L"_" + + std::to_wstring(GetTickCount64())); + fs::create_directories(root); + } + + ~TempSandbox() { + std::error_code ec; + fs::remove_all(root, ec); + } +}; + +void touch_file(const fs::path& file) { + fs::create_directories(file.parent_path()); + std::ofstream out(file, std::ios::binary); + out << "test"; +} + +bool same_existing_file(const fs::path& left, const fs::path& right) { + std::error_code ec; + return fs::equivalent(left, right, ec); +} + +struct ScopedCurrentDirectory { + std::wstring original; + + explicit ScopedCurrentDirectory(const fs::path& path) { + wchar_t buffer[MAX_PATH]; + auto len = GetCurrentDirectoryW(MAX_PATH, buffer); + if(len > 0) { + original.assign(buffer, len); + } + SetCurrentDirectoryW(path.wstring().c_str()); + } + + ~ScopedCurrentDirectory() { + if(!original.empty()) { + SetCurrentDirectoryW(original.c_str()); + } + } +}; + +struct ScopedEnvVar { + std::wstring name; + std::wstring original; + bool had_original = false; + + ScopedEnvVar(std::wstring name_value, const std::wstring& value) : name(std::move(name_value)) { + wchar_t buffer[32768]; + auto len = GetEnvironmentVariableW(name.c_str(), buffer, 32768); + if(len > 0 && len < 32768) { + had_original = true; + original.assign(buffer, len); + } + SetEnvironmentVariableW(name.c_str(), value.c_str()); + } + + ~ScopedEnvVar() { + if(had_original) { + SetEnvironmentVariableW(name.c_str(), original.c_str()); + } else { + SetEnvironmentVariableW(name.c_str(), nullptr); + } + } +}; + +TEST_SUITE(win_payload_resolver) { + TEST_CASE(app_name_resolver_appends_exe_and_searches_current_directory) { + TempSandbox sandbox; + ScopedCurrentDirectory scope(sandbox.root); + + auto app_path = sandbox.root / "clang.exe"; + touch_file(app_path); + + auto resolved = ct::win::payload::create_app_name_resolver().resolve("clang"); + EXPECT_TRUE(same_existing_file(fs::path(resolved), app_path)); + }; + + TEST_CASE(app_name_resolver_does_not_append_exe_when_input_has_path) { + TempSandbox sandbox; + ScopedCurrentDirectory scope(sandbox.root); + + auto app_path = sandbox.root / "bin" / "lld"; + touch_file(app_path); + + auto resolved = ct::win::payload::create_app_name_resolver().resolve("bin\\lld"); + EXPECT_TRUE(same_existing_file(fs::path(resolved), app_path)); + EXPECT_TRUE(fs::path(resolved).filename() == "lld"); + }; + + TEST_CASE(command_line_resolver_searches_path_variable) { + TempSandbox sandbox; + auto cwd = sandbox.root / "cwd"; + fs::create_directories(cwd); + ScopedCurrentDirectory scope(cwd); + + auto app_dir = sandbox.root / "pathbin"; + auto app_path = app_dir / "runner.exe"; + touch_file(app_path); + + ScopedEnvVar path_scope(L"PATH", app_dir.wstring()); + + auto resolved = ct::win::payload::create_command_line_resolver().resolve("runner"); + EXPECT_TRUE(same_existing_file(fs::path(resolved), app_path)); + }; + + TEST_CASE(resolve_abspath_supports_quoted_command_line) { + TempSandbox sandbox; + ScopedCurrentDirectory scope(sandbox.root); + + auto app_path = sandbox.root / "quoted.exe"; + touch_file(app_path); + + auto resolved = ct::win::payload::resolve_abspath(nullptr, "\"quoted\" --help"); + EXPECT_TRUE(same_existing_file(fs::path(resolved), app_path)); + }; +}; + +} // namespace diff --git a/tests/unit/catter-hook/win/payload/util.cc b/tests/unit/catter-hook/win/payload/util.cc index 9d66abd4..1b7a83ae 100644 --- a/tests/unit/catter-hook/win/payload/util.cc +++ b/tests/unit/catter-hook/win/payload/util.cc @@ -3,153 +3,50 @@ #include -#include -#include #include #include namespace ct = catter; -namespace fs = std::filesystem; namespace { -struct TempSandbox { - fs::path root; - - TempSandbox() { - root = fs::temp_directory_path() / - (L"catter_win_payload_ut_" + std::to_wstring(GetCurrentProcessId()) + L"_" - + std::to_wstring(GetTickCount64())); - fs::create_directories(root); - } - - ~TempSandbox() { - std::error_code ec; - fs::remove_all(root, ec); - } -}; - -void touch_file(const fs::path& file) { - fs::create_directories(file.parent_path()); - std::ofstream out(file, std::ios::binary); - out << "test"; -} - -bool same_existing_file(const fs::path& left, const fs::path& right) { - std::error_code ec; - if(fs::equivalent(left, right, ec)) { - return true; - } - return false; -} - -struct ScopedCurrentDirectory { - std::wstring original; - - ScopedCurrentDirectory(const fs::path& path) { - wchar_t buffer[MAX_PATH]; - auto len = GetCurrentDirectoryW(MAX_PATH, buffer); - if(len > 0) { - original.assign(buffer, len); - } - SetCurrentDirectoryW(path.wstring().c_str()); - } - - ~ScopedCurrentDirectory() { - if(!original.empty()) { - SetCurrentDirectoryW(original.c_str()); - } - } -}; - struct ScopedEnvVar { - std::wstring name; - std::wstring original; - bool had_original = false; - - ScopedEnvVar(std::wstring name_value, const std::wstring& value) : name(std::move(name_value)) { - wchar_t buffer[32768]; - auto len = GetEnvironmentVariableW(name.c_str(), buffer, 32768); - if(len > 0 && len < 32768) { - had_original = true; - original.assign(buffer, len); - } - SetEnvironmentVariableW(name.c_str(), value.c_str()); - } - - ~ScopedEnvVar() { - if(had_original) { - SetEnvironmentVariableW(name.c_str(), original.c_str()); - } else { - SetEnvironmentVariableW(name.c_str(), nullptr); - } - } + std::wstring name; + std::wstring original; + bool had_original = false; + + ScopedEnvVar(std::wstring name_value, const std::wstring& value) : name(std::move(name_value)) { + wchar_t buffer[32768]; + auto len = GetEnvironmentVariableW(name.c_str(), buffer, 32768); + if(len > 0 && len < 32768) { + had_original = true; + original.assign(buffer, len); + } + SetEnvironmentVariableW(name.c_str(), value.c_str()); + } + + ~ScopedEnvVar() { + if(had_original) { + SetEnvironmentVariableW(name.c_str(), original.c_str()); + } else { + SetEnvironmentVariableW(name.c_str(), nullptr); + } + } }; TEST_SUITE(win_payload_util) { - TEST_CASE(resolve_abspath_from_application_name_appends_exe_and_searches_current_directory) { - TempSandbox sandbox; - ScopedCurrentDirectory scope(sandbox.root); - - auto app_path = sandbox.root / "clang.exe"; - touch_file(app_path); - - auto resolved = ct::win::payload::resolve_abspath("clang", "clang -c main.cc"); - EXPECT_TRUE(same_existing_file(fs::path(resolved), app_path)); - }; - - TEST_CASE(resolve_abspath_from_application_name_with_path_does_not_append_exe) { - TempSandbox sandbox; - ScopedCurrentDirectory scope(sandbox.root); - - auto app_path = sandbox.root / "bin" / "lld"; - touch_file(app_path); - - auto resolved = ct::win::payload::resolve_abspath("bin\\lld", "bin\\lld /v"); - EXPECT_TRUE(same_existing_file(fs::path(resolved), app_path)); - EXPECT_TRUE(fs::path(resolved).filename() == "lld"); - }; - - TEST_CASE(resolve_abspath_from_command_line_searches_path_variable) { - TempSandbox sandbox; - auto cwd = sandbox.root / "cwd"; - fs::create_directories(cwd); - ScopedCurrentDirectory scope(cwd); - - auto app_dir = sandbox.root / "pathbin"; - auto app_path = app_dir / "runner.exe"; - touch_file(app_path); - - ScopedEnvVar path_scope(L"PATH", app_dir.wstring()); - - auto resolved = ct::win::payload::resolve_abspath(nullptr, "runner --version"); - EXPECT_TRUE(same_existing_file(fs::path(resolved), app_path)); - }; - - TEST_CASE(resolve_abspath_from_command_line_supports_quoted_token) { - TempSandbox sandbox; - ScopedCurrentDirectory scope(sandbox.root); - - auto app_path = sandbox.root / "quoted.exe"; - touch_file(app_path); - - auto resolved = ct::win::payload::resolve_abspath(nullptr, "\"quoted\" --help"); - EXPECT_TRUE(same_existing_file(fs::path(resolved), app_path)); - }; - - TEST_CASE(get_proxy_path_reads_environment_variable) { - ScopedEnvVar scope(L"CATTER_PROXY_PATH", L"C:\\tmp\\proxy.exe"); - EXPECT_TRUE(ct::win::payload::get_proxy_path() == "C:\\tmp\\proxy.exe"); - EXPECT_TRUE(ct::win::payload::get_proxy_path_wide() == L"C:\\tmp\\proxy.exe"); - }; - - TEST_CASE(get_ipc_id_reads_environment_variable) { - ScopedEnvVar scope(L"CATTER_IPC_ID", L"12345"); - EXPECT_TRUE(ct::win::payload::get_ipc_id() == "12345"); - EXPECT_TRUE(ct::win::payload::get_ipc_id_wide() == L"12345"); - }; - + TEST_CASE(get_proxy_path_reads_environment_variable) { + ScopedEnvVar scope(L"CATTER_PROXY_PATH", L"C:\\tmp\\proxy.exe"); + EXPECT_TRUE(ct::win::payload::get_proxy_path() == "C:\\tmp\\proxy.exe"); + EXPECT_TRUE(ct::win::payload::get_proxy_path() == L"C:\\tmp\\proxy.exe"); + }; + + TEST_CASE(get_ipc_id_reads_environment_variable) { + ScopedEnvVar scope(L"CATTER_IPC_ID", L"12345"); + EXPECT_TRUE(ct::win::payload::get_ipc_id() == "12345"); + EXPECT_TRUE(ct::win::payload::get_ipc_id() == L"12345"); + }; }; } // namespace diff --git a/xmake.lua b/xmake.lua index 04009230..83afe05d 100644 --- a/xmake.lua +++ b/xmake.lua @@ -269,6 +269,7 @@ target("ut-catter-hook-win64") add_rules("ut-base") add_includedirs("src/catter-hook/") + add_files("src/catter-hook/win/payload/resolver.cc") add_files("src/catter-hook/win/payload/util.cc") add_files("tests/unit/catter-hook/win/**.cc") From d34ac630418f9a21e40d107e8d09d5e59c16c6e3 Mon Sep 17 00:00:00 2001 From: seele Date: Mon, 23 Mar 2026 17:13:09 +0800 Subject: [PATCH 3/4] refactor: streamline resolver functions and improve code organization --- src/catter-hook/win/payload/resolver.cc | 131 ++++++++++++- src/catter-hook/win/payload/resolver.h | 120 +----------- .../unit/catter-hook/win/payload/resolver.cc | 184 +++++++++--------- 3 files changed, 223 insertions(+), 212 deletions(-) diff --git a/src/catter-hook/win/payload/resolver.cc b/src/catter-hook/win/payload/resolver.cc index a0689cba..92180ce3 100644 --- a/src/catter-hook/win/payload/resolver.cc +++ b/src/catter-hook/win/payload/resolver.cc @@ -38,8 +38,7 @@ template bool Resolver::contains_path(std::basic_string_view value) { return std::find_if(value.begin(), value.end(), [](char_t c) { return is_path_sep(c) || c == char_t(':'); - }) - != value.end(); + }) != value.end(); } template @@ -93,7 +92,135 @@ bool Resolver::is_file(std::basic_string_view path) { return attrs != INVALID_FILE_ATTRIBUTES && (attrs & FILE_ATTRIBUTE_DIRECTORY) == 0; } +namespace detail { + +template +constexpr char_t k_path_delimiter = char_t(';'); + +template +constexpr char_t k_path_sep = char_t('\\'); + +template +std::vector> split_path_var(std::basic_string_view value) { + std::vector> segments; + size_t start = 0; + while(start <= value.size()) { + auto split_pos = value.find(k_path_delimiter, start); + auto token = + value.substr(start, + split_pos == std::basic_string_view::npos ? value.size() - start + : split_pos - start); + if(!token.empty()) { + segments.emplace_back(token); + } + + if(split_pos == std::basic_string_view::npos) { + break; + } + start = split_pos + 1; + } + return segments; +} + +template +void push_if_non_empty(std::vector>& out, + std::basic_string value) { + if(!value.empty()) { + out.push_back(std::move(value)); + } +} + +template +std::basic_string get_current_drive_root() { + auto cwd = GetCurrentDirectoryDynamic(); + if(cwd.size() < 2 || cwd[1] != char_t(':')) { + return {}; + } + + std::basic_string root; + root.push_back(cwd[0]); + root.push_back(char_t(':')); + root.push_back(k_path_sep); + return root; +} + +template +constexpr std::basic_string_view k_path_env_name = {}; + +template <> +constexpr std::basic_string_view k_path_env_name = "PATH"; + +template <> +constexpr std::basic_string_view k_path_env_name = L"PATH"; + +template +constexpr std::basic_string_view k_system16_name = {}; + +template <> +constexpr std::basic_string_view k_system16_name = "System"; + +template <> +constexpr std::basic_string_view k_system16_name = L"System"; + +} // namespace detail + +template +Resolver create_app_name_resolver() { + std::vector> search_paths; + search_paths.reserve(2); + // Search order for explicit application_name: + // 1) current drive root, 2) current directory. + detail::push_if_non_empty(search_paths, detail::get_current_drive_root()); + detail::push_if_non_empty(search_paths, GetCurrentDirectoryDynamic()); + return Resolver(std::move(search_paths)); +} + +template +Resolver create_command_line_resolver() { + std::vector> search_paths; + search_paths.reserve(64); + // Search order for command line token: + // 1) directory of current process image. + // 2) current directory. + // 3) 32-bit system directory. + // 4) 16-bit system directory named "System" under Windows directory. + // 5) Windows directory. + // 6) directories listed in PATH. + auto module_dir = GetModuleDirectory(nullptr); + auto current_dir = GetCurrentDirectoryDynamic(); + auto system_dir = GetSystemDirectoryDynamic(); + auto windows_dir = GetWindowsDirectoryDynamic(); + + detail::push_if_non_empty(search_paths, std::move(module_dir)); + detail::push_if_non_empty(search_paths, std::move(current_dir)); + detail::push_if_non_empty(search_paths, std::move(system_dir)); + if(!windows_dir.empty()) { + auto system16_dir = windows_dir; + if(system16_dir.back() != detail::k_path_sep && + system16_dir.back() != char_t('/')) { + system16_dir.push_back(detail::k_path_sep); + } + system16_dir.append(detail::k_system16_name); + detail::push_if_non_empty(search_paths, std::move(system16_dir)); + } + detail::push_if_non_empty(search_paths, windows_dir); + + auto path_value = + GetEnvironmentVariableDynamic(detail::k_path_env_name.data(), 4096); + auto path_segments = detail::split_path_var(path_value); + for(auto& segment: path_segments) { + detail::push_if_non_empty(search_paths, std::move(segment)); + } + + return Resolver(std::move(search_paths)); +} + template class Resolver; template class Resolver; +template Resolver create_app_name_resolver(); +template Resolver create_app_name_resolver(); +template Resolver create_command_line_resolver(); +template Resolver create_command_line_resolver(); + } // namespace catter::win::payload diff --git a/src/catter-hook/win/payload/resolver.h b/src/catter-hook/win/payload/resolver.h index 85fd02e2..9983e6ce 100644 --- a/src/catter-hook/win/payload/resolver.h +++ b/src/catter-hook/win/payload/resolver.h @@ -47,126 +47,10 @@ Resolver create_command_line_resolver(); extern template class Resolver; extern template class Resolver; -namespace detail { - -template -constexpr char_t k_path_delimiter = char_t(';'); - -template -constexpr char_t k_path_sep = char_t('\\'); - -template -std::vector> split_path_var(std::basic_string_view value) { - std::vector> segments; - size_t start = 0; - while(start <= value.size()) { - auto split_pos = value.find(k_path_delimiter, start); - auto token = - value.substr(start, - split_pos == std::basic_string_view::npos ? value.size() - start - : split_pos - start); - if(!token.empty()) { - segments.emplace_back(token); - } - - if(split_pos == std::basic_string_view::npos) { - break; - } - start = split_pos + 1; - } - return segments; -} - template -void push_if_non_empty(std::vector>& out, - std::basic_string value) { - if(!value.empty()) { - out.push_back(std::move(value)); - } -} - -template -std::basic_string get_current_drive_root() { - auto cwd = GetCurrentDirectoryDynamic(); - if(cwd.size() < 2 || cwd[1] != char_t(':')) { - return {}; - } - - std::basic_string root; - root.push_back(cwd[0]); - root.push_back(char_t(':')); - root.push_back(k_path_sep); - return root; -} - -template -constexpr std::basic_string_view k_path_env_name = {}; - -template <> -constexpr inline std::basic_string_view k_path_env_name = "PATH"; - -template <> -constexpr inline std::basic_string_view k_path_env_name = L"PATH"; - -template -constexpr std::basic_string_view k_system16_name = {}; - -template <> -constexpr inline std::basic_string_view k_system16_name = "System"; - -template <> -constexpr inline std::basic_string_view k_system16_name = L"System"; - -} // namespace detail - -template -Resolver create_app_name_resolver() { - std::vector> search_paths; - // Search order for explicit application_name: - // 1) current drive root, 2) current directory. - detail::push_if_non_empty(search_paths, detail::get_current_drive_root()); - detail::push_if_non_empty(search_paths, GetCurrentDirectoryDynamic()); - return Resolver(std::move(search_paths)); -} +Resolver create_app_name_resolver(); template -Resolver create_command_line_resolver() { - std::vector> search_paths; - - // Search order for command line token: - // 1) directory of current process image. - // 2) current directory. - // 3) 32-bit system directory. - // 4) 16-bit system directory named "System" under Windows directory. - // 5) Windows directory. - // 6) directories listed in PATH. - auto module_dir = GetModuleDirectory(nullptr); - auto current_dir = GetCurrentDirectoryDynamic(); - auto system_dir = GetSystemDirectoryDynamic(); - auto windows_dir = GetWindowsDirectoryDynamic(); - - detail::push_if_non_empty(search_paths, std::move(module_dir)); - detail::push_if_non_empty(search_paths, std::move(current_dir)); - detail::push_if_non_empty(search_paths, std::move(system_dir)); - if(!windows_dir.empty()) { - auto system16_dir = windows_dir; - if(system16_dir.back() != detail::k_path_sep && - system16_dir.back() != char_t('/')) { - system16_dir.push_back(detail::k_path_sep); - } - system16_dir.append(detail::k_system16_name); - detail::push_if_non_empty(search_paths, std::move(system16_dir)); - } - detail::push_if_non_empty(search_paths, windows_dir); - - auto path_value = - GetEnvironmentVariableDynamic(detail::k_path_env_name.data(), 4096); - auto path_segments = detail::split_path_var(path_value); - for(auto& segment: path_segments) { - detail::push_if_non_empty(search_paths, std::move(segment)); - } - - return Resolver(std::move(search_paths)); -} +Resolver create_command_line_resolver(); } // namespace catter::win::payload diff --git a/tests/unit/catter-hook/win/payload/resolver.cc b/tests/unit/catter-hook/win/payload/resolver.cc index 016ae69a..5d1c9bec 100644 --- a/tests/unit/catter-hook/win/payload/resolver.cc +++ b/tests/unit/catter-hook/win/payload/resolver.cc @@ -15,125 +15,125 @@ namespace fs = std::filesystem; namespace { struct TempSandbox { - fs::path root; - - TempSandbox() { - root = fs::temp_directory_path() / - (L"catter_win_payload_resolver_ut_" + std::to_wstring(GetCurrentProcessId()) + L"_" - + std::to_wstring(GetTickCount64())); - fs::create_directories(root); - } - - ~TempSandbox() { - std::error_code ec; - fs::remove_all(root, ec); - } + fs::path root; + + TempSandbox() { + root = fs::temp_directory_path() / + (L"catter_win_payload_resolver_ut_" + std::to_wstring(GetCurrentProcessId()) + L"_" + + std::to_wstring(GetTickCount64())); + fs::create_directories(root); + } + + ~TempSandbox() { + std::error_code ec; + fs::remove_all(root, ec); + } }; void touch_file(const fs::path& file) { - fs::create_directories(file.parent_path()); - std::ofstream out(file, std::ios::binary); - out << "test"; + fs::create_directories(file.parent_path()); + std::ofstream out(file, std::ios::binary); + out << "test"; } bool same_existing_file(const fs::path& left, const fs::path& right) { - std::error_code ec; - return fs::equivalent(left, right, ec); + std::error_code ec; + return fs::equivalent(left, right, ec); } struct ScopedCurrentDirectory { - std::wstring original; - - explicit ScopedCurrentDirectory(const fs::path& path) { - wchar_t buffer[MAX_PATH]; - auto len = GetCurrentDirectoryW(MAX_PATH, buffer); - if(len > 0) { - original.assign(buffer, len); - } - SetCurrentDirectoryW(path.wstring().c_str()); - } - - ~ScopedCurrentDirectory() { - if(!original.empty()) { - SetCurrentDirectoryW(original.c_str()); - } - } + std::wstring original; + + explicit ScopedCurrentDirectory(const fs::path& path) { + wchar_t buffer[MAX_PATH]; + auto len = GetCurrentDirectoryW(MAX_PATH, buffer); + if(len > 0) { + original.assign(buffer, len); + } + SetCurrentDirectoryW(path.wstring().c_str()); + } + + ~ScopedCurrentDirectory() { + if(!original.empty()) { + SetCurrentDirectoryW(original.c_str()); + } + } }; struct ScopedEnvVar { - std::wstring name; - std::wstring original; - bool had_original = false; - - ScopedEnvVar(std::wstring name_value, const std::wstring& value) : name(std::move(name_value)) { - wchar_t buffer[32768]; - auto len = GetEnvironmentVariableW(name.c_str(), buffer, 32768); - if(len > 0 && len < 32768) { - had_original = true; - original.assign(buffer, len); - } - SetEnvironmentVariableW(name.c_str(), value.c_str()); - } - - ~ScopedEnvVar() { - if(had_original) { - SetEnvironmentVariableW(name.c_str(), original.c_str()); - } else { - SetEnvironmentVariableW(name.c_str(), nullptr); - } - } + std::wstring name; + std::wstring original; + bool had_original = false; + + ScopedEnvVar(std::wstring name_value, const std::wstring& value) : name(std::move(name_value)) { + wchar_t buffer[32768]; + auto len = GetEnvironmentVariableW(name.c_str(), buffer, 32768); + if(len > 0 && len < 32768) { + had_original = true; + original.assign(buffer, len); + } + SetEnvironmentVariableW(name.c_str(), value.c_str()); + } + + ~ScopedEnvVar() { + if(had_original) { + SetEnvironmentVariableW(name.c_str(), original.c_str()); + } else { + SetEnvironmentVariableW(name.c_str(), nullptr); + } + } }; TEST_SUITE(win_payload_resolver) { - TEST_CASE(app_name_resolver_appends_exe_and_searches_current_directory) { - TempSandbox sandbox; - ScopedCurrentDirectory scope(sandbox.root); + TEST_CASE(app_name_resolver_appends_exe_and_searches_current_directory) { + TempSandbox sandbox; + ScopedCurrentDirectory scope(sandbox.root); - auto app_path = sandbox.root / "clang.exe"; - touch_file(app_path); + auto app_path = sandbox.root / "clang.exe"; + touch_file(app_path); - auto resolved = ct::win::payload::create_app_name_resolver().resolve("clang"); - EXPECT_TRUE(same_existing_file(fs::path(resolved), app_path)); - }; + auto resolved = ct::win::payload::create_app_name_resolver().resolve("clang"); + EXPECT_TRUE(same_existing_file(fs::path(resolved), app_path)); + }; - TEST_CASE(app_name_resolver_does_not_append_exe_when_input_has_path) { - TempSandbox sandbox; - ScopedCurrentDirectory scope(sandbox.root); + TEST_CASE(app_name_resolver_does_not_append_exe_when_input_has_path) { + TempSandbox sandbox; + ScopedCurrentDirectory scope(sandbox.root); - auto app_path = sandbox.root / "bin" / "lld"; - touch_file(app_path); + auto app_path = sandbox.root / "bin" / "lld"; + touch_file(app_path); - auto resolved = ct::win::payload::create_app_name_resolver().resolve("bin\\lld"); - EXPECT_TRUE(same_existing_file(fs::path(resolved), app_path)); - EXPECT_TRUE(fs::path(resolved).filename() == "lld"); - }; + auto resolved = ct::win::payload::create_app_name_resolver().resolve("bin\\lld"); + EXPECT_TRUE(same_existing_file(fs::path(resolved), app_path)); + EXPECT_TRUE(fs::path(resolved).filename() == "lld"); + }; - TEST_CASE(command_line_resolver_searches_path_variable) { - TempSandbox sandbox; - auto cwd = sandbox.root / "cwd"; - fs::create_directories(cwd); - ScopedCurrentDirectory scope(cwd); + TEST_CASE(command_line_resolver_searches_path_variable) { + TempSandbox sandbox; + auto cwd = sandbox.root / "cwd"; + fs::create_directories(cwd); + ScopedCurrentDirectory scope(cwd); - auto app_dir = sandbox.root / "pathbin"; - auto app_path = app_dir / "runner.exe"; - touch_file(app_path); + auto app_dir = sandbox.root / "pathbin"; + auto app_path = app_dir / "runner.exe"; + touch_file(app_path); - ScopedEnvVar path_scope(L"PATH", app_dir.wstring()); + ScopedEnvVar path_scope(L"PATH", app_dir.wstring()); - auto resolved = ct::win::payload::create_command_line_resolver().resolve("runner"); - EXPECT_TRUE(same_existing_file(fs::path(resolved), app_path)); - }; + auto resolved = ct::win::payload::create_command_line_resolver().resolve("runner"); + EXPECT_TRUE(same_existing_file(fs::path(resolved), app_path)); + }; - TEST_CASE(resolve_abspath_supports_quoted_command_line) { - TempSandbox sandbox; - ScopedCurrentDirectory scope(sandbox.root); + TEST_CASE(resolve_abspath_supports_quoted_command_line) { + TempSandbox sandbox; + ScopedCurrentDirectory scope(sandbox.root); - auto app_path = sandbox.root / "quoted.exe"; - touch_file(app_path); + auto app_path = sandbox.root / "quoted.exe"; + touch_file(app_path); - auto resolved = ct::win::payload::resolve_abspath(nullptr, "\"quoted\" --help"); - EXPECT_TRUE(same_existing_file(fs::path(resolved), app_path)); - }; + auto resolved = ct::win::payload::resolve_abspath(nullptr, "\"quoted\" --help"); + EXPECT_TRUE(same_existing_file(fs::path(resolved), app_path)); + }; }; } // namespace From 717ba75bbd88ad7d5e89c0c4e906f70ba759369b Mon Sep 17 00:00:00 2001 From: seele Date: Mon, 23 Mar 2026 17:57:12 +0800 Subject: [PATCH 4/4] refactor: simplify resolver logic and enhance search path functionality --- src/catter-hook/win/payload/resolver.cc | 133 ++++++++---------------- src/catter-hook/win/payload/resolver.h | 23 +--- src/catter-hook/win/payload/winapi.h | 47 +++++++-- 3 files changed, 83 insertions(+), 120 deletions(-) diff --git a/src/catter-hook/win/payload/resolver.cc b/src/catter-hook/win/payload/resolver.cc index 92180ce3..544c0e61 100644 --- a/src/catter-hook/win/payload/resolver.cc +++ b/src/catter-hook/win/payload/resolver.cc @@ -1,97 +1,6 @@ #include "resolver.h" namespace catter::win::payload { - -template -std::basic_string Resolver::resolve(std::basic_string_view app_name) const { - if(app_name.empty()) { - return {}; - } - - auto original_name = std::basic_string(app_name); - auto fixed_name = fix_app_name(app_name); - if(contains_path(fixed_name)) { - auto direct = to_absolute(fixed_name); - if(!direct.empty() && is_file(direct)) { - return direct; - } - // Keep original input when the path cannot be resolved to an existing file. - return original_name; - } - - for(const auto& search_path: m_search_paths) { - auto candidate = join_path(search_path, fixed_name); - auto absolute_candidate = to_absolute(candidate); - if(absolute_candidate.empty()) { - continue; - } - if(is_file(absolute_candidate)) { - return absolute_candidate; - } - } - - // Keep original input when search paths do not resolve an existing file. - return original_name; -} - -template -bool Resolver::contains_path(std::basic_string_view value) { - return std::find_if(value.begin(), value.end(), [](char_t c) { - return is_path_sep(c) || c == char_t(':'); - }) != value.end(); -} - -template -bool Resolver::has_extension(std::basic_string_view file_name) { - return file_name.find(char_t('.')) != std::basic_string_view::npos; -} - -template -std::basic_string_view Resolver::exe_suffix() { - if constexpr(std::is_same_v) { - return ".exe"; - } else { - return L".exe"; - } -} - -template -std::basic_string Resolver::fix_app_name(std::basic_string_view app_name) { - if(contains_path(app_name) || app_name.back() == char_t('.') || has_extension(app_name)) { - return std::basic_string(app_name); - } - - std::basic_string fixed(app_name); - fixed.append(exe_suffix()); - return fixed; -} - -template -std::basic_string Resolver::join_path(std::basic_string_view directory, - std::basic_string_view file_name) { - if(directory.empty()) { - return std::basic_string(file_name); - } - - std::basic_string out(directory); - if(!is_path_sep(out.back())) { - out.push_back(path_sep); - } - out.append(file_name); - return out; -} - -template -std::basic_string Resolver::to_absolute(std::basic_string_view path) { - return GetFullPathNameDynamic(path); -} - -template -bool Resolver::is_file(std::basic_string_view path) { - auto attrs = FixGetFileAttributes(std::basic_string(path).c_str()); - return attrs != INVALID_FILE_ATTRIBUTES && (attrs & FILE_ATTRIBUTE_DIRECTORY) == 0; -} - namespace detail { template @@ -162,8 +71,50 @@ constexpr std::basic_string_view k_system16_name = "System"; template <> constexpr std::basic_string_view k_system16_name = L"System"; +template +constexpr std::basic_string_view exe_suffix; +template <> +constexpr std::basic_string_view exe_suffix = ".exe"; +template <> +constexpr std::basic_string_view exe_suffix = L".exe"; } // namespace detail +template +std::basic_string Resolver::resolve(std::basic_string_view app_name) const { + if(app_name.empty()) { + return {}; + } + + auto original_name = std::basic_string(app_name); + + auto search_paths = join_search_paths(m_search_paths); + auto resolved = SearchPathDynamic(search_paths.empty() ? nullptr : search_paths.c_str(), + app_name, + detail::exe_suffix.data()); + if(!resolved.empty()) { + return resolved; + } + + // Keep original input when search paths do not resolve an existing file. + return original_name; +} + +template +std::basic_string Resolver::join_search_paths( + const std::vector>& search_paths) { + std::basic_string merged; + for(const auto& path: search_paths) { + if(path.empty()) { + continue; + } + if(!merged.empty()) { + merged.push_back(char_t(';')); + } + merged.append(path); + } + return merged; +} + template Resolver create_app_name_resolver() { std::vector> search_paths; diff --git a/src/catter-hook/win/payload/resolver.h b/src/catter-hook/win/payload/resolver.h index 9983e6ce..0b47c496 100644 --- a/src/catter-hook/win/payload/resolver.h +++ b/src/catter-hook/win/payload/resolver.h @@ -20,21 +20,8 @@ class Resolver { std::basic_string resolve(std::basic_string_view app_name) const; private: - constexpr static char_t path_sep = char_t('\\'); - - constexpr static bool is_path_sep(char_t c) { - return c == char_t('\\') || c == char_t('/'); - } - - static bool contains_path(std::basic_string_view value); - static bool has_extension(std::basic_string_view file_name); - static std::basic_string_view exe_suffix(); - static std::basic_string fix_app_name(std::basic_string_view app_name); - static std::basic_string join_path(std::basic_string_view directory, - std::basic_string_view file_name); - static std::basic_string to_absolute(std::basic_string_view path); - static bool is_file(std::basic_string_view path); - + static std::basic_string + join_search_paths(const std::vector>& search_paths); std::vector> m_search_paths; }; @@ -47,10 +34,4 @@ Resolver create_command_line_resolver(); extern template class Resolver; extern template class Resolver; -template -Resolver create_app_name_resolver(); - -template -Resolver create_command_line_resolver(); - } // namespace catter::win::payload diff --git a/src/catter-hook/win/payload/winapi.h b/src/catter-hook/win/payload/winapi.h index b601cc86..97db5dce 100644 --- a/src/catter-hook/win/payload/winapi.h +++ b/src/catter-hook/win/payload/winapi.h @@ -1,5 +1,6 @@ #pragma once +#include #include #include @@ -106,6 +107,20 @@ DWORD FixGetModuleFileName(HMODULE module, char_t* buffer, DWORD size) { } } +template +DWORD FixSearchPath(const char_t* path, + const char_t* file_name, + const char_t* extension, + DWORD buffer_size, + char_t* buffer, + char_t** file_part) { + if constexpr(std::is_same_v) { + return SearchPathA(path, file_name, extension, buffer_size, buffer, file_part); + } else { + return SearchPathW(path, file_name, extension, buffer_size, buffer, file_part); + } +} + template UINT FixGetSystemDirectory(char_t* buffer, UINT size) { if constexpr(std::is_same_v) { @@ -177,15 +192,11 @@ std::basic_string GetModuleDirectory(HMODULE module, size_t initial_size if(module_path.empty()) { return {}; } - - auto pos = module_path.find_last_of(char_t('\\')); - if(pos == std::basic_string::npos) { - pos = module_path.find_last_of(char_t('/')); - } - if(pos == std::basic_string::npos) { - return {}; + if constexpr(std::is_same_v) { + return std::filesystem::path(module_path).parent_path().string(); + } else { + return std::filesystem::path(module_path).parent_path().wstring(); } - return module_path.substr(0, pos); } template @@ -233,4 +244,24 @@ std::basic_string GetFullPathNameDynamic(std::basic_string_view return WinApiBufferDecision{.next_size = static_cast(result) + 1}; }); } + +template +std::basic_string SearchPathDynamic(const char_t* path, + std::basic_string_view file_name, + const char_t* extension = nullptr, + size_t initial_size = MAX_PATH) { + auto input = std::basic_string(file_name); + return CallWinApiWithGrowingBuffer( + initial_size, + [&](char_t* buffer, DWORD size) { + return FixSearchPath(path, input.c_str(), extension, size, buffer, nullptr); + }, + [](DWORD result, size_t buffer_size) -> WinApiBufferDecision { + if(static_cast(result) < buffer_size) { + return WinApiBufferDecision{.done = true, + .output_size = static_cast(result)}; + } + return WinApiBufferDecision{.next_size = static_cast(result) + 1}; + }); +} } // namespace catter::win::payload