diff options
Diffstat (limited to 'src/envfs.cpp')
| -rw-r--r-- | src/envfs.cpp | 450 |
1 files changed, 450 insertions, 0 deletions
diff --git a/src/envfs.cpp b/src/envfs.cpp new file mode 100644 index 00000000..a67faf84 --- /dev/null +++ b/src/envfs.cpp @@ -0,0 +1,450 @@ +#include "envfs.h" +#include "env.h" +#include "util.h" +#include <utility.h> +#include <log.h> + +using namespace MOBase; + +typedef struct _UNICODE_STRING { + USHORT Length; + USHORT MaximumLength; + PWSTR Buffer; +} UNICODE_STRING, *PUNICODE_STRING; +typedef const UNICODE_STRING *PCUNICODE_STRING; + +typedef struct _OBJECT_ATTRIBUTES { + ULONG Length; + HANDLE RootDirectory; + PUNICODE_STRING ObjectName; + ULONG Attributes; + PVOID SecurityDescriptor; + PVOID SecurityQualityOfService; +} OBJECT_ATTRIBUTES, *POBJECT_ATTRIBUTES; + + +typedef struct _FILE_DIRECTORY_INFORMATION { + ULONG NextEntryOffset; + ULONG FileIndex; + LARGE_INTEGER CreationTime; + LARGE_INTEGER LastAccessTime; + LARGE_INTEGER LastWriteTime; + LARGE_INTEGER ChangeTime; + LARGE_INTEGER EndOfFile; + LARGE_INTEGER AllocationSize; + ULONG FileAttributes; + ULONG FileNameLength; + WCHAR FileName[1]; +} FILE_DIRECTORY_INFORMATION, *PFILE_DIRECTORY_INFORMATION; + + +#define FILE_SHARE_VALID_FLAGS 0x00000007 + +// copied from ntstatus.h +#define STATUS_SUCCESS ((NTSTATUS)0x00000000L) +#define STATUS_BUFFER_OVERFLOW ((NTSTATUS)0x80000005L) +#define STATUS_NO_MORE_FILES ((NTSTATUS)0x80000006L) +#define STATUS_NO_SUCH_FILE ((NTSTATUS)0xC000000FL) + +typedef struct _IO_STATUS_BLOCK IO_STATUS_BLOCK; + +typedef struct _IO_STATUS_BLOCK *PIO_STATUS_BLOCK; +// typedef VOID (NTAPI *PIO_APC_ROUTINE )(__in PVOID ApcContext, __in +// PIO_STATUS_BLOCK IoStatusBlock, __in ULONG Reserved); +typedef VOID(NTAPI *PIO_APC_ROUTINE)(PVOID ApcContext, + PIO_STATUS_BLOCK IoStatusBlock, + ULONG Reserved); + + +typedef enum _FILE_INFORMATION_CLASS { + FileDirectoryInformation = 1 +} FILE_INFORMATION_CLASS; + +typedef NTSTATUS(WINAPI *NtQueryDirectoryFile_type)( + HANDLE, HANDLE, PIO_APC_ROUTINE, PVOID, PIO_STATUS_BLOCK, PVOID, ULONG, + FILE_INFORMATION_CLASS, BOOLEAN, PUNICODE_STRING, BOOLEAN); + +typedef NTSTATUS(WINAPI *NtOpenFile_type)(PHANDLE, ACCESS_MASK, + POBJECT_ATTRIBUTES, PIO_STATUS_BLOCK, + ULONG, ULONG); + +typedef NTSTATUS(WINAPI *NtClose_type)(HANDLE); + + +NtOpenFile_type NtOpenFile = nullptr; +NtQueryDirectoryFile_type NtQueryDirectoryFile = nullptr; +extern NtClose_type NtClose = nullptr; + + +#define FILE_DIRECTORY_FILE 0x00000001 +#define FILE_WRITE_THROUGH 0x00000002 +#define FILE_SEQUENTIAL_ONLY 0x00000004 +#define FILE_NO_INTERMEDIATE_BUFFERING 0x00000008 + +#define FILE_SYNCHRONOUS_IO_ALERT 0x00000010 +#define FILE_SYNCHRONOUS_IO_NONALERT 0x00000020 +#define FILE_NON_DIRECTORY_FILE 0x00000040 +#define FILE_CREATE_TREE_CONNECTION 0x00000080 + +#define FILE_COMPLETE_IF_OPLOCKED 0x00000100 +#define FILE_NO_EA_KNOWLEDGE 0x00000200 +#define FILE_OPEN_REMOTE_INSTANCE 0x00000400 +#define FILE_RANDOM_ACCESS 0x00000800 + +#define FILE_DELETE_ON_CLOSE 0x00001000 +#define FILE_OPEN_BY_FILE_ID 0x00002000 +#define FILE_OPEN_FOR_BACKUP_INTENT 0x00004000 +#define FILE_NO_COMPRESSION 0x00008000 + +#if (_WIN32_WINNT >= _WIN32_WINNT_WIN7) +#define FILE_OPEN_REQUIRING_OPLOCK 0x00010000 +#endif + +#define FILE_RESERVE_OPFILTER 0x00100000 +#define FILE_OPEN_REPARSE_POINT 0x00200000 +#define FILE_OPEN_NO_RECALL 0x00400000 +#define FILE_OPEN_FOR_FREE_SPACE_QUERY 0x00800000 + +#define FILE_VALID_OPTION_FLAGS 0x00ffffff +#define FILE_VALID_PIPE_OPTION_FLAGS 0x00000032 +#define FILE_VALID_MAILSLOT_OPTION_FLAGS 0x00000032 +#define FILE_VALID_SET_FLAGS 0x00000036 + +typedef struct _IO_STATUS_BLOCK { +#pragma warning(push) +#pragma warning(disable: 4201) // we'll always use the Microsoft compiler + union { + NTSTATUS Status; + PVOID Pointer; + } DUMMYUNIONNAME; +#pragma warning(pop) + + ULONG_PTR Information; +} IO_STATUS_BLOCK, *PIO_STATUS_BLOCK; + + + +namespace env +{ + +std::wstring_view toStringView(const UNICODE_STRING* s) +{ + if (s && s->Buffer) { + return {s->Buffer, (s->Length / sizeof(wchar_t))}; + } else { + return {}; + } +} + +std::wstring_view toStringView(POBJECT_ATTRIBUTES poa) +{ + if (poa->ObjectName) { + return toStringView(poa->ObjectName); + } + + return {}; +} + +QString toString(POBJECT_ATTRIBUTES poa) +{ + const auto sv = toStringView(poa); + return QString::fromWCharArray(sv.data(), static_cast<int>(sv.size())); +} + + +class HandleCloserThread +{ +public: + HandleCloserThread() + : m_ready(false) + { + m_handles.reserve(50'000); + } + + void shrink() + { + m_handles.shrink_to_fit(); + } + + void add(HANDLE h) + { + m_handles.push_back(h); + } + + void wakeup() + { + m_ready = true; + m_cv.notify_one(); + } + + void run() + { + MOShared::SetThisThreadName("HandleCloserThread"); + + std::unique_lock lock(m_mutex); + m_cv.wait(lock, [&]{ return m_ready; }); + + closeHandles(); + } + +private: + std::vector<HANDLE> m_handles; + std::condition_variable m_cv; + std::mutex m_mutex; + bool m_ready; + + void closeHandles() + { + for (auto& h : m_handles) { + NtClose(h); + } + + m_handles.clear(); + m_ready = false; + } +}; + +constexpr std::size_t AllocSize = 1024 * 1024; +static ThreadPool<HandleCloserThread> g_handleClosers; + +void setHandleCloserThreadCount(std::size_t n) +{ + g_handleClosers.setMax(n); +} + +void forEachEntryImpl( + void* cx, HandleCloserThread& hc, std::vector<std::unique_ptr<unsigned char[]>>& buffers, + POBJECT_ATTRIBUTES poa, std::size_t depth, + DirStartF* dirStartF, DirEndF* dirEndF, FileF* fileF) +{ + IO_STATUS_BLOCK iosb; + UNICODE_STRING ObjectName; + OBJECT_ATTRIBUTES oa = { sizeof(oa), 0, &ObjectName }; + NTSTATUS status; + + status = NtOpenFile( + &oa.RootDirectory, FILE_GENERIC_READ, poa, &iosb, FILE_SHARE_VALID_FLAGS, + FILE_SYNCHRONOUS_IO_NONALERT|FILE_OPEN_REPARSE_POINT|FILE_OPEN_FOR_BACKUP_INTENT); + + if (status < 0) { + log::error( + "NtOpenFile() failed for '{}', {}", + toString(poa), formatSystemMessage(status)); + + return; + } + + hc.add(oa.RootDirectory); + unsigned char* buffer; + + if (depth >= buffers.size()) { + buffers.emplace_back(std::make_unique<unsigned char[]>(AllocSize)); + buffer = buffers.back().get(); + } else { + buffer = buffers[depth].get(); + } + + union + { + PVOID pv; + PBYTE pb; + PFILE_DIRECTORY_INFORMATION DirInfo; + }; + + for (;;) { + status = NtQueryDirectoryFile( + oa.RootDirectory, NULL, NULL, NULL, &iosb, + buffer, AllocSize, FileDirectoryInformation, FALSE, NULL, FALSE); + + if (status == STATUS_NO_MORE_FILES) { + break; + } else if (status < 0) { + log::error( + "NtQueryDirectoryFile() failed for '{}', {}", + toString(poa), formatSystemMessage(status)); + break; + } + + ULONG NextEntryOffset = 0; + + pv = buffer; + + auto isDotDir = [](auto* o) { + if (o->Length == 2 && o->Buffer[0] == '.') { + return true; + } + + if (o->Length == 4 && o->Buffer[0] == '.' && o->Buffer[1] == '.') { + return true; + } + + return false; + }; + + std::size_t count = 0; + + for (;;) { + ++count; + pb += NextEntryOffset; + + ObjectName.Buffer = DirInfo->FileName; + ObjectName.Length = (USHORT)DirInfo->FileNameLength; + + if (!isDotDir(&ObjectName)) { + ObjectName.MaximumLength = ObjectName.Length; + + if (DirInfo->FileAttributes & FILE_ATTRIBUTE_DIRECTORY) { + dirStartF(cx, toStringView(&oa)); + forEachEntryImpl(cx, hc, buffers, &oa, depth+1, dirStartF, dirEndF, fileF); + dirEndF(cx, toStringView(&oa)); + } else { + FILETIME ft; + ft.dwLowDateTime = DirInfo->LastWriteTime.LowPart; + ft.dwHighDateTime = DirInfo->LastWriteTime.HighPart; + fileF(cx, toStringView(&oa), ft); + } + } + + NextEntryOffset = DirInfo->NextEntryOffset; + + if (NextEntryOffset == 0) { + break; + } + } + + if (AllocSize - iosb.Information > (ULONG)FIELD_OFFSET(FILE_DIRECTORY_INFORMATION, FileName[256])) { + // NO_MORE_FILES + break; + } + } +} + + +void DirectoryWalker::forEachEntry( + const std::wstring& path, void* cx, + DirStartF* dirStartF, DirEndF* dirEndF, FileF* fileF) +{ + auto& hc = g_handleClosers.request(); + + if (!NtOpenFile) { + LibraryPtr m(::LoadLibraryW(L"ntdll.dll")); + NtOpenFile = (NtOpenFile_type)::GetProcAddress(m.get(), "NtOpenFile"); + NtQueryDirectoryFile = (NtQueryDirectoryFile_type)::GetProcAddress(m.get(), "NtQueryDirectoryFile"); + NtClose = (NtClose_type)::GetProcAddress(m.get(), "NtClose"); + } + + const std::wstring ntpath = std::wstring(L"\\??\\") + path; + + UNICODE_STRING ObjectName = {}; + ObjectName.Buffer = const_cast<wchar_t*>(ntpath.c_str()); + ObjectName.Length = (USHORT)ntpath.size() * sizeof(wchar_t); + ObjectName.MaximumLength = ObjectName.Length; + + OBJECT_ATTRIBUTES oa = {}; + oa.Length = sizeof(oa); + oa.ObjectName = &ObjectName; + + forEachEntryImpl(cx, hc, m_buffers, &oa, 0, dirStartF, dirEndF, fileF); + hc.wakeup(); +} + + +void forEachEntry( + const std::wstring& path, void* cx, + DirStartF* dirStartF, DirEndF* dirEndF, FileF* fileF) +{ + DirectoryWalker().forEachEntry(path, cx, dirStartF, dirEndF, fileF); +} + +Directory getFilesAndDirs(const std::wstring& path) +{ + struct Context + { + std::stack<Directory*> current; + }; + + Directory root; + + Context cx; + cx.current.push(&root); + + env::forEachEntry(path, &cx, + [](void* pcx, std::wstring_view path) { + Context* cx = (Context*)pcx; + + cx->current.top()->dirs.push_back(Directory(path)); + cx->current.push(&cx->current.top()->dirs.back()); + }, + + [](void* pcx, std::wstring_view path) { + Context* cx = (Context*)pcx; + cx->current.pop(); + }, + + [](void* pcx, std::wstring_view path, FILETIME ft) { + Context* cx = (Context*)pcx; + + cx->current.top()->files.push_back(File(path, ft)); + } + ); + + return root; +} + +File::File(std::wstring_view n, FILETIME ft) : + name(n.begin(), n.end()), + lcname(MOShared::ToLowerCopy(name)), + lastModified(ft) +{ +} + +Directory::Directory() +{ +} + +Directory::Directory(std::wstring_view n) + : name(n.begin(), n.end()), lcname(MOShared::ToLowerCopy(name)) +{ +} + + +void getFilesAndDirsWithFindImpl(const std::wstring& path, Directory& d) +{ + const std::wstring searchString = path + L"\\*"; + + WIN32_FIND_DATAW findData; + + HANDLE searchHandle = ::FindFirstFileExW( + searchString.c_str(), FindExInfoBasic, &findData, FindExSearchNameMatch, + nullptr, FIND_FIRST_EX_LARGE_FETCH); + + if (searchHandle != INVALID_HANDLE_VALUE) { + BOOL result = true; + + while (result) { + if (findData.dwFileAttributes & FILE_ATTRIBUTE_DIRECTORY) { + if ((wcscmp(findData.cFileName, L".") != 0) && + (wcscmp(findData.cFileName, L"..") != 0)) { + const std::wstring newPath = path + L"\\" + findData.cFileName; + d.dirs.push_back(Directory(findData.cFileName)); + getFilesAndDirsWithFindImpl(newPath, d.dirs.back()); + } + } else { + d.files.push_back(File(findData.cFileName, findData.ftLastWriteTime)); + } + + result = ::FindNextFileW(searchHandle, &findData); + } + } + + ::FindClose(searchHandle); +} + +Directory getFilesAndDirsWithFind(const std::wstring& path) +{ + Directory d; + getFilesAndDirsWithFindImpl(path, d); + return d; +} + +} // namespace |