vdkctl.c

Virtual Disk Driver

DWORD VdkClose(
	HANDLE	hDevice,
	ULONG	disk_number,
	ULONG	graceful)
{
	DWORD	return_len;
	BOOL	close_device = FALSE;
	DWORD	ret = ERROR_SUCCESS;

	if (!hDevice || hDevice == INVALID_HANDLE_VALUE) {
		if ((hDevice = VdkOpenDevice(disk_number, 0)) == INVALID_HANDLE_VALUE) {
			ret = GetLastError();

			VDKTRACE(0, ("VdkClose: CreateFile(" VDK_DEV_TEMPLATE ") - %s\n",
				disk_number, 0, VdkStatusStr(ret)));

			goto cleanup;
		}

		close_device = TRUE;
	}

	if (!DeviceIoControl(hDevice, IOCTL_VDK_CLOSE_FILE,
		&graceful, sizeof(graceful), NULL, 0, &return_len, NULL)) {

		ret = GetLastError();

		VDKTRACE(0, ("VdkClose: IOCTL_VDK_CLOSE_FILE - %s\n",
			VdkStatusStr(ret)));

		goto cleanup;
	}

cleanup:

	if (close_device && hDevice != INVALID_HANDLE_VALUE) {
		CloseHandle(hDevice);
	}

	return ret;
}

//
//	Get Virtual Disk image file info
//
DWORD VdkGetFileInfo(
	HANDLE	hDevice,
	ULONG	disk_number,
	PVDK_OPEN_FILE_INFO	*file_info)
{
	ULONG	info_size;
	DWORD	return_len;
	BOOL	close_device = FALSE;
	DWORD	ret = ERROR_SUCCESS;

	if (!file_info) {
		return ERROR_INVALID_PARAMETER;
	}

	*file_info = NULL;

	//
	//	Open device
	//
	if (!hDevice || hDevice == INVALID_HANDLE_VALUE) {
		if ((hDevice = VdkOpenDevice(disk_number, 0)) == INVALID_HANDLE_VALUE) {
			ret = GetLastError();

			VDKTRACE(0, (
				"VdkGetFileInfo: CreateFile(" VDK_DEV_TEMPLATE ") - %s\n",
				disk_number, 0, VdkStatusStr(ret)));

			goto cleanup;;
		}

		close_device = TRUE;
	}

	//
	//	Query file information size
	//
	if (!DeviceIoControl(hDevice, IOCTL_VDK_QUERY_FILE_SIZE,
		NULL, 0, &info_size, sizeof(info_size), &return_len, NULL)) {

		ret = GetLastError();

		VDKTRACE(0, ("VdkGetFileInfo: DeviceIoControl(IOCTL_VDK_QUERY_FILE_SIZE) - %s\n",
			VdkStatusStr(ret)));

		goto cleanup;
	}

	if (info_size == 0) {
		goto cleanup;
	}

	//
	//	Allocate information area
	//
	if ((*file_info = (PVDK_OPEN_FILE_INFO)VdkAllocMem(info_size)) == NULL) {
		ret = GetLastError();

		VDKTRACE(0, ("VdkGetFileInfo: VdkAllocMem(%lu) - %s\n",
			info_size, VdkStatusStr(ret)));

		goto cleanup;
	}

	memset(*file_info, 0, info_size);

	//
	//	Query file information
	//
	if (!DeviceIoControl(hDevice, IOCTL_VDK_QUERY_FILE,
		NULL, 0, *file_info, info_size, &return_len, NULL)) {

		ret = GetLastError();

		VDKTRACE(0, ("VdkGetFileInfo: DeviceIoControl(IOCTL_VDK_QUERY_FILE) - %s\n",
			VdkStatusStr(ret)));

		VdkFreeMem(*file_info);
		*file_info = NULL;

		goto cleanup;
	}

cleanup:
	if (close_device && hDevice != INVALID_HANDLE_VALUE) {
		CloseHandle(hDevice);
	}

	return ret;
}

