vdkcmd.c

Virtual Disk Driver

/*
	vdkcmd.c

	Virtual Disk drive control program (console version)
	Copyright (C) 2003 Ken Kato
*/

#include "vdkbase.h"
#include "vdkutil.h"
#include "vdkfile.h"
#include "vdkioctl.h"
#include "vdkaccess.h"
#include "vdkpart.h"
#include "vdkmsg.h"
#include "vdkver.h"
#include "vdkctl.h"

#include "cowdisk.h"
#include "vmdisk.h"

#include "VDiskUtil.h"

//
//	current driver state
//
static DWORD driver_state = VDK_NOT_INSTALLED;

//
//	command processing functions
//
typedef int (*cmdfnc)(char **args);

static int	Install(char **args);
static int	Remove(char **args);
static int	Start(char **args);
static int	Stop(char **args);
static int	Driver(char **args);
static int	DiskNum(char **args);
static int	Create(char **args);
static int	Delete(char **args);
static int	Device(char **args);
static int	Open(char **args);
static int	Close(char **args);
static int	Link(char **args);
static int	Unlink(char **args);
static int	Image(char **args);
static int	View(char **args);
static int	Help(char **args);
#ifdef _DEBUG
static int	Trace(char **args);
#endif	// _DEBUG

static const struct {
	char	*cmd;
	int		min_argc;
	int		max_argc;
	cmdfnc	func;
	DWORD	usage;
	DWORD	helpmsg;
}
Commands[] = {
	{ "INSTALL",	2, 4, Install,	MSG_USAGE_INSTALL	, MSG_HELP_INSTALL	},
	{ "REMOVE",		2, 2, Remove,	MSG_USAGE_REMOVE	, MSG_HELP_REMOVE	},
	{ "START",		2, 2, Start,	MSG_USAGE_START		, MSG_HELP_START	},
	{ "STOP",		2, 2, Stop,		MSG_USAGE_STOP		, MSG_HELP_STOP		},
	{ "DRIVER",		2, 3, Driver,	MSG_USAGE_DRIVER	, MSG_HELP_DRIVER	},
	{ "DISK",		3, 3, DiskNum,	MSG_USAGE_DISK		, MSG_HELP_DISK		},
	{ "CREATE",		2, 2, Create,	MSG_USAGE_CREATE	, MSG_HELP_CREATE	},
	{ "DELETE",		2, 2, Delete,	MSG_USAGE_DELETE	, MSG_HELP_DELETE	},
	{ "VIEW",		3, 3, View,		MSG_USAGE_VIEW		, MSG_HELP_VIEW		},
	{ "OPEN",		4, 7, Open,		MSG_USAGE_OPEN		, MSG_HELP_OPEN		},
	{ "CLOSE",		3, 4, Close,	MSG_USAGE_CLOSE		, MSG_HELP_CLOSE	},
	{ "LINK",		4, 5, Link,		MSG_USAGE_LINK		, MSG_HELP_LINK		},
	{ "ULINK",		3, 4, Unlink,	MSG_USAGE_ULINK		, MSG_HELP_ULINK	},
	{ "IMAGE",		2, 3, Image,	MSG_USAGE_IMAGE		, MSG_HELP_IMAGE	},
	{ "HELP",		2, 3, Help,		MSG_USAGE_HELP		, MSG_HELP_HELP		},
#ifdef _DEBUG
	{ "DEVICE",		2, 2, Device,	0, 0 },
	{ "TRACE",		2, 20, Trace,	0, 0 },
#endif	// _DEBUG
	{ NULL,			0, 0, NULL,		0, 0 }
};

//
//	local functions
//
static BOOL IsUserAdmin();
static DWORD GetDiskFromDrive(CHAR drive, PULONG disk);
static void PrintFileInfo(PVDK_OPEN_FILE_INFO file_info);
static DWORD Retry_Callback(PVOID param, DWORD err);
static DWORD Continue_Callback(PVOID param, DWORD err);
static void PartList_Callback(PPARTITION_ITEM pitem, PVOID param);
static void Assign_Callback(PPARTITION_ITEM pitem, PVOID param);

//
//	struct used for Assign_Callback parameter
//
typedef struct _ASSIGN_PARAM {
	ULONG	disk_number;
	ULONG	part_number;
	BOOL	read_only;
	PCHAR	drive_letters;
}
ASSIGN_PARAM, *PASSIGN_PARAM;

#ifdef VDK_DEBUG
#define VDK_DEBUG_TAG	" (debug)"
#else
#define VDK_DEBUG_TAG
#endif

