Buffer Overflow in ZwQueryDirectoryFile

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;
}

Google found a good answer by Craig:

https://stackoverflow.com/a/68991525