//
//	Dismount a single virtual partition
//
DWORD VdkDismount(
	HANDLE	hDevice,
	ULONG	disk_number,
	ULONG	part_number,
	ULONG	unlock)
{
	DWORD	return_len;
	BOOL	close_device = FALSE;
	DWORD	ret = ERROR_SUCCESS;

	if (!hDevice || hDevice == INVALID_HANDLE_VALUE) {
		hDevice = VdkOpenDevice(disk_number, part_number);

		if (hDevice == INVALID_HANDLE_VALUE) {
			ret = GetLastError();

			VDKTRACE(0, ("VdkDismount: CreateFile(" VDK_DEV_TEMPLATE ") - %s\n",
				disk_number, part_number, VdkStatusStr(ret)));

			return ret;
		}

		close_device = TRUE;
	}

	if (!DeviceIoControl(hDevice, FSCTL_LOCK_VOLUME,
		NULL, 0, NULL, 0, &return_len, NULL)) {

		ret = GetLastError();

		VDKTRACE(0, (
			"VdkDismount: DeviceIoControl(FSCTL_LOCK_VOLUME) - %s\n",
			VdkStatusStr(ret)));

		goto cleanup;
	}

	if (!DeviceIoControl(hDevice, FSCTL_DISMOUNT_VOLUME,
		NULL, 0, NULL, 0, &return_len, NULL)) {

		ret = GetLastError();

		VDKTRACE(0, ("VdkDismount: DeviceIoControl(FSCTL_DISMOUNT_VOLUME) - %s\n",
			VdkStatusStr(ret)));
	}

	if (ret == ERROR_SUCCESS) {
		if (!DeviceIoControl(hDevice, IOCTL_VDK_NOTIFY_DISMOUNT,
			NULL, 0, NULL, 0, &return_len, NULL)) {

			VDKTRACE(0, ("VdkDismount: DeviceIoControl(IOCTL_VDK_NOTIFY_DISMOUNT) - %s\n",
				VdkStatusStr(GetLastError())));
		}
	}

	if (unlock && !DeviceIoControl(hDevice, FSCTL_UNLOCK_VOLUME,
		NULL, 0, NULL, 0, &return_len, NULL)) {

		VDKTRACE(0, ("VdkDismount: DeviceIoControl(FSCTL_UNLOCK_VOLUME) - %s\n",
			VdkStatusStr(GetLastError())));
	}

cleanup:
	if (hDevice != INVALID_HANDLE_VALUE && close_device) {
		CloseHandle(hDevice);
	}

	return ret;
}

//
//	Dismount all virtual partition devices
//
DWORD VdkDismountAll(BOOL zombie_only)
{
	HANDLE hDevice;
	ULONG device_num = 0;
	PVDK_DEVICE_INFO device_info = NULL;
	DWORD tmp;
	DWORD ret;

	//
	//	Get information about all devices
	//
	ret = VdkGetDeviceList(&device_num, &device_info);

	if (ret != VDK_OK) {
		return ret;
	}

	//
	//	dismount existing devices
	//
	for (tmp = 0; tmp < device_num; tmp++) {
		if (!device_info[tmp].ReferenceCount ||
			(zombie_only && !device_info[tmp].Zombie)) {
			continue;
		}

		DefineDosDevice(DDD_RAW_TARGET_PATH, "VDKTMP", device_info[tmp].DeviceName);

		hDevice = CreateFile(
			"\\\\.\\VDKTMP",
			GENERIC_READ | GENERIC_WRITE,
			FILE_SHARE_READ | FILE_SHARE_WRITE,
			NULL,
			OPEN_EXISTING,
			FILE_FLAG_NO_BUFFERING,
			NULL);

		if (hDevice == INVALID_HANDLE_VALUE) {
			ret = GetLastError();

			VDKTRACE(0, ("VdkDismountAll: CreateFile(\\\\.\\VDKTMP) - %s\n",
				VdkStatusStr(ret)));
		}

		DefineDosDevice(DDD_REMOVE_DEFINITION, "VDKTMP", NULL);

		if (ret == ERROR_SUCCESS) {
			VdkDismount(hDevice, 0, 0, TRUE);
			CloseHandle(hDevice);
		}
	}

	//
	//	check if all devices are released
	//
	ret = VdkGetDriverInfo(NULL, NULL, NULL, NULL, &tmp);

	if (ret == ERROR_SUCCESS && tmp > 1) {
		ret = ERROR_BUSY;
	}

	if (device_info) {
		VdkFreeMem(device_info);
	}

	return ret;
}