//
//	main
//
int main(int argc, char **argv)
{
	char	*cmd;
	DWORD	ret;
	int		idx;

	//
	//	Reports memory leaks at process termination
	//
#ifdef VDK_DEBUG
	_CrtSetDbgFlag(_CrtSetDbgFlag(_CRTDBG_REPORT_FLAG) | _CRTDBG_LEAK_CHECK_DF);
#endif

	printf(
		VDK_PRODUCT_NAME " version "
		VDK_PRODUCT_VERSION_STR VDK_DEBUG_TAG "\n"
		"http://chitchat.at.infoseek.co.jp/vmware/\n\n"
	);

	if (!IsUserAdmin()) {
		PrintMessage(MSG_MUST_BE_ADMIN);
		return -1;
	}

	//
	//	At least one parameter (command) is required
	//
	if (argc < 2) {
		PrintMessage(MSG_HELP_USAGE);
		return -1;
	}

	//
	//	Get Current Driver state
	//
	if ((ret = VdkGetDriverState(&driver_state)) != ERROR_SUCCESS) {
		PrintMessage(MSG_GET_STAT_NG);
		printf("%s\n", VdkStatusStr(ret));
		return -1;
	}

	//
	//	Try to dismount all zombie devices
	//
	VdkDismountAll(TRUE);

	//
	//	Process vdk command parameter
	//
	cmd = *(argv + 1);

	idx = 0;

	while (Commands[idx].cmd) {
		if (_stricmp(cmd, Commands[idx].cmd) == 0) {

			if (argc < Commands[idx].min_argc ||
				argc > Commands[idx].max_argc) {
				PrintMessage(Commands[idx].usage);
				return -1;
			}

			ret = (*Commands[idx].func)(argv + 2);

			return ret;
		}

		idx++;
	}

	//
	//	unknown command
	//
	PrintMessage(MSG_UNKNOWN_COMMAND, cmd);

	return -1;
}

//
//	Installs the Virtual Disk Driver Dynamically
//	Command Line Parameters:
//	vdk install [path] [/auto]
//	(optional) driver file path	- default to executive's dir
//	(optional) auto start switch - default to demand start
//
int	Install(char **args)
{
	char *install_path = NULL;
	BOOL auto_start = FALSE;

	DWORD ret;

	//	process parameters

	while (*args) {
		if (_stricmp(*args, "/auto") == 0) {
			if (auto_start) {
				PrintMessage(MSG_DUPLICATE_ARGS, *args);
				return -1;
			}
			else {
				auto_start = TRUE;
			}
		}
		else if (**args == '/') {
			PrintMessage(MSG_UNKNOWN_OPTION, *args, "INSTALL");
			return -1;
		}
		else {
			if (install_path) {
				PrintMessage(MSG_DUPLICATE_ARGS, "path");
				return -1;
			}
			else {
				install_path = *args;
			}
		}
		args++;
	}

	//	already installed?

	if (driver_state != VDK_NOT_INSTALLED) {
		PrintMessage(MSG_ALREADY_INSTALLED);
		return -1;
	}

	//	install the driver

	if ((ret = VdkInstall(install_path, auto_start)) != ERROR_SUCCESS) {
		PrintMessage(MSG_INSTALL_NG);
		printf("%s\n", VdkStatusStr(ret));
		return -1;
	}

	//	operation successfull

	PrintMessage(MSG_INSTALL_OK);

	return 0;
}

//
//	Remove Virtual Disk Driver from system
//	Command Line Parameters: None
//
int	Remove(char **args)
{
	DWORD	ret;
	int		i;

	UNREFERENCED_PARAMETER(args);

	//	ensure the driver is installed

	if (driver_state == VDK_NOT_INSTALLED) {
		PrintMessage(MSG_NOT_INSTALLED);
		return -1;
	}

	//	ensure the driver is stopped

	if (driver_state != SERVICE_STOPPED) {
		if (Stop(NULL) == -1) {
			return -1;
		}
	}

	//	remove the driver

	if ((ret = VdkRemove()) != ERROR_SUCCESS) {
		PrintMessage(MSG_REMOVE_NG);
		printf("%s\n", VdkStatusStr(ret));
		return -1;
	}

	// Wait for the driver to be actually removed for 3 secs Max.

	for (i = 0; i < 10; i++) {
		ret = VdkGetDriverState(&driver_state);

		if (ret != ERROR_SUCCESS) {
			PrintMessage(MSG_GET_STAT_NG);
			printf("%s\n", VdkStatusStr(ret));
			return -1;
		}

		if (driver_state == VDK_NOT_INSTALLED) {
			break;
		}

		Sleep(300);
	}

	if (driver_state == VDK_NOT_INSTALLED) {
		PrintMessage(MSG_REMOVE_OK);
	}
	else {
		PrintMessage(MSG_REMOVE_PENDING);
	}

	return 0;
}

