udiskclassify.cpp

Get Product Model and Product Vendor of UDisk. Dll & test program.

/**
* File
*	UDiskClassify.cpp
* Brief
*	Implementation of CUDiskClassify.
* Author
*	Zo Yong Min
*************************************************************/


// UDiskClassify.cpp : Defines the entry point for the DLL application.
//

#include "stdafx.h"
#include <tchar.h>
#include <stdio.h>
#include "UDiskClassify.h"


#ifdef _MANAGED
#pragma managed(push, off)
#endif

BOOL APIENTRY DllMain( HMODULE hModule,
                       DWORD  ul_reason_for_call,
                       LPVOID lpReserved
					 )
{
	switch (ul_reason_for_call)
	{
	case DLL_PROCESS_ATTACH:
	case DLL_THREAD_ATTACH:
	case DLL_THREAD_DETACH:
	case DLL_PROCESS_DETACH:
		break;
	}
    return TRUE;
}

#ifdef _MANAGED
#pragma managed(pop)
#endif



CUDiskClassify::CUDiskClassify()
{
}

BOOL CUDiskClassify::GetDisksProperty(HANDLE hDevice, PSTORAGE_DEVICE_DESCRIPTOR pDevDesc)
{
	STORAGE_PROPERTY_QUERY	Query;	
	DWORD dwOutBytes;				
	BOOL bResult;					

	Query.PropertyId = StorageDeviceProperty;
	Query.QueryType = PropertyStandardQuery;

	bResult = ::DeviceIoControl(hDevice,			
		IOCTL_STORAGE_QUERY_PROPERTY,			
		&Query, sizeof(STORAGE_PROPERTY_QUERY),	
		pDevDesc, pDevDesc->Size,				
		&dwOutBytes,							
		(LPOVERLAPPED)NULL);					

	return bResult;
}

char CUDiskClassify::chFirstDriveFromMask(ULONG unitmask)
{
	char i;
	for (i=0; i<MAX_UD; ++i)  
	{
		if (unitmask & 0x1) 
			break;
		unitmask = unitmask >> 1;
	}
	return 'A' + i;
}

int CUDiskClassify::GetUDiskCount()
{
	return m_udnum;
}

UDISK_INFO *CUDiskClassify::GetUDiskInfo()
{
	return &m_udi[0];
}

void CUDiskClassify::ScanUDisk()
{
	int				k = 1;
	DWORD			MaxDriveSet = 0, CurDriveSet = 0;
	TCHAR			szBuf[300];
	HANDLE			hDevice;
	TCHAR			szDrvName[33];
	PSTORAGE_DEVICE_DESCRIPTOR pDevDesc;
	char *p;

	// initialize
	m_udnum = 0;
	ZeroMemory(m_udi, sizeof(UDISK_INFO) * MAX_UD);

	// rescan
	MaxDriveSet = GetLogicalDrives();
	CurDriveSet = MaxDriveSet;
	for (DWORD drive = 0; drive < MAX_UD; ++drive)
	{
		if ( MaxDriveSet & (1 << drive) )  
		{
			DWORD temp = 1<<drive;
			ZeroMemory(szDrvName, sizeof(szDrvName));
			_stprintf_s(szDrvName, _T("%c:\\"), 'A'+drive );
			if(GetDriveType(szDrvName)== DRIVE_REMOVABLE)
			{
				_stprintf_s(szBuf, _T("\\\\?\\%c:"), 'A'+drive);
				hDevice = CreateFile(szBuf, GENERIC_READ,
					FILE_SHARE_READ | FILE_SHARE_WRITE, NULL,
					OPEN_EXISTING, NULL, NULL);

				if (hDevice != INVALID_HANDLE_VALUE)
				{

					pDevDesc = (PSTORAGE_DEVICE_DESCRIPTOR)new BYTE[sizeof(STORAGE_DEVICE_DESCRIPTOR) + 512 - 1];
					pDevDesc->Size = sizeof(STORAGE_DEVICE_DESCRIPTOR) + 512 - 1;

					if(GetDisksProperty(hDevice, pDevDesc))
					{
						// set info 1
						m_udi[m_udnum].MoveDiskName = chFirstDriveFromMask(temp);
						p = (char*)pDevDesc;

						// set info 2
						strcpy_s(m_udi[m_udnum].VendorId, (pDevDesc->VendorIdOffset ? &p[pDevDesc->VendorIdOffset]:"(NULL)"));
						// set info 3
						strcpy_s(m_udi[m_udnum].ProductId, (pDevDesc->ProductIdOffset ? &p[pDevDesc->ProductIdOffset]:"(NULL)"));
						// set info 4
						strcpy_s(m_udi[m_udnum].ProductRevision, (pDevDesc->ProductRevisionOffset ? &p[pDevDesc->ProductRevisionOffset]:"(NULL)"));
						// set info 5
						strcpy_s(m_udi[m_udnum].SerialNumber, (pDevDesc->SerialNumberOffset ? &p[pDevDesc->SerialNumberOffset]:"(NULL)"));

						// count
						m_udnum++;
					}

					delete pDevDesc;
					CloseHandle(hDevice);
				}
			}
		}
	}
}