//
//	Get VDK device name in the kernel namespace
//
DWORD VdkGetDeviceName(
	ULONG disk_number,
	ULONG part_number,
	LPTSTR buf)
{
	TCHAR dos_name[sizeof(VDK_DOS_TEMPLATE) + 10];
	TCHAR dir_name[MAX_DEVNAME_LEN];

	*buf = '\0';

	sprintf(dos_name, VDK_DEVICE_BASENAME "%u", disk_number);

	if (!QueryDosDevice(dos_name, dir_name, sizeof(dir_name))) {
		return GetLastError();
	}

	if (part_number != (ULONG)-1) {
		sprintf(buf, "%s\\Partition%lu", dir_name, part_number);
	}
	else {
		strcpy(buf, dir_name);
	}

	return ERROR_SUCCESS;
}

//
//	Get Virtual Disk drive letter
//
DWORD VdkGetDriveLetter(
	ULONG	disk_number,
	ULONG	part_number,
	TCHAR	*drive_letter)
{
	DWORD	logical_drives;
	TCHAR	dos_device[] = " :";
	TCHAR	device_name[MAX_PATH], dos_name[MAX_PATH];
	DWORD	ret = ERROR_SUCCESS;

	if (!drive_letter) {
		return ERROR_INVALID_PARAMETER;
	}

	*drive_letter = '\0';

	ret = VdkGetDeviceName(disk_number, part_number, device_name);

	if (ret != ERROR_SUCCESS) {
		return ret;
	}

	logical_drives = GetLogicalDrives();

	if (logical_drives == 0) {
		ret = GetLastError();

		VDKTRACE(0, (
			"VdkGetDriveLetter: GetLogicalDrives - %s\n", VdkStatusStr(ret)));

		return ret;
	}

	dos_device[0] = 'A';

	while (logical_drives) {
		if (logical_drives & 0x01) {
			if (QueryDosDevice(dos_device, dos_name, sizeof(dos_name))) {
				if (_stricmp(device_name, dos_name) == 0) {
					*drive_letter = dos_device[0];
					return ERROR_SUCCESS;
				}
			}
			else {
				ret = GetLastError();

				VDKTRACE(0, (
					"VdkGetDriveLetter: QueryDosDevice(%s) - %s\n",
					dos_device, VdkStatusStr(ret)));

				return ret;
			}
		}
		logical_drives >>= 1;
		dos_device[0]++;
	}

	return ret;
}

//
//	Assign a DOS drive letter to Virtual Disk Drive
//
DWORD VdkSetDriveLetter(
	ULONG	disk_number,
	ULONG	part_number,
	TCHAR	drive_letter)
{
	DWORD ret = ERROR_SUCCESS;
	TCHAR dos_device[] = " :";
	TCHAR device_name[MAX_PATH], dos_name[MAX_PATH];

	if (!isalpha(drive_letter)) {
		return ERROR_INVALID_PARAMETER;
	}

	//	check if device is active

	ret = VdkCheckDeviceState(NULL, disk_number, part_number);

	if (ret != ERROR_SUCCESS) {
		return ret;
	}

	//	Check if the drive letter is already in use

	dos_device[0] = (TCHAR)toupper(drive_letter);

	ret = VdkGetDeviceName(disk_number, part_number, device_name);

	if (ret != ERROR_SUCCESS) {
		return ret;
	}

	if (QueryDosDevice(dos_device, dos_name, sizeof(dos_name))) {
		if (strcmp(dos_name, device_name) == 0) {
			return ERROR_SUCCESS;
		}
		else {
			VDKTRACE(0, (
				"VdkSetDriveLetter: Drive letter '%c' is linked to '%s'\n",
				dos_device[0], dos_name));

			return ERROR_ALREADY_ASSIGNED;
		}
	}
	else {
		if ((ret = GetLastError()) != ERROR_FILE_NOT_FOUND) {
			VDKTRACE(0, (
				"VdkSetDriveLetter: QueryDosDevice(%s) - %lu\n",
				dos_device, VdkStatusStr(ret)));

			return ret;
		}
	}

	//	Check if the Virtual Disk drive has a DriveLetter already assigned

	ret = VdkGetDriveLetter(disk_number, part_number, &drive_letter);

	if (ret != ERROR_SUCCESS) {
		return ret;
	}

	if (isalpha(drive_letter)) {
		if (drive_letter == dos_device[0]) {
			return ERROR_SUCCESS;
		}
		else {
			VDKTRACE(0, ("VdkSetDriveLetter: Drive Letter '%c' is already assigned to %s\n",
				drive_letter, device_name));

			return ERROR_ALREADY_ASSIGNED;
		}
	}

	//	Assign the new drive letter

	if (!DefineDosDevice(DDD_RAW_TARGET_PATH, dos_device, device_name)) {
		ret = GetLastError();

		VDKTRACE(0, ("VdkSetDriveLetter: DefineDosDevice(%s) - %s\n",
			dos_device, VdkStatusStr(ret)));
	}

	return ret;
}

//
//	Remove Dos Drive Letter
//
DWORD VdkDelDriveLetter(
	TCHAR	drive_letter)
{
	HANDLE	hDevice;
	ULONG	tmp;
	TCHAR	dos_device[] = "\\\\.\\ :";
	DWORD	ret = ERROR_SUCCESS;

	if (!isalpha(drive_letter)) {
		return ERROR_INVALID_DRIVE;
	}

	dos_device[4] = drive_letter;

	//
	//	Check if the target device is a VDK device
	//
	hDevice = CreateFile(
		dos_device,
		GENERIC_READ,
		FILE_SHARE_READ | FILE_SHARE_WRITE,
		NULL,
		OPEN_EXISTING,
		0,
		NULL);

	if (hDevice == INVALID_HANDLE_VALUE) {
		ret = GetLastError();

		if (ret != ERROR_FILE_NOT_FOUND &&
			ret != ERROR_PATH_NOT_FOUND) {

			return ret;
		}

		goto execute;
	}

	if (!DeviceIoControl(hDevice, IOCTL_VDK_GET_VERSION,
		NULL, 0, NULL, 0, &tmp, NULL)) {

		ret = ERROR_INVALID_PARAMETER;
	}

	CloseHandle(hDevice);

	if (ret != ERROR_SUCCESS) {
		return ret;
	}

execute:
	if (!DefineDosDevice(DDD_REMOVE_DEFINITION, &dos_device[4], NULL)) {
		ret = GetLastError();

		VDKTRACE(0, (
			"VdkUnlink: DefineDosDevice(%s) - %s\n",
			&dos_device[4], VdkStatusStr(ret)));
	}

	return ret;
}

//
// choose first available drive letter
//
char ChooseDriveLetter()
{
	DWORD	logical_drives = GetLogicalDrives();
	char	drive_letter = 'C';

	if (logical_drives == 0) {
		return '\0';
	}

	//
	//	Do not assign A and B to Virtual Disk even if they are not used
	//
	logical_drives >>= 2;

	while (logical_drives & 0x1) {
		logical_drives >>= 1;
		drive_letter++;
	}

	if (drive_letter > 'Z') {
		return '\0';
	}

	return drive_letter;
}

//
//	Ensure Virtual Disk image is closed
//
DWORD VdkCloseDrive(
	ULONG			disk_number,
	VDK_CALLBACK	retrycb,
	VDK_CALLBACK	contcb,
	PVOID			param)
{
	HANDLE	hDisk;
	TCHAR	vdk_dev[MAX_PATH];
	ULONG	dos_dev[26];
	ULONG	part_num;
	ULONG	max_parts;
	ULONG	graceful;
	ULONG	i;
	DWORD	ret = ERROR_SUCCESS;

	VdkGetDeviceName(disk_number, 0, vdk_dev);

	//
	//	Delete all drive letters
	//
	for (i = 0; i < 26; i++) {
		TCHAR	dos_link[] = " :";
		TCHAR	dos_device[MAX_PATH];

		dos_link[0] = (TCHAR)(i + 'a');
		dos_device[0] = '\0';

		if (QueryDosDevice(dos_link, dos_device, sizeof(dos_device)) &&
			!_strnicmp(vdk_dev, dos_device, strlen(vdk_dev) - 1)) {

			dos_dev[i] = atol(dos_device + strlen(vdk_dev) - 1);

			DefineDosDevice(
				(DDD_RAW_TARGET_PATH | DDD_REMOVE_DEFINITION | DDD_EXACT_MATCH_ON_REMOVE),
				dos_link,
				dos_device);
		}
		else {
			dos_dev[i] = (ULONG)-1;
		}
	}

	hDisk = VdkOpenDevice(disk_number, 0);

	if (hDisk == INVALID_HANDLE_VALUE) {
		ret = VdkLastError();

		VDKTRACE(0, ("VdkCloseDrive: CreateFile(" VDK_DEV_TEMPLATE ") - %s\n",
			disk_number, 0, VdkStatusStr(ret)));

		return ret;
	}

	if (!DeviceIoControl(hDisk, IOCTL_VDK_NUMBER_OF_PARTS,
		NULL, 0, &max_parts, sizeof(max_parts), &part_num, NULL)) {

		ret = VdkLastError();

		VDKTRACE(0, ("VdkCloseDrive: IOCTL_VDK_NUMBER_OF_PARTS - %s\n",
			VdkStatusStr(ret)));

		CloseHandle(hDisk);
		return ret;
	}

	//
	//	Dismount all partitions on this disk
	//
	graceful = TRUE;

	for (part_num = 0; part_num <= max_parts; part_num++) {

		ret = VdkCheckDeviceState(
			part_num ? NULL : hDisk, disk_number, part_num);

		if (ret != ERROR_SUCCESS) {
			break;
		}

dismount_retry:

		SetCursor(LoadCursor(NULL, IDC_WAIT));

		for(i = 0;;) {
			ret = VdkDismount(
				part_num ? NULL : hDisk, disk_number, part_num, FALSE);

			if (ret != ERROR_ACCESS_DENIED || ++i > 10) {
				break;
			}

			Sleep(500);			//	retry after 0.5 sec.
		}

		SetCursor(LoadCursor(NULL, IDC_ARROW));

		if (ret == ERROR_ACCESS_DENIED &&
			retrycb && (*retrycb)(param, ret)) {

			goto dismount_retry;
		}

		if (ret != ERROR_SUCCESS) {

			// Failed to dismount
			if (contcb && (*contcb)(param, ret)) {

				//	proceed anyway
				if (!part_num) {
					graceful = FALSE;
				}

				ret = ERROR_SUCCESS;
			}
			else {

				//	abort operation
				break;
			}
		}
	}	// next partition

	//
	//	couldn't dismount partition 0 -> file is not opened
	//
	if (!part_num) {
		CloseHandle(hDisk);
		return ret;
	}

	//
	//	Close virtual disk file
	//
	if (ret == ERROR_SUCCESS) {
		ret = VdkClose(hDisk, disk_number, graceful);
	}

	if (ret != ERROR_SUCCESS) {
		//
		//	Close failed -- restore drive letters
		//
		for (i = 0; i < 26; i++) {

			if (dos_dev[i] != -1) {
				if (VdkSetDriveLetter(disk_number, dos_dev[i], (char)(i + 'a')) != ERROR_SUCCESS) {
				}
			}
		}
	}

	CloseHandle(hDisk);

	return ret;
}

// End Of File