mirrors/Sandboxie
mirrors
/
Sandboxie
mirror of https://github.com/sandboxie-plus/Sandboxie.git
1
0
Fork 0
Code Releases Activity
Sandboxie/SandboxieTools/ImBox/ImDiskIO.cpp

765 lines
22 KiB
C++
Raw Blame History

#include "framework.h"
#include <stdio.h>
#include "..\Common\helpers.h"
#include "..\ImDisk\inc\imdproxy.h"
#include "..\ImDisk\inc\imdisk.h"
#include "ImDiskIO.h"
#include "ImBox.h"
extern "C" {
#define ObjectNameInformation 1
typedef struct _OBJECT_NAME_INFORMATION {
UNICODE_STRING Name;
} OBJECT_NAME_INFORMATION, *POBJECT_NAME_INFORMATION;
#define FILE_SHARE_VALID_FLAGS 0x00000007
#define IO_REPARSE_TAG_SYMLINK (0xA000000CL)
#define SYMLINK_FLAG_RELATIVE 1
typedef struct _REPARSE_DATA_BUFFER {
ULONG ReparseTag;
USHORT ReparseDataLength;
USHORT Reserved;
union {
struct {
USHORT SubstituteNameOffset;
USHORT SubstituteNameLength;
USHORT PrintNameOffset;
USHORT PrintNameLength;
ULONG Flags;
WCHAR PathBuffer[1];
} SymbolicLinkReparseBuffer;
struct {
USHORT SubstituteNameOffset;
USHORT SubstituteNameLength;
USHORT PrintNameOffset;
USHORT PrintNameLength;
WCHAR PathBuffer[1];
} MountPointReparseBuffer;
struct {
UCHAR DataBuffer[1];
} GenericReparseBuffer;
};
} REPARSE_DATA_BUFFER, *PREPARSE_DATA_BUFFER;
typedef enum _FSINFOCLASS {
FileFsVolumeInformation = 1, // FILE_FS_VOLUME_INFORMATION
FileFsLabelInformation, // 2
FileFsSizeInformation, // 3
FileFsDeviceInformation, // 4
FileFsAttributeInformation, // 5 // FILE_FS_ATTRIBUTE_INFORMATION
FileFsControlInformation, // 6
FileFsFullSizeInformation, // 7
FileFsObjectIdInformation, // 8
FileFsDriverPathInformation, // 9
FileFsMaximumInformation
} FS_INFORMATION_CLASS, *PFS_INFORMATION_CLASS;
NTSYSCALLAPI NTSTATUS NTAPI NtQueryVolumeInformationFile(
IN HANDLE FileHandle,
OUT PIO_STATUS_BLOCK IoStatusBlock,
OUT PVOID FsInformation,
IN ULONG Length,
IN ULONG FsInformationClass
);
typedef struct _FILE_FS_VOLUME_INFORMATION {
LARGE_INTEGER VolumeCreationTime;
ULONG VolumeSerialNumber;
ULONG VolumeLabelLength;
BOOLEAN SupportsObjects;
WCHAR VolumeLabel[1];
} FILE_FS_VOLUME_INFORMATION, *PFILE_FS_VOLUME_INFORMATION;
typedef struct _FILE_FS_ATTRIBUTE_INFORMATION {
ULONG FileSystemAttributes;
LONG MaximumComponentNameLength;
ULONG FileSystemNameLength;
WCHAR FileSystemName[1];
} FILE_FS_ATTRIBUTE_INFORMATION, *PFILE_FS_ATTRIBUTE_INFORMATION;
NTSTATUS NTAPI RtlSetThreadErrorMode(IN ULONG NewMode, OUT PULONG OldMode);
bool IsVolumeUnRecognized(std::wstring NtPath);
bool FormatVolume(LPCWSTR root, LPCWSTR fs, LPCWSTR label);
WCHAR WINAPI MyImDiskFindFreeDriveLetter();
HANDLE WINAPI MyImDiskOpenDeviceByMountPoint(LPCWSTR MountPoint, DWORD AccessMode);
BOOL WINAPI MyImDiskCliValidateDriveLetterTarget(LPCWSTR DriveLetter, LPCWSTR ValidTargetPath);
}
struct SImDiskIO
{
std::wstring Mount;
std::wstring Format;
std::wstring Params;
HANDLE hImDisk;
std::wstring Proxy;
HANDLE hEvent;
HANDLE hMapping;
WCHAR* pSection;
};
CImDiskIO::CImDiskIO(CAbstractIO* pIO, const std::wstring& Mount, const std::wstring& Format, const std::wstring& Params)
{
m = new SImDiskIO;
m->Mount = Mount;
m->Format = Format;
m->Params = Params;
m->hImDisk = INVALID_HANDLE_VALUE;
m->hEvent = NULL;
m->hMapping = NULL;
m->pSection = NULL;
m_pIO = pIO;
m_hThread = INVALID_HANDLE_VALUE;
}
CImDiskIO::~CImDiskIO()
{
if (m_hThread != INVALID_HANDLE_VALUE) {
if (WaitForSingleObject(m_hThread, 60 * 1000) == WAIT_TIMEOUT)
TerminateThread(m_hThread, -1);
CloseHandle(m_hThread);
}
if(m) delete m;
}
void CImDiskIO::SetProxyName(const std::wstring& Name)
{
if (m) m->Proxy = Name;
}
void CImDiskIO::SetMountEvent(HANDLE hEvent)
{
if (m) m->hEvent = hEvent;
}
void CImDiskIO::SetMountSection(HANDLE hMapping, WCHAR* pSection)
{
if (m) m->hMapping = hMapping;
if (m) m->pSection = pSection;
}
DWORD WINAPI CImDiskIO_Thread(LPVOID lpThreadParameter)
{
//
// !!! NOTE !!! After this thread was created the rest of the code does not touch
// the m members which we use here so we don't bother with explicit synchronization
// hence this thread is responsible freeing m
//
SImDiskIO* m = ((SImDiskIO*)lpThreadParameter);
std::wstring Device;
DWORD exit_code;
do {
Sleep(100);
HANDLE handle = (HANDLE)MyImDiskOpenDeviceByMountPoint(m->Mount.c_str(), 0);
if (handle != INVALID_HANDLE_VALUE) {
BYTE buffer[MAX_PATH];
DWORD length = sizeof(buffer);
if (NT_SUCCESS(NtQueryObject(handle, (OBJECT_INFORMATION_CLASS)ObjectNameInformation, buffer, length, &length))) {
UNICODE_STRING* uni = &((OBJECT_NAME_INFORMATION*)buffer)->Name;
length = uni->Length / sizeof(WCHAR);
if (uni->Buffer) {
uni->Buffer[length] = 0;
Device = uni->Buffer;
}
}
CloseHandle(handle);
break;
}
GetExitCodeProcess(m->hImDisk, &exit_code);
} while (exit_code == STILL_ACTIVE);
CloseHandle(m->hImDisk); m->hImDisk = INVALID_HANDLE_VALUE;
if (!Device.empty() && !m->Format.empty()) {
std::wstring Drive;
if (!IsVolumeUnRecognized(Device)) {
DbgPrint(L"The volume: %s was recognized, format skipped.\n", Device.c_str());
}
else
if ((m->Mount.length() == 2 && m->Mount[1] == L':') // check if mount is a drive letter
|| (m->Mount.length() == 3 && m->Mount[1] == L':' && m->Mount[2] == L'\\')) {
Drive = m->Mount;
}
else {
WCHAR drive = MyImDiskFindFreeDriveLetter();
if (!drive) {
DbgPrint(L"No free drive letter found.\n");
}
else {
Drive = L" :";
Drive[0] = drive;
}
if (!DefineDosDevice(DDD_RAW_TARGET_PATH, Drive.c_str(), Device.c_str())) {
DbgPrint(L"Failed to Mount drive letter.\n");
Drive.clear();
}
}
if (!Drive.empty()) {
if (MyImDiskCliValidateDriveLetterTarget(Drive.c_str(), Device.c_str())) {
LPCWSTR fs = m->Format.c_str();
LPCWSTR label = wcschr(fs, L':');
if (!label) label = L"";
else *(*(LPWSTR*)&label)++ = L'\0';
//for (int i = 0; i < 3; i++) {
if (FormatVolume(Drive.c_str(), fs, label)) {
if (!IsVolumeUnRecognized(Device)) // check success
DbgPrint(L"Successfully Formatted: %s\n", m->Mount.c_str());
else {
DbgPrint(L"Failed to Format: %s\n", m->Mount.c_str());
// Sleep(1000);
// continue; // retry
}
}
else // fails only when lib is not available
DbgPrint(L"Can not Format: %s\n", m->Mount.c_str());
// break;
//}
}
if (Drive != m->Mount) {
if (!DefineDosDevice(DDD_REMOVE_DEFINITION | DDD_EXACT_MATCH_ON_REMOVE | DDD_RAW_TARGET_PATH, Drive.c_str(), Device.c_str())) {
DbgPrint(L"Failed to Unmount drive letter.\n");
}
}
}
}
if (m->pSection) {
wmemcpy(m->pSection, Device.c_str(), Device.length() + 1);
UnmapViewOfFile(m->pSection);
}
if(m->hMapping)
CloseHandle(m->hMapping);
if (m->hEvent) {
SetEvent(m->hEvent);
CloseHandle(m->hEvent);
}
delete m;
return 0;
}
int CImDiskIO::DoComm()
{
HANDLE hFileMap;
ULARGE_INTEGER map_size;
BYTE *shm_view, *main_buf;
struct s_req_block {unsigned char request_code, pad[7]; ULONGLONG offset; ULONGLONG length;} *req_block;
struct s_resp_block {unsigned char errorno, pad[7]; ULONGLONG length;} *resp_block;
struct s_trim_block {unsigned char request_code, pad[7]; unsigned int length;} *trim_block;
HANDLE shm_request_event, shm_response_event;
if (m->Proxy.empty()){
#ifdef _M_ARM64
ULONG64 ctr = _ReadStatusReg(ARM64_CNTVCT);
#else
ULONG64 ctr = __rdtsc();
#endif
m->Proxy = L"ImBox" + std::to_wstring(ctr);
}
map_size.QuadPart = DEF_BUFFER_SIZE + IMDPROXY_HEADER_SIZE;
if (!(hFileMap = CreateFileMapping(INVALID_HANDLE_VALUE, NULL, PAGE_READWRITE | SEC_COMMIT, map_size.HighPart, map_size.LowPart, (L"Global\\" + m->Proxy).c_str())))
return ERR_FILE_MAPPING;
if (!(shm_view = (BYTE*)MapViewOfFile(hFileMap, FILE_MAP_WRITE, 0, 0, 0)))
return ERR_FILE_MAPPING;
m_pIO->PrepViewOfFile(shm_view);
main_buf = shm_view + IMDPROXY_HEADER_SIZE;
req_block = (s_req_block*)shm_view;
resp_block = (s_resp_block*)shm_view;
trim_block = (s_trim_block*)shm_view;
if (!(shm_request_event = CreateEvent(NULL, FALSE, FALSE, (L"Global\\" + m->Proxy + L"_Request").c_str())) || GetLastError() == ERROR_ALREADY_EXISTS)
return ERR_CREATE_EVENT;
if (!(shm_response_event = CreateEvent(NULL, FALSE, FALSE, (L"Global\\" + m->Proxy + L"_Response").c_str())))
return ERR_CREATE_EVENT;
STARTUPINFO si = {sizeof si};
PROCESS_INFORMATION pi;
std::wstring cmd = L"imdisk -a -t proxy -o shm -f " + m->Proxy;
if (!m->Mount.empty()) cmd += L" -m \"" + m->Mount + L"\"";
if (!m->Params.empty())cmd += L" " + m->Params;
if (!CreateProcess(NULL, (WCHAR*)cmd.c_str(), NULL, NULL, FALSE, CREATE_NO_WINDOW, NULL, NULL, &si, &pi)) {
DbgPrint(L"Failed to run imdisk.exe.\n");
return ERR_IMDISK_FAILED;
}
NtClose(pi.hThread);
if (!m->Mount.empty()) {
//
// clear format directive if formatting is not allowed for this disk e.g. non empty image file
// to force format you can use imdisk: params="-p \"/fs:ntfs /q /y\""
//
if (!m_pIO->CanBeFormated())
m->Format.clear();
m->hImDisk = pi.hProcess;
m_hThread = CreateThread(NULL, 0, CImDiskIO_Thread, m, 0, NULL);
m = NULL;
}
else {
NtClose(pi.hProcess);
}
LARGE_INTEGER t;
t.QuadPart = -100000000; // A negative value specifies an interval relative to the current time in units of 100 nanoseconds.
if (NtWaitForSingleObject(shm_request_event, FALSE, &t) != STATUS_SUCCESS || req_block->request_code != IMDPROXY_REQ_INFO)
return ERR_IMDISK_TIMEOUT;
IMDPROXY_INFO_RESP proxy_info = {0};
proxy_info.file_size = m_pIO->GetDiskSize();
proxy_info.req_alignment = 1;
proxy_info.flags = IMDPROXY_FLAG_SUPPORTS_UNMAP; // TRIM
memcpy(shm_view, &proxy_info, sizeof proxy_info);
for (;;) {
NtSignalAndWaitForSingleObject(shm_response_event, shm_request_event, FALSE, NULL);
if (req_block->request_code == IMDPROXY_REQ_READ) {
if (!m_pIO->DiskRead(main_buf, req_block->length, req_block->offset)) {
DbgPrint(L"DiskRead error.\n");
}
}
else if (req_block->request_code == IMDPROXY_REQ_WRITE) {
if (!m_pIO->DiskWrite(main_buf, req_block->length, req_block->offset)) {
DbgPrint(L"DiskWrite error, SOME DATA WILL BE LOST.");
}
}
else if (req_block->request_code == IMDPROXY_REQ_UNMAP) {
m_pIO->TrimProcess((DEVICE_DATA_SET_RANGE*)main_buf, trim_block->length / sizeof(DEVICE_DATA_SET_RANGE));
}
else if (req_block->request_code == IMDPROXY_REQ_CLOSE) {
return ERR_OK;
}
else { // unknown command
DbgPrint(L"Unknown Command: %d\n", req_block->request_code);
return ERR_UNKNOWN_COMMAND;
}
resp_block->errorno = 0;
resp_block->length = req_block->length;
}
}
//static void disp_message(WCHAR *disp_text, WCHAR *arg, BOOL wait)
//{
// DWORD dw;
//
// _snwprintf(txt, _countof(txt) - 1, disp_text, arg);
// WTSSendMessage(WTS_CURRENT_SERVER_HANDLE, WTSGetActiveConsoleSessionId(), L"ImDisk", 14, txt, (wcslen(txt) + 1) * sizeof(WCHAR), MB_OK | MB_ICONERROR, 0, &dw, wait);
//}
//
//static void disp_err_mem()
//{
// if (GetTickCount() - err_time >= 10000) {
// disp_message(L"Not enough memory to write data into %s.\nSome data will be lost.", drive_arg, TRUE);
// err_time = GetTickCount();
// }
//}
extern "C" {
bool IsVolumeUnRecognized(std::wstring NtPath)
{
if (NtPath.back() != L'\\') NtPath.push_back(L'\\');
HANDLE handle;
IO_STATUS_BLOCK iosb;
UNICODE_STRING objname;
RtlInitUnicodeString(&objname, NtPath.c_str());
OBJECT_ATTRIBUTES objattrs;
InitializeObjectAttributes(
&objattrs, &objname, OBJ_CASE_INSENSITIVE, NULL, NULL);
ULONG OldMode;
RtlSetThreadErrorMode(0x10u, &OldMode);
NTSTATUS status = NtCreateFile(
&handle, GENERIC_READ | SYNCHRONIZE, &objattrs,
&iosb, NULL, 0, FILE_SHARE_VALID_FLAGS,
FILE_OPEN,
FILE_DIRECTORY_FILE | FILE_SYNCHRONOUS_IO_NONALERT,
NULL, 0);
RtlSetThreadErrorMode(OldMode, 0i64);
if (NT_SUCCESS(status))
{
union {
FILE_FS_ATTRIBUTE_INFORMATION fsInfo;
BYTE fsInfoBuff[64];
} u;
if (NT_SUCCESS(NtQueryVolumeInformationFile(handle, &iosb, &u.fsInfo, sizeof(u), FileFsAttributeInformation))) {
u.fsInfo.FileSystemName[u.fsInfo.FileSystemNameLength / sizeof(wchar_t)] = 0;
DbgPrint(L"Recognized FileSystem: %s\n", u.fsInfo.FileSystemName);
}
NtClose(handle);
}
if (status == STATUS_UNRECOGNIZED_VOLUME)
return true;
if (!NT_SUCCESS(status))
DbgPrint(L"NtQueryVolumeInformationFile failed 0x%08X.\n", status);
return false;
}
// types from winfile
typedef BOOLEAN(WINAPI *FMIFS_CALLBACK)(ULONG PacketType, ULONG PacketLength, PVOID PacketData);
typedef void (WINAPI* PFORMAT) (PWSTR DriveName, ULONG MediaType, PWSTR FileSystemName, PWSTR Label, BOOLEAN Quick, FMIFS_CALLBACK Callback);
static BOOLEAN WINAPI my_format_callback(ULONG PacketType, ULONG PacketLength, PVOID PacketData) {
return 1;
}
bool FormatVolume(LPCWSTR root, LPCWSTR fs, LPCWSTR label)
{
bool ret = false;
HMODULE fmifs = LoadLibrary(L"fmifs");
if (fmifs != NULL) {
PFORMAT Format = (PFORMAT)GetProcAddress(fmifs, "Format");
if (Format != NULL) {
Format((PWSTR)root, 0, (PWSTR)fs, (PWSTR)label, TRUE, my_format_callback);
ret = true;
}
FreeLibrary(fmifs);
}
return ret;
}
// ImDisk
HANDLE WINAPI MyImDiskOpenDeviceByName(PUNICODE_STRING FileName, DWORD AccessMode)
{
NTSTATUS status;
HANDLE handle;
OBJECT_ATTRIBUTES object_attrib;
IO_STATUS_BLOCK io_status;
InitializeObjectAttributes(&object_attrib,
FileName,
OBJ_CASE_INSENSITIVE,
NULL,
NULL);
status = NtOpenFile(&handle,
SYNCHRONIZE | AccessMode,
&object_attrib,
&io_status,
FILE_SHARE_READ | FILE_SHARE_WRITE,
FILE_NON_DIRECTORY_FILE | FILE_SYNCHRONOUS_IO_NONALERT);
if (!NT_SUCCESS(status))
{
SetLastError(RtlNtStatusToDosError(status));
return INVALID_HANDLE_VALUE;
}
return handle;
}
HANDLE WINAPI MyImDiskOpenDeviceByNumber(DWORD DeviceNumber, DWORD AccessMode)
{
WCHAR device_path[_countof(IMDISK_DEVICE_BASE_NAME) + 16];
UNICODE_STRING file_name;
// Build device path, e.g. \Device\ImDisk2
_snwprintf_s(device_path, ARRAYSIZE(device_path), _countof(device_path),
IMDISK_DEVICE_BASE_NAME L"%u", DeviceNumber);
device_path[_countof(device_path) - 1] = 0;
RtlInitUnicodeString(&file_name, device_path);
return MyImDiskOpenDeviceByName(&file_name, AccessMode);
}
HANDLE WINAPI MyImDiskOpenDeviceByMountPoint(LPCWSTR MountPoint, DWORD AccessMode)
{
UNICODE_STRING DeviceName;
WCHAR DriveLetterPath[] = L"\\DosDevices\\ :";
PREPARSE_DATA_BUFFER ReparseData = NULL;
HANDLE h;
if ((MountPoint[0] != 0) &&
((wcscmp(MountPoint + 1, L":") == 0) ||
(wcscmp(MountPoint + 1, L":\\") == 0)))
{
DriveLetterPath[12] = MountPoint[0];
RtlInitUnicodeString(&DeviceName, DriveLetterPath);
}
else if (((wcsncmp(MountPoint, L"\\\\?\\", 4) == 0) ||
(wcsncmp(MountPoint, L"\\\\.\\", 4) == 0)) &&
(wcschr(MountPoint + 4, L'\\') == NULL))
{
return CreateFile(MountPoint, AccessMode,
FILE_SHARE_READ | FILE_SHARE_WRITE, NULL, OPEN_EXISTING, 0, NULL);
}
else
{
HANDLE hDir;
DWORD dw;
DWORD buffer_size =
FIELD_OFFSET(REPARSE_DATA_BUFFER, MountPointReparseBuffer) +
MAXIMUM_REPARSE_DATA_BUFFER_SIZE;
hDir = CreateFile(MountPoint, GENERIC_READ,
FILE_SHARE_READ | FILE_SHARE_WRITE, NULL,
OPEN_EXISTING,
FILE_FLAG_BACKUP_SEMANTICS |
FILE_FLAG_OPEN_REPARSE_POINT, NULL);
if (hDir == INVALID_HANDLE_VALUE)
return INVALID_HANDLE_VALUE;
ReparseData = (PREPARSE_DATA_BUFFER)HeapAlloc(GetProcessHeap(),
HEAP_GENERATE_EXCEPTIONS | HEAP_ZERO_MEMORY,
buffer_size);
if (!DeviceIoControl(hDir, FSCTL_GET_REPARSE_POINT,
NULL, 0,
ReparseData, buffer_size,
&dw, NULL))
{
DWORD last_error = GetLastError();
CloseHandle(hDir);
HeapFree(GetProcessHeap(), 0, ReparseData);
SetLastError(last_error);
return INVALID_HANDLE_VALUE;
}
CloseHandle(hDir);
if (ReparseData->ReparseTag != IO_REPARSE_TAG_MOUNT_POINT)
{
HeapFree(GetProcessHeap(), 0, ReparseData);
SetLastError(ERROR_NOT_A_REPARSE_POINT);
return INVALID_HANDLE_VALUE;
}
DeviceName.Length =
ReparseData->MountPointReparseBuffer.SubstituteNameLength;
DeviceName.Buffer = (PWSTR)
((PUCHAR)ReparseData->MountPointReparseBuffer.PathBuffer +
ReparseData->MountPointReparseBuffer.SubstituteNameOffset);
DeviceName.MaximumLength = DeviceName.Length;
}
if (DeviceName.Buffer[(DeviceName.Length >> 1) - 1] == L'\\')
{
DeviceName.Buffer[(DeviceName.Length >> 1) - 1] = 0;
DeviceName.Length -= 2;
}
h = MyImDiskOpenDeviceByName(&DeviceName, AccessMode);
if (ReparseData != NULL)
HeapFree(GetProcessHeap(), 0, ReparseData);
return h;
}
BOOL WINAPI MyImDiskCliValidateDriveLetterTarget(LPCWSTR DriveLetter, LPCWSTR ValidTargetPath)
{
WCHAR target[MAX_PATH];
if (QueryDosDevice(DriveLetter, target, _countof(target)))
{
if (wcscmp(target, ValidTargetPath) == 0)
{
return TRUE;
}
DbgPrint(L"Drive letter %s points to %s instead of expected %s.\n", DriveLetter, target, ValidTargetPath);
}
else if (GetLastError() != ERROR_FILE_NOT_FOUND)
{
DbgPrint(L"Error verifying temporary drive letter:");
}
return FALSE;
}
//BOOL WINAPI MyImDiskGetDeviceListEx(IN ULONG ListLength, OUT ULONG *DeviceList)
//{
// UNICODE_STRING file_name;
// HANDLE driver;
// ULONG dw;
//
// RtlInitUnicodeString(&file_name, IMDISK_CTL_DEVICE_NAME);
//
// driver = MyImDiskOpenDeviceByName(&file_name, GENERIC_READ);
// if (driver == INVALID_HANDLE_VALUE)
// return 0;
//
// if (!DeviceIoControl(driver,
// IOCTL_IMDISK_QUERY_DRIVER,
// NULL, 0,
// DeviceList, ListLength << 2,
// &dw, NULL))
// {
// DWORD dwLastError = GetLastError();
// NtClose(driver);
// SetLastError(dwLastError);
// return FALSE;
// }
//
// NtClose(driver);
//
// if ((dw == sizeof(ULONG)) &
// (*DeviceList > 0))
// {
// SetLastError(ERROR_MORE_DATA);
// return FALSE;
// }
//
// SetLastError(NO_ERROR);
// return TRUE;
//}
WCHAR WINAPI MyImDiskFindFreeDriveLetter()
{
DWORD logical_drives = GetLogicalDrives();
WCHAR search;
for (search = L'Z'; search >= L'I'; search--)
{
if ((logical_drives & (1 << (search - L'A'))) == 0)
{
return search;
}
}
return 0;
}
}
#include <intrin.h>
#ifndef _M_ARM64
#include <emmintrin.h>
#endif
bool data_search_std(void *_ptr, int size)
{
unsigned char* ptr = (unsigned char*)_ptr;
long *scan_ptr;
if (!size) return FALSE;
scan_ptr = (long*)ptr;
ptr = (BYTE*)ptr + size - sizeof(long);
if (*(long*)ptr) return TRUE; // check if the last long not 0
*(long*)ptr = 1; // set last long to 1 to ensure termination
while (!*(scan_ptr++));
*(long*)ptr = 0; // restore last long to 0
return --scan_ptr != (long*)ptr;
}
#ifdef _M_ARM64
bool (*data_search)(void* ptr, int size) = data_search_std;
#else
bool data_search_sse2(void *_ptr, int size)
{
unsigned char* ptr = (unsigned char*)_ptr;
unsigned char *end_ptr;
__m128i zero;
if (!size) return FALSE;
zero = _mm_setzero_si128();
end_ptr = ptr + size - sizeof(__m128i);
if ((unsigned short)_mm_movemask_epi8(_mm_cmpeq_epi8(*(__m128i*)end_ptr, zero)) != 0xffff) return TRUE;
*end_ptr = 1;
while ((unsigned short)_mm_movemask_epi8(_mm_cmpeq_epi8(*(__m128i*)ptr, zero)) == 0xffff) ptr += sizeof(__m128i);
*end_ptr = 0;
return ptr != end_ptr;
}
bool data_search_avx(void *_ptr, int size)
{
unsigned char* ptr = (unsigned char*)_ptr;
unsigned char *end_ptr;
__m256i one;
if (!size) return FALSE;
one = _mm256_set1_epi8(0xff);
end_ptr = ptr + size - sizeof(__m256i);
if (!_mm256_testz_si256(*(__m256i*)end_ptr, one)) return TRUE;
*end_ptr = 1;
while (_mm256_testz_si256(*(__m256i*)ptr, one)) ptr += sizeof(__m256i);
*end_ptr = 0;
return ptr != end_ptr;
}
extern "C" {
void* pick_data_search()
{
bool (*ret)(void*, int);
int cpuInfo[4];
__cpuid(cpuInfo, 1);
#ifndef _WIN64
ret = data_search_std;
if (cpuInfo[3] & 0x4000000)
#endif
ret = data_search_sse2;
bool osUsesXSAVE_XRSTORE = cpuInfo[2] & (1 << 27);
bool cpuAVXSuport = cpuInfo[2] & (1 << 28);
if (osUsesXSAVE_XRSTORE && cpuAVXSuport) {
unsigned long long xcrFeatureMask = _xgetbv(_XCR_XFEATURE_ENABLED_MASK);
if ((xcrFeatureMask & 0x6) == 0x6)
ret = data_search_avx;
}
return ret;
}
}
bool (*data_search)(void* ptr, int size) = (bool (*)(void*, int))pick_data_search();
#endif
View Git Blame Copy Permalink