//
//	Start the Virtual Disk Driver
//	Command Line Parameters: None
//
int	Start(char **args)
{
	DWORD	ret;

	UNREFERENCED_PARAMETER(args);

	//	ensure that the driver is installed

	if (driver_state == VDK_NOT_INSTALLED) {
		if ((ret = VdkInstall(NULL, FALSE)) == ERROR_SUCCESS) {
			PrintMessage(MSG_INSTALL_OK);
		}
		else {
			PrintMessage(MSG_INSTALL_NG);
			printf("%s\n", VdkStatusStr(ret));
			return -1;
		}
	}

	//	ensure that the driver is not started yet

	if (driver_state == SERVICE_RUNNING) {
		PrintMessage(MSG_ALREADY_RUNNING);
		return -1;
	}

	//	start the driver

	if ((ret = VdkStart(&driver_state)) != ERROR_SUCCESS) {
		PrintMessage(MSG_START_NG);
		printf("%s\n", VdkStatusStr(ret));
		return -1;
	}

	//	operation successfull

	PrintMessage(MSG_START_OK);

	return 0;
}

//
//	Stop the Virtual Disk Driver
//	Command Line Parameters: None
//
int	Stop(char **args)
{
	DWORD ret;

	UNREFERENCED_PARAMETER(args);

	//	ensure that the driver is installed

	if (driver_state == VDK_NOT_INSTALLED) {
		PrintMessage(MSG_NOT_INSTALLED);
		return -1;
	}

	//	ensure that the driver is running

	if (driver_state == SERVICE_STOPPED) {
		PrintMessage(MSG_NOT_RUNNING);
		return -1;
	}

	//	ensure that all image are unmounted

	if (driver_state == SERVICE_RUNNING) {
		char *param[] = { "*", NULL };

		if (Close(param) == -1) {
			return -1;
		}
	}

	ret = VdkDismountAll(FALSE);

	if (ret != ERROR_SUCCESS) {
		PrintMessage(MSG_STOP_NG);
		PrintMessage(MSG_DRIVE_IN_USE);
		return -1;
	}

	//	stop the driver

	if ((ret = VdkStop(&driver_state)) != ERROR_SUCCESS) {
		PrintMessage(MSG_STOP_NG);
		printf("%s\n", VdkStatusStr(ret));
		return -1;
	}

	if (driver_state != SERVICE_STOPPED) {
		PrintMessage(MSG_STOP_PENDING);
		PrintMessage(MSG_RESTART_WARN);
		return -1;
	}

	PrintMessage(MSG_STOP_OK);

	return 0;
}

//
//	Show current driver status
//	Command Line Parameters: None
//
int	Driver(char **args)
{
	TCHAR	path[MAX_PATH];
	DWORD	start_type;
	ULONG	initial_dev;
	ULONG	version;
	ULONG	disk_device;
	ULONG	attached_part;
	ULONG	orphaned_part;
	ULONG	reference_count;
	DWORD	ret;

	// if the driver is not installed, no more info to show

	if (driver_state == VDK_NOT_INSTALLED) {
		PrintMessage(MSG_CURRENT_STATUS);
		PrintMessage(MSG_STATUS_NOT_INST);
		return 0;
	}

	//	get current driver config

	if ((ret = VdkGetDriverConfig(path, &start_type, &initial_dev)) != ERROR_SUCCESS) {
		PrintMessage(MSG_GET_CONFIG_NG);
		printf("%s\n", VdkStatusStr(ret));
		return -1;
	}

	//	print driver file path

	PrintMessage(MSG_DRIVER_LOCATION, path);

	//	print running driver version
	version = 0;

	if (driver_state == SERVICE_RUNNING) {
		ret = VdkCheckVersion(NULL, &version);
	}
	else {
		ret = VdkCheckFileVersion(path, &version);
	}

	if (ret != ERROR_SUCCESS &&
		ret != ERROR_REVISION_MISMATCH) {
		PrintMessage(MSG_GET_CONFIG_NG);
		printf("%s\n", VdkStatusStr(ret));
		return -1;
	}

	if (version) {
		PrintMessage(MSG_DRIVER_VERSION,
			HIWORD(version),
			(version & 0x00007fff),
			(version & 0x00008000) ? " (debug)" : "");
	}

	//	print driver start type

	PrintMessage(MSG_START_TYPE);

	switch (start_type) {
	case SERVICE_AUTO_START:
		PrintMessage(MSG_START_AUTO);
		break;

	case SERVICE_BOOT_START:
		PrintMessage(MSG_START_BOOT);
		break;

	case SERVICE_DEMAND_START:
		PrintMessage(MSG_START_DEMAND);
		break;

	case SERVICE_DISABLED:
		PrintMessage(MSG_START_DISABLED);
		break;

	case SERVICE_SYSTEM_START :
		PrintMessage(MSG_START_SYSTEM);
		break;

	default:
		PrintMessage(MSG_UNKNOWN_ULONG, start_type);
		break;
	}

	//	print current status

	PrintMessage(MSG_CURRENT_STATUS);

	switch (driver_state) {
	case SERVICE_STOPPED:
		PrintMessage(MSG_STATUS_STOPPED);
		break;

	case SERVICE_START_PENDING:
		PrintMessage(MSG_STATUS_START_P);
		break;

	case SERVICE_STOP_PENDING:
		PrintMessage(MSG_STATUS_STOP_P);
		break;

	case SERVICE_RUNNING:
		PrintMessage(MSG_STATUS_RUNNING);
		break;

	case SERVICE_CONTINUE_PENDING:
		PrintMessage(MSG_STATUS_CONT_P);
		break;

	case SERVICE_PAUSE_PENDING:
		PrintMessage(MSG_STATUS_PAUSE_P);
		break;

	case SERVICE_PAUSED:
		PrintMessage(MSG_STATUS_PAUSED);
		break;

	default:
		PrintMessage(MSG_UNKNOWN_ULONG, driver_state);
		break;
	}

	if (version < VDK_DRIVER_VERSION_VAL) {
		return 0;
	}

	//	if driver is not running, no more info

	if (driver_state != SERVICE_RUNNING) {
		PrintMessage(MSG_DISK_DEVICE, initial_dev);
		return 0;
	}

	ret = VdkGetDriverInfo(NULL,
		&disk_device,
		&attached_part,
		&orphaned_part,
		&reference_count);

	if (ret != ERROR_SUCCESS) {
		PrintMessage(MSG_GET_CONFIG_NG);
		printf("%s\n", VdkStatusStr(ret));
		return -1;
	}

	PrintMessage(MSG_DISK_DEVICE, disk_device);

	//
	//	hidden option
	//
	if (*args && !_stricmp(*args, "full")) {
		PrintMessage(MSG_ATTACHED_PART, attached_part);
		PrintMessage(MSG_ORPHANED_PART, orphaned_part);
		PrintMessage(MSG_REFERENCE_COUNT, reference_count);
	}

	return 0;
}

//
//	Set initial number of disk devices
//	Command Line Parameters:
//		device number
//
int DiskNum(char **args)
{
	ULONG	dev_num;
	DWORD	ret;

	//	ensure that the driver is installed

	if (driver_state == VDK_NOT_INSTALLED) {
		PrintMessage(MSG_NOT_INSTALLED);
		return -1;
	}

	if (!isdigit(**args)) {
		PrintMessage(MSG_INVALID_DISKNUM, *args);
		return -1;
	}

	dev_num = atol(*args);

	if (dev_num == 0 || dev_num > VDK_MAXIMUM_DISK_NUM) {
		PrintMessage(MSG_INVALID_DISKNUM, *args);
		return -1;
	}

	ret = VdkSetDeviceNum(dev_num);

	if (ret != ERROR_SUCCESS) {
		PrintMessage(MSG_DISKNUM_NG);
		printf("%s\n", VdkStatusStr(ret));
		return -1;
	}

	PrintMessage(MSG_DISKNUM_OK, dev_num);

	return 0;
}

//
//	Create a new disk device
//	Command Line Parameters: none
//
int Create(char **args)
{
	ULONG disk_device;
	DWORD ret;

	UNREFERENCED_PARAMETER(args);

	//	ensure that the driver is installed

	if (driver_state == VDK_NOT_INSTALLED) {
		PrintMessage(MSG_NOT_INSTALLED);
		return -1;
	}

	//	ensure that the driver is running

	if (driver_state != SERVICE_RUNNING) {
		PrintMessage(MSG_NOT_RUNNING);
		return -1;
	}

	//	get current number of disk devices

	ret = VdkGetDriverInfo(NULL, &disk_device, NULL, NULL, NULL);

	if (ret != ERROR_SUCCESS) {
		PrintMessage(MSG_GET_CONFIG_NG);
		printf("%s\n", VdkStatusStr(ret));
		return -1;
	}

	if (disk_device >= VDK_MAXIMUM_DISK_NUM) {
		PrintMessage(MSG_NO_MORE_DISK);
		return -1;
	}

	//	create a disk

	ret = VdkCreateDisk(NULL);

	if (ret != ERROR_SUCCESS) {
		PrintMessage(MSG_CREATEDISK_NG, disk_device);
		printf("%s\n", VdkStatusStr(ret));
		return -1;
	}

	PrintMessage(MSG_CREATEDISK_OK, disk_device);

	//	print new number of disk devices

	ret = VdkGetDriverInfo(NULL, &disk_device, NULL, NULL, NULL);

	if (ret != ERROR_SUCCESS) {
		PrintMessage(MSG_GET_CONFIG_NG);