windows-nt/Source/XPSP1/NT/ds/security/base/keymgr/diskio.cpp
2020-09-26 16:20:57 +08:00

446 lines
13 KiB
C++

// Dependencies: shellapi.h, shell32.lib for SHGetFileInfo()
// windows.h, kernel32.lib for GetDiskFreeSpace()
// io.h for _waccess()
#include <windef.h>
#include <windows.h>
#include <string.h>
#include <io.h>
#include <stdio.h>
//#include <shellapi.h>
#include <shlwapi.h>
//#include <shlobjp.h>
#include "switches.h"
#include "wizres.h"
extern HINSTANCE g_hInstance;
#if !defined(SHFMT_OPT_FULL)
#if defined (__cplusplus)
extern "C" {
#endif
DWORD WINAPI SHFormatDrive(HWND,UINT,UINT,UINT);
#define SHFMT_ID_DEFAULT 0xffff
#define SHFMT_OPT_FULL 0x0001
#define SHFMT_OPT_SYSONLY 0x0002
#define SHFMT_ERROR 0xffffffffL
#define SHFMT_CANCEL 0xfffffffeL
#define SHFMT_NOFORMAT 0xffffffdL
#if defined (__cplusplus)
}
#endif
#endif
// Miscellaneous declarations not contain in header files
// These will be miscellaneous items found in other files within this project
int RMessageBox(HWND hw,UINT_PTR uiResIDTitle, UINT_PTR uiResIDText, UINT uiType);
extern HWND c_hDlg;
extern WCHAR pszFileName[];
INT g_iFileSize = 0;
INT g_iBufferSize = 0;
INT g_iSectorSize = 0;
HANDLE g_hFile = NULL;
BOOL GetFileSize(WCHAR *pszFilePath,INT *icbSize)
{
WIN32_FILE_ATTRIBUTE_DATA stWFAD = {0};
if (NULL == pszFilePath) return FALSE;
if (NULL == icbSize) return FALSE;
if (!GetFileAttributesEx(pszFilePath,GetFileExInfoStandard,&stWFAD)) return FALSE;
#ifdef LOUDLY
#ifdef LOUDLY
WCHAR rgc[100];
swprintf(rgc,L"GetFileSize returns %d\n",stWFAD.nFileSizeLow);
OutputDebugString(rgc);
#endif
#endif
*icbSize = stWFAD.nFileSizeLow;
return TRUE;
}
DWORD GetDriveFreeSpace(WCHAR *pszFilePath)
{
WCHAR *pwc;
WCHAR rgcModel[]={L"A:"};
DWORD dwSpc,dwBps,dwCfc,dwTcc,dwFree;
if (NULL == pszFilePath) return 0;
rgcModel[0] = *pszFilePath;
if (!GetDiskFreeSpace(rgcModel,&dwSpc,&dwBps,&dwCfc,&dwTcc))
{
#ifdef LOUDLY
WCHAR rgwc[100];
swprintf(rgwc,L"GetDriveFreeSpace encountered error %x\n",GetLastError());
OutputDebugString(rgwc);
OutputDebugString(L"GetDriveFreeSpace returning 0\n");
#endif
return 0;
}
dwFree = dwBps * dwCfc * dwSpc;
#ifdef LOUDLY
WCHAR rgc[100];
swprintf(rgc,L"GetDriveFreeSpace returns %d\n",dwFree);
OutputDebugString(rgc);
#endif
return dwFree;
}
DWORD GetDriveSectorSize(WCHAR *pszFilePath)
{
WCHAR *pwc;
WCHAR rgcModel[]={L"A:"};
DWORD dwSpc,dwBps,dwCfc,dwTcc;
if (NULL == pszFilePath) return 0;
rgcModel[0] = *pszFilePath;
if (!GetDiskFreeSpace(rgcModel,&dwSpc,&dwBps,&dwCfc,&dwTcc))
{
#ifdef LOUDLY
WCHAR rgwc[100];
swprintf(rgwc,L"GetDriveSectorSize encountered error %x\n",GetLastError());
OutputDebugString(rgwc);
OutputDebugString(L"GetDriveSectorSize returning 0\n");
#endif
return 0;
}
#ifdef LOUDLY
WCHAR rgc[100];
swprintf(rgc,L"GetDriveSectorSize returns %d\n",dwBps);
OutputDebugString(rgc);
#endif
return dwBps;
}
// take data size, sector size, return ptr to finished buffer
LPVOID CreateFileBuffer(INT iDataSize,INT iSectorSize)
{
INT iSize;
LPVOID lpv;
if (iDataSize == iSectorSize) iSize = iDataSize;
else
{
iSize = iDataSize/iSectorSize;
iSize += 1;
iSize *= iSectorSize;
}
g_iBufferSize = iSize;
lpv = VirtualAlloc(NULL,iSize,MEM_COMMIT,PAGE_READWRITE | PAGE_NOCACHE);
#ifdef LOUDLY
WCHAR rgc[100];
if (lpv) OutputDebugString(L"CreateFileBuffer succeeded\n");
else OutputDebugString(L"CreateFileBuffer failed ******\n");
swprintf(rgc,L"File Buffer size is %d\n",g_iBufferSize);
OutputDebugString(rgc);
#endif
return lpv;
}
// take ptr to buffer, release using VirtualFree()
void ReleaseFileBuffer(LPVOID lpv)
{
ZeroMemory(lpv,g_iBufferSize);
VirtualFree(lpv,0,MEM_RELEASE);
#ifdef LOUDLY
OutputDebugString(L"ReleaseFileBuffer called\n");
#endif
return;
}
/*
MediumIsPresent() returns true if there is a readable medium present in the drive.
*/
BOOL FileMediumIsPresent(TCHAR *pszPath) {
UINT uMode = 0;
BOOL bResult = FALSE;
TCHAR rgcModel[]=TEXT("A:");
DWORD dwError = 0;
if (*pszPath == 0) return FALSE;
rgcModel[0] = *pszPath;
uMode = SetErrorMode(SEM_FAILCRITICALERRORS);
if (0 == _waccess(rgcModel,0)) {
bResult = TRUE;
}
else dwError = GetLastError();
#ifdef LOUDLY
WCHAR rgwc[100];
swprintf(rgwc,L"_waccess returns error %x for %s\n",dwError,rgcModel);
OutputDebugString(rgwc);
if (!bResult) OutputDebugString(L"FileMediumIsPresent returning FALSE\n");
else OutputDebugString(L"FileMediumIsPresent returning TRUE\n");
#endif
// Correct certain obvious errors with the user's help
if (ERROR_UNRECOGNIZED_MEDIA == dwError)
{
// unformatted disk
WCHAR rgcFmt[200] = {0};
WCHAR rgcMsg[200] = {0};
WCHAR rgcTitle[200] = {0};
#ifdef LOUDLY
OutputDebugString(L"FileMediumIsPresent found an unformatted medium\n");
#endif
INT iCount = LoadString(g_hInstance,IDS_MBTFORMAT,rgcTitle,200 - 1);
iCount = LoadString(g_hInstance,IDS_MBMFORMAT,rgcFmt,200 - 1);
if (0 == iCount) goto LblNoBox;
swprintf(rgcMsg,rgcFmt,rgcModel);
INT iDrive = PathGetDriveNumber(rgcModel);
int iRet = MessageBox(c_hDlg,rgcMsg,rgcTitle,MB_YESNO);
if (IDYES == iRet)
{
dwError = SHFormatDrive(c_hDlg,iDrive,SHFMT_ID_DEFAULT,0);
if (0 == bResult) bResult = TRUE;
}
}
LblNoBox:
uMode = SetErrorMode(uMode);
return bResult;
}
//
// On save, create file if
// absent. Return handle on success, NULL on fail. FileName will be in
// c_rgcFileName.
//
HANDLE GetInputFile(void) {
HANDLE hFile = INVALID_HANDLE_VALUE;
DWORD dwErr;
WIN32_FILE_ATTRIBUTE_DATA stAttributes = {0};
#ifdef LOUDLY
OutputDebugString(L"GetInputFile() opening input file ");
OutputDebugString(pszFileName);
OutputDebugString(L"\n");
#endif
if (FileMediumIsPresent(pszFileName)) {
g_iSectorSize = GetDriveSectorSize(pszFileName);
if (0 == g_iSectorSize) return NULL;
if (GetFileAttributesEx(pszFileName,GetFileExInfoStandard,&stAttributes))
{
// file exists and we have a size for it.
g_iFileSize = stAttributes.nFileSizeLow;
}
else
{
dwErr = GetLastError();
if (dwErr == ERROR_FILE_NOT_FOUND)
RMessageBox(c_hDlg,IDS_MBTWRONGDISK ,IDS_MBMWRONGDISK ,MB_ICONEXCLAMATION);
else
{
#ifdef LOUDLY
{
WCHAR rgs[200] = {0};
swprintf(rgs,L"GetFileAttributesEx() failed, error = %x\n",dwErr);
OutputDebugString(rgs);
}
#endif
RMessageBox(c_hDlg,IDS_MBTDISKERROR ,IDS_MBMDISKERROR ,MB_ICONEXCLAMATION);
}
g_hFile = NULL;
return NULL;
} // end GetFileAttributes
hFile = CreateFileW(pszFileName,
GENERIC_READ,
0,
NULL,
OPEN_EXISTING,
FILE_FLAG_NO_BUFFERING,
NULL);
if (INVALID_HANDLE_VALUE == hFile) {
dwErr = GetLastError();
if (dwErr == ERROR_FILE_NOT_FOUND)
RMessageBox(c_hDlg,IDS_MBTWRONGDISK ,IDS_MBMWRONGDISK ,MB_ICONEXCLAMATION);
else
RMessageBox(c_hDlg,IDS_MBTDISKERROR ,IDS_MBMDISKERROR ,MB_ICONEXCLAMATION);
}
}
else {
RMessageBox(c_hDlg,IDS_MBTNODISK ,IDS_MBMNODISK ,MB_ICONEXCLAMATION);
}
if ((NULL == hFile) || (INVALID_HANDLE_VALUE == hFile)) {
g_hFile = NULL;
return NULL;
}
g_hFile = hFile;
return hFile;
}
void CloseInputFile(void)
{
#ifdef LOUDLY
OutputDebugString(L"Input file closed\n");
#endif
if (g_hFile)
{
CloseHandle(g_hFile);
g_hFile = NULL;
}
return;
}
HANDLE GetOutputFile(void) {
//HANDLE hFile = INVALID_HANDLE_VALUE;
HANDLE hFile = NULL;
DWORD dwErr;
if (FileMediumIsPresent(pszFileName)) {
g_iSectorSize = GetDriveSectorSize(pszFileName);
if (0 == g_iSectorSize) return NULL;
hFile = CreateFileW(pszFileName,
GENERIC_WRITE,
0,
NULL,
CREATE_NEW,
FILE_FLAG_NO_BUFFERING,
NULL);
if ((NULL == hFile) || (INVALID_HANDLE_VALUE == hFile)) {
dwErr = GetLastError();
#ifdef LOUDLY
TCHAR rgct[500];
swprintf(rgct,L"File create returns %x\n",dwErr);
OutputDebugString(rgct);
#endif
if ((dwErr == ERROR_FILE_EXISTS)) {
if (IDYES != RMessageBox(c_hDlg,IDS_MBTOVERWRITE ,IDS_MBMOVERWRITE ,MB_YESNO)) {
// Overwrite abandoned.
g_hFile = NULL;
return NULL;
}
else {
SetFileAttributes(pszFileName,FILE_ATTRIBUTE_NORMAL);
hFile = CreateFileW(pszFileName,
GENERIC_WRITE,
0,
NULL,
CREATE_ALWAYS,
FILE_FLAG_NO_BUFFERING,
NULL);
#ifdef LOUDLY
dwErr = GetLastError();
swprintf(rgct,L"File create failed %x\n",dwErr);
OutputDebugString(rgct);
#endif
}
} // end if already exists error
} // end if NULL == hFile
}
else {
RMessageBox(c_hDlg,IDS_MBTNODISK ,IDS_MBMNODISK ,MB_ICONEXCLAMATION);
}
if (INVALID_HANDLE_VALUE == hFile) {
g_hFile = NULL;
return NULL;
}
#ifdef LOUDLY
OutputDebugString(L"File successfully created\n");
#endif
g_hFile = hFile;
return hFile;
}
/*
DWORD ReadPrivateData(PWSTR,LPBYTE *,INT *)
DWORD WritePrivateData(PWSTR,LPBYTE,INT)
These functions read or write a reasonably short block of data to a disk
device, avoiding buffering of the data. This allows the data to be wiped
by the client before the buffers are released.
The DWORD return value is that which would return from GetLastError() and
can be handled accordingly.
ReadPrivateData() returns a malloc'd pointer which must be freed by the client. It
also returns the count of bytes read from the medium to the INT *. The file is
closed following the read operation before the function returns.
WritePrivateData() writes a count of bytes from LPBYTE to the disk. When it returns,
the buffer used to do so has been flushed and the file is closed.
*/
/*
prgb = byte ptr to data returned from the read
piCount = size of active data field within the buffer
Note that even if the read fails (file not found, read error, etc.) the buffer
ptr is still valid.
*/
//ReadFile(c_hFile,c_pPrivate,c_cbPrivate,&c_cbPrivate,NULL) return bytes read
INT ReadPrivateData(BYTE **prgb,INT *piCount)
{
LPVOID lpv;
DWORD dwBytesRead;
if (NULL == prgb) return 0;
if (NULL == piCount) return 0;
if (g_hFile)
{
lpv = CreateFileBuffer(g_iFileSize,g_iSectorSize);
if (NULL == lpv)
{
*prgb = 0; // indicate no need to free this buffer
*piCount = 0;
return 0;
}
*prgb = (BYTE *)lpv; // even if no data, gotta free using VirtualFree()
*piCount = 0;
if (0 == ReadFile(g_hFile,lpv,g_iBufferSize,&dwBytesRead,NULL)) return 0;
*piCount = g_iFileSize;
#ifdef LOUDLY
OutputDebugString(L"ReadPrivateData success\n");
#endif
return g_iFileSize;
}
return 0;
}
/*
Write data to file at g_rgwczFileName
Convert size to multiple of sector size
Alloc buffer
Copy data to buffer
Write data
Scrub Buffer, release
return 0
*/
BOOL WritePrivateData(BYTE *lpData,INT icbData) {
DWORD dwcb = 0;
LPVOID lpv;
if (NULL == g_hFile) return FALSE;
if (NULL == lpData) return FALSE;
if (0 == icbData) return FALSE;
if (g_hFile)
{
g_iFileSize = icbData;
lpv = CreateFileBuffer(g_iFileSize,g_iSectorSize);
if (NULL == lpv)
{
return FALSE;
}
ZeroMemory(lpv,g_iBufferSize);
memcpy(lpv,lpData,icbData);
WriteFile(g_hFile,lpv,g_iBufferSize,&dwcb,NULL);
VirtualFree(lpv,g_iBufferSize,MEM_RELEASE);
}
// ret TRUE iff file write succeeds and count of bytes is correct
#ifdef LOUDLY
if (dwcb) OutputDebugString(L"WritePrivateData succeeded\n");
else OutputDebugString(L"WritePrivateData failed ***\n");
#endif
if (dwcb != g_iBufferSize) return FALSE;
else return TRUE;
}