// // DISK.CPP - Classes for partition table manipulation // // Revision History: // #include "disk.h" #include #include #include #include HRESULT CDrive::Initialize( LPCTSTR lpszLogicalDrive ) { HANDLE handle; HRESULT hr; DISK_GEOMETRY geom; ULONG bps; PVOID unalignedBuffer; PVOID buffer; /* * Open the disk device */ handle = LowOpenPartition(lpszLogicalDrive); if (handle == INVALID_HANDLE_VALUE) { return E_INVALIDARG; } /* * Get the geometry */ hr = LowGetGeometry(handle, &geom); if (FAILED(hr)) { CloseHandle(handle); return hr; } m_geometry.cylinders = geom.Cylinders.QuadPart; m_geometry.mediaType = geom.MediaType; m_geometry.tracksPerCylinder = geom.TracksPerCylinder; m_geometry.sectorsPerTrack = geom.SectorsPerTrack; m_geometry.bytesPerSector = geom.BytesPerSector; m_geometry.bytesPerTrack = m_geometry.sectorsPerTrack * m_geometry.bytesPerSector; m_geometry.bytesPerCylinder = m_geometry.tracksPerCylinder * m_geometry.bytesPerTrack; m_geometry.totalSectorCount = (ULONGLONG)(m_geometry.cylinders * m_geometry.bytesPerCylinder); m_length = m_geometry.cylinders * m_geometry.bytesPerCylinder; if (m_length == 0) { // Probably no media is present in the drive return E_INVALIDARG; } /* * Get the true length of the drive */ hr = LowGetLength(handle, &m_trueLength); if (FAILED(hr)) { CloseHandle(handle); return hr; } /* * Check whether this is a NEC98 disk */ m_isNEC98 = FALSE; bps = m_geometry.bytesPerSector; unalignedBuffer = (PVOID) new char[2 * bps]; if (!unalignedBuffer) { CloseHandle(handle); return E_OUTOFMEMORY; } buffer = (PVOID) (((ULONG_PTR)unalignedBuffer + bps) & ~((ULONG_PTR)(bps - 1))); hr = LowReadSectors(handle, bps, 0, 1, buffer); if (FAILED(hr)) { delete [] (char*) unalignedBuffer; CloseHandle(handle); return hr; } if (IsNEC_98) { if (((unsigned char *)buffer)[0x1fe] == 0x55 && ((unsigned char *)buffer)[0x1ff] == 0xaa) { if (((unsigned char *)buffer)[4] == 'I' && ((unsigned char *)buffer)[5] == 'P' && ((unsigned char *)buffer)[6] == 'L' && ((unsigned char *)buffer)[7] == '1') { m_isNEC98 = TRUE; } } else { m_isNEC98 = TRUE; } } delete [] (char*) unalignedBuffer; /* * We have all the information we need. Return. */ CloseHandle(handle); return S_OK; } CDrive::~CDrive( ) { } HRESULT CDrive::ReadBootRecord( LPCTSTR lpszLogicalDrive, UINT nSectors, PBYTE *buffer ) { HANDLE hPartition; HRESULT hr; *buffer = new BYTE[m_geometry.bytesPerSector * nSectors]; // Do disk ops, read bootcode // hPartition = LowOpenPartition(lpszLogicalDrive); if ( hPartition == INVALID_HANDLE_VALUE ) { delete[] *buffer; *buffer = NULL; throw new W32Error(); } hr = LowReadSectors(hPartition, m_geometry.bytesPerSector, 0, nSectors, *buffer); if ( S_OK != hr ) { delete[] *buffer; *buffer = NULL; CloseHandle( hPartition ); throw new W32Error(); } CloseHandle( hPartition ); // The calling function is responsible for cleaning up the buffer. // return hr; } HRESULT CDrive::WriteBootRecord( LPCTSTR lpszLogicalDrive, UINT nSectors, PBYTE *buffer ) { HANDLE hPartition; HRESULT hr; UINT *uiBackupSector = NULL; // Do disk ops, write bootcode // hPartition = LowOpenPartition(lpszLogicalDrive); if ( INVALID_HANDLE_VALUE == hPartition ) { throw new W32Error(); } // Figure out where the backup boot sector is. It is at offset 0x32 in the boot record. // uiBackupSector = (UINT *) &((*buffer)[0x32]); hr = LowWriteSectors(hPartition, m_geometry.bytesPerSector, 0, nSectors, *buffer); if ( S_OK != hr ) { CloseHandle(hPartition); throw new W32Error(); } if ( uiBackupSector ) { hr = LowWriteSectors(hPartition, m_geometry.bytesPerSector, *uiBackupSector, nSectors, *buffer); } if ( S_OK != hr ) { CloseHandle(hPartition); throw new W32Error(); } CloseHandle(hPartition); return hr; } HRESULT CDrive::WriteBootRecordXP( LPCTSTR lpszLogicalDrive, UINT nSectors, PBYTE *buffer ) { HANDLE hPartition; HRESULT hr; UINT *uiBackupSector = NULL; // Do disk ops, write bootcode // hPartition = LowOpenPartition(lpszLogicalDrive); if ( INVALID_HANDLE_VALUE == hPartition ) { throw new W32Error(); } // Figure out where the backup boot sector is. It is at offset 0x32 in the boot record. // uiBackupSector = (UINT *) &((*buffer)[0x32]); // Write out the first 2 sectors to sectors 0 and 1 on the disk. We will write the // third sector to sector 12. // hr = LowWriteSectors(hPartition, m_geometry.bytesPerSector, 0, nSectors - 1, *buffer); if ( S_OK != hr ) { CloseHandle(hPartition); throw new W32Error(); } // For NT we need to write out the third sector of the bootcode to sector 12. // hr = LowWriteSectors(hPartition, m_geometry.bytesPerSector, 12, 1, *buffer + (2 * m_geometry.bytesPerSector)); if ( S_OK != hr ) { CloseHandle(hPartition); throw new W32Error(); } if ( uiBackupSector ) { hr = LowWriteSectors(hPartition, m_geometry.bytesPerSector, *uiBackupSector, nSectors, *buffer); } if ( S_OK != hr ) { CloseHandle(hPartition); throw new W32Error(); } CloseHandle(hPartition); return hr; } HANDLE LowOpenDisk( ULONG diskNumber ) { HANDLE handle = NULL; int err = 0; int i = 0; WCHAR buffer[64]; swprintf(buffer, L"\\\\.\\PHYSICALDRIVE%lu", diskNumber); for ( i = 0; i < 5; i++ ) { handle = CreateFile( buffer, GENERIC_READ | GENERIC_WRITE, FILE_SHARE_DELETE | FILE_SHARE_READ | FILE_SHARE_WRITE, NULL, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL); if (handle == INVALID_HANDLE_VALUE) { err = GetLastError(); if (err == ERROR_SHARING_VIOLATION) Sleep(2000); else break; } else break; } return handle; } HANDLE LowOpenPartition( ULONG diskNumber, ULONG partitionNumber ) { WCHAR buffer[64]; swprintf(buffer, L"\\\\?\\GLOBALROOT\\Device\\Harddisk%lu\\Partition%lu", diskNumber, partitionNumber); return CreateFile( buffer, GENERIC_READ | GENERIC_WRITE, FILE_SHARE_DELETE, NULL, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL); } HANDLE LowOpenPartition( LPCTSTR lpszLogicalDrive ) { return CreateFile( lpszLogicalDrive, GENERIC_READ | GENERIC_WRITE, FILE_SHARE_DELETE, NULL, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL); } HRESULT LowGetGeometry( HANDLE handle, PDISK_GEOMETRY geometry ) { ULONG size; if (!DeviceIoControl( handle, IOCTL_DISK_GET_DRIVE_GEOMETRY, NULL, 0, geometry, sizeof(DISK_GEOMETRY), &size, NULL)) { return E_FAIL; } return S_OK; } HRESULT LowGetLength( HANDLE handle, PLONGLONG length ) { PARTITION_INFORMATION_EX partInfoEx; PARTITION_INFORMATION partInfo; ULONG size; /* * Try first the new ioctl */ if (DeviceIoControl( handle, IOCTL_DISK_GET_PARTITION_INFO_EX, NULL, 0, &partInfoEx, sizeof(PARTITION_INFORMATION_EX), &size, NULL)) { *length = partInfoEx.PartitionLength.QuadPart; return S_OK; } /* * For Win2K systems we should use the old ioctl */ if (DeviceIoControl( handle, IOCTL_DISK_GET_PARTITION_INFO, NULL, 0, &partInfo, sizeof(PARTITION_INFORMATION), &size, NULL)) { *length = partInfo.PartitionLength.QuadPart; return S_OK; } return E_FAIL; } HRESULT LowReadSectors( HANDLE handle, ULONG sectorSize, ULONG startingSector, ULONG numberOfSectors, PVOID buffer ) { IO_STATUS_BLOCK statusBlock; LARGE_INTEGER byteOffset; byteOffset.QuadPart = UInt32x32To64(startingSector, sectorSize); statusBlock.Status = 0; statusBlock.Information = 0; if (!NT_SUCCESS(NtReadFile( handle, 0, NULL, NULL, &statusBlock, buffer, numberOfSectors * sectorSize, &byteOffset, NULL))) { return E_FAIL; } return S_OK; } HRESULT LowWriteSectors( HANDLE handle, ULONG sectorSize, ULONG startingSector, ULONG numberOfSectors, PVOID buffer ) { IO_STATUS_BLOCK statusBlock; LARGE_INTEGER byteOffset; byteOffset.QuadPart = UInt32x32To64(startingSector, sectorSize); statusBlock.Status = 0; statusBlock.Information = 0; if (!NT_SUCCESS(NtWriteFile( handle, 0, NULL, NULL, &statusBlock, buffer, numberOfSectors * sectorSize, &byteOffset, NULL))) { return E_FAIL; } return S_OK; } HRESULT LowFsLock( HANDLE handle ) { ULONG size; if (!DeviceIoControl( handle, FSCTL_LOCK_VOLUME, NULL, 0, NULL, 0, &size, NULL)) { return E_FAIL; } return S_OK; } HRESULT LowFsUnlock( HANDLE handle ) { ULONG size; if (!DeviceIoControl( handle, FSCTL_UNLOCK_VOLUME, NULL, 0, NULL, 0, &size, NULL)) { return E_FAIL; } return S_OK; } HRESULT LowFsDismount( HANDLE handle ) { ULONG size; if (!DeviceIoControl( handle, FSCTL_DISMOUNT_VOLUME, NULL, 0, NULL, 0, &size, NULL)) { return E_FAIL; } return S_OK; } LONGLONG RoundUp( LONGLONG value, LONGLONG factor ) /* * Rounds a value up to a multiple of a given number */ { // This is the most common case so treat it separately if (value % factor == 0) { return value; } // And this is the generic formula return ((LONGLONG)((value + factor - 1) / factor)) * factor; } LONGLONG RoundDown( LONGLONG value, LONGLONG factor ) /* * Rounds a value down to a multiple of a given number */ { // This is the most common case so treat it separately if (value % factor == 0) { return value; } //And this the generic formula return ((LONGLONG)(value / factor)) * factor; }