Hi! I am currently trying to develop a driver to list the files in a folder, but when the folder contains too many files, it runs into the BUFFER_OVERFLOW function, How can I fix this?
Thank you in advance for any response!
void* MemAlloc(size_t size, const unsigned long tag)
{
void* MemBlock = NULL;
if (KeGetCurrentIrql() <= APC_LEVEL && tag != 0)
{
MemBlock = ExAllocatePoolWithTag(PagedPool, size, tag);
if (MemBlock != NULL)
RtlZeroMemory(MemBlock, size);
}
return MemBlock;
}
void MemFree(void* Buffer, const unsigned long tag)
{
if (KeGetCurrentIrql() <= APC_LEVEL && tag != 0 && Buffer != NULL)
{
try
{
ExFreePoolWithTag(Buffer, tag);
}
__except (EXCEPTION_EXECUTE_HANDLER)
{
DbgPrint("Invalid Memory\r\n");
}
}
}
BOOLEAN IsFileDirectory(unsigned long dwAttributes)
{
return (BOOLEAN)((dwAttributes & FILE_ATTRIBUTE_DIRECTORY) == FILE_ATTRIBUTE_DIRECTORY);
}
/// <summary>
/// Lista todos os arquivos em um diretorio
/// </summary>
/// <param name="szFileDirectoryName"></param>
/// <returns></returns>
NTSTATUS ListarPasta(_In_ PUNICODE_STRING Pasta)
{
NTSTATUS Status = STATUS_SUCCESS;
// BOOLEAN
BOOLEAN Iniciou = FALSE;
// Informações
FILE_BOTH_DIR_INFORMATION* InfoPasta;
unsigned long Tamanho;
// Variáveis
HANDLE Alca;
OBJECT_ATTRIBUTES Atributos;
ANSI_STRING PastaAnsi;
UNICODE_STRING PastaUnicode;
// Inicie o ANSI e o UNICODE
RtlInitAnsiString(&PastaAnsi, Pasta);
RtlAnsiStringToUnicodeString(&PastaUnicode, &PastaAnsi, TRUE);
if (PastaUnicode.Buffer == NULL)
{
// Falha
return STATUS_UNSUCCESSFUL;
}
// IO
IO_STATUS_BLOCK IoStatusBlock;
// IQL
KIRQL kIrql = KeGetCurrentIrql();
Tamanho = (sizeof(FILE_BOTH_DIR_INFORMATION) * 2);
// Se for Passive
if (kIrql == PASSIVE_LEVEL)
{
InitializeObjectAttributes(&Atributos, &PastaUnicode, OBJ_CASE_INSENSITIVE | OBJ_KERNEL_HANDLE, NULL, NULL);
// Abra a pasta
Status = ZwOpenFile(
&Alca,
GENERIC_READ | SYNCHRONIZE,
// Atributos e IO
&Atributos,
&IoStatusBlock,
// Permissão que vamos compartilhar para os outros usuários
FILE_SHARE_READ,
// É uma pasta
FILE_OPEN_REPARSE_POINT | FILE_SYNCHRONOUS_IO_NONALERT
);
// Se não tiver sucesso
if (!NT_SUCCESS(Status) && IoStatusBlock.Information != FILE_OPENED)
{
// Pare
RtlFreeUnicodeString(&PastaUnicode);
return Status;
}
// Aloque espaço
InfoPasta = MemAlloc(Tamanho, 'JEWS');
if (InfoPasta != NULL)
{
/*
Isso fará um loop até que todos os arquivos sejam gravados no buffer,
claro, podemos obter todos os nomes de arquivos em uma chamada se nosso
buffer era grande o suficiente (se soubéssemos quantos arquivos existem no diretório)
*/
while (TRUE)
{
// Obtenha os arquivos na pasta
Status = ZwQueryDirectoryFile(Alca,
NULL,
NULL,
NULL,
&IoStatusBlock,
InfoPasta,
Tamanho,
FileBothDirectoryInformation,
FALSE,
NULL,
Iniciou
);
// Status
switch (Status)
{
// Caso falhe
case STATUS_BUFFER_OVERFLOW:
// Libere a memória
MemFree(InfoPasta, 'JEWS');
// Aumente o espaço
Tamanho *= 4;
// Aloque novamente
InfoPasta = MemAlloc(Tamanho, 'JEWS');
// Se falhar
if (InfoPasta == NULL)
{
DbgPrint("STATUS_BUFFER_OVERFLOW - 2x");
RtlFreeUnicodeString(&PastaUnicode);
return Status;
}
// Se não tiver pasta
case STATUS_NO_MORE_FILES:
DbgPrint("STATUS_NO_MORE_FILES");
RtlFreeUnicodeString(&PastaUnicode);
CleanupRoutine(Alca, InfoPasta, 'JEWS');
return STATUS_SUCCESS;
// Se for sucesso
case STATUS_SUCCESS:
{
if (Iniciou)
Iniciou = FALSE;
while (TRUE)
{
// Se for NULL
if ((InfoPasta->FileAttributes & FILE_ATTRIBUTE_DIRECTORY) &&
(InfoPasta->FileName)[0] == L'.') break;
// Nome do arquivo completo
UCHAR* NomeChar = MemAlloc(InfoPasta->FileNameLength * sizeof(WCHAR), 'KIKE');
if (NomeChar)
{
// Copie
RtlCopyMemory(NomeChar, InfoPasta->FileName, InfoPasta->FileNameLength);
PCHAR NomeCompleto = MemAlloc(1024, 'KIKE');
// Se foi alocado
if (NomeCompleto)
{
// Converta
sprintf(
NomeCompleto,
"%s%S\r\n",
Pasta,
NomeChar
);
/*
// Se for um arquivo
if (IsFileDirectory(InfoPasta->FileAttributes))
{
DbgPrint("File allocation buffer\r\n");
}
*/
//KdPrint((NomeArquivo));
EscreverNoArquivo(ARQUIVOS_LOG, NomeCompleto, FILE_APPEND_DATA);
MemFree(NomeCompleto, 'KIKE');
}
MemFree(NomeChar, 'KIKE');
}
if (InfoPasta->NextEntryOffset == 0)
break;
InfoPasta += InfoPasta->NextEntryOffset;
}
}
break;
default:
DbgPrint("DEFAULT");
CleanupRoutine(Alca, InfoPasta, 'JEWS');
RtlFreeUnicodeString(&PastaUnicode);
return Status;
}
}
}
CleanupRoutine(Alca, InfoPasta, 'JEWS');
RtlFreeUnicodeString(&PastaUnicode);
}
// Termine
return Status;
}