Driver: fortutl: Use heap in fort_reg_value()

This commit is contained in:
Nodir Temirkhodjaev 2023-04-11 13:46:06 +03:00
parent 123a0b766f
commit e9655e328f
2 changed files with 21 additions and 8 deletions

View File

@ -40,8 +40,10 @@ static NTSTATUS fort_driver_load(PDRIVER_OBJECT driver, PUNICODE_STRING reg_path
ExInitializeDriverRuntime(DrvRtPoolNxOptIn); ExInitializeDriverRuntime(DrvRtPoolNxOptIn);
status = fort_system32_path_init(driver, reg_path); status = fort_system32_path_init(driver, reg_path);
if (!NT_SUCCESS(status)) if (!NT_SUCCESS(status)) {
LOG("Driver Path Init: Error: %x\n", status);
return status; return status;
}
UNICODE_STRING device_name; UNICODE_STRING device_name;
RtlInitUnicodeString(&device_name, FORT_NT_DEVICE_NAME); RtlInitUnicodeString(&device_name, FORT_NT_DEVICE_NAME);
@ -50,13 +52,16 @@ static NTSTATUS fort_driver_load(PDRIVER_OBJECT driver, PUNICODE_STRING reg_path
PDEVICE_OBJECT device_obj; PDEVICE_OBJECT device_obj;
status = IoCreateDevice(driver, sizeof(FORT_DEVICE), &device_name, FORT_DEVICE_TYPE, 0, status = IoCreateDevice(driver, sizeof(FORT_DEVICE), &device_name, FORT_DEVICE_TYPE, 0,
/*exclusive=*/TRUE, &device_obj); /*exclusive=*/TRUE, &device_obj);
if (!NT_SUCCESS(status)) if (!NT_SUCCESS(status)) {
LOG("Create Device: Error: %x\n", status);
return status; return status;
}
device_obj->Flags |= DO_BUFFERED_IO; device_obj->Flags |= DO_BUFFERED_IO;
status = IoRegisterShutdownNotification(device_obj); status = IoRegisterShutdownNotification(device_obj);
if (!NT_SUCCESS(status)) { if (!NT_SUCCESS(status)) {
LOG("Register Shutdown: Error: %x\n", status);
fort_driver_delete_device(driver); fort_driver_delete_device(driver);
return status; return status;
} }
@ -65,8 +70,10 @@ static NTSTATUS fort_driver_load(PDRIVER_OBJECT driver, PUNICODE_STRING reg_path
RtlInitUnicodeString(&device_link, FORT_DOS_DEVICE_NAME); RtlInitUnicodeString(&device_link, FORT_DOS_DEVICE_NAME);
status = IoCreateSymbolicLink(&device_link, &device_name); status = IoCreateSymbolicLink(&device_link, &device_name);
if (!NT_SUCCESS(status)) if (!NT_SUCCESS(status)) {
LOG("Create Link: Error: %x\n", status);
return status; return status;
}
driver->DriverUnload = &fort_driver_unload; driver->DriverUnload = &fort_driver_unload;
driver->MajorFunction[IRP_MJ_CREATE] = &fort_device_create; driver->MajorFunction[IRP_MJ_CREATE] = &fort_device_create;

View File

@ -7,7 +7,7 @@
#define FORT_MAX_FILE_SIZE (4 * 1024 * 1024) #define FORT_MAX_FILE_SIZE (4 * 1024 * 1024)
#define FORT_KEY_INFO_PATH_SIZE \ #define FORT_KEY_INFO_PATH_SIZE \
(2 + (MAX_PATH * sizeof(WCHAR)) / sizeof(KEY_VALUE_FULL_INFORMATION)) (2 * sizeof(KEY_VALUE_FULL_INFORMATION) + (MAX_PATH * sizeof(WCHAR)))
static WCHAR g_system32PathBuffer[64]; static WCHAR g_system32PathBuffer[64];
static UNICODE_STRING g_system32Path; static UNICODE_STRING g_system32Path;
@ -33,11 +33,15 @@ static NTSTATUS fort_reg_value(HANDLE regKey, PUNICODE_STRING valueName, PUNICOD
{ {
NTSTATUS status; NTSTATUS status;
KEY_VALUE_FULL_INFORMATION keyInfo[FORT_KEY_INFO_PATH_SIZE]; PKEY_VALUE_FULL_INFORMATION keyInfo =
ULONG keyInfoSize; fort_mem_alloc(FORT_KEY_INFO_PATH_SIZE, FORT_UTL_POOL_TAG);
if (keyInfo == NULL)
return STATUS_INSUFFICIENT_RESOURCES;
status = ZwQueryValueKey( ULONG keyInfoSize = 0;
regKey, valueName, KeyValueFullInformation, keyInfo, sizeof(keyInfo), &keyInfoSize);
status = ZwQueryValueKey(regKey, valueName, KeyValueFullInformation, keyInfo,
FORT_KEY_INFO_PATH_SIZE, &keyInfoSize);
if (NT_SUCCESS(status)) { if (NT_SUCCESS(status)) {
const PUCHAR src = ((const PUCHAR) keyInfo + keyInfo->DataOffset); const PUCHAR src = ((const PUCHAR) keyInfo + keyInfo->DataOffset);
@ -46,6 +50,8 @@ static NTSTATUS fort_reg_value(HANDLE regKey, PUNICODE_STRING valueName, PUNICOD
status = fort_string_new(len, (PCWSTR) src, outData); status = fort_string_new(len, (PCWSTR) src, outData);
} }
fort_mem_free(keyInfo, FORT_UTL_POOL_TAG);
return status; return status;
} }