windows-nt/Source/XPSP1/NT/admin/wmi/wbem/winmgmt/roswell/a51tools.cpp
2020-09-26 16:20:57 +08:00

646 lines
14 KiB
C++

/*++
Copyright (C) 2000-2001 Microsoft Corporation
--*/
#include <windows.h>
#include <stdio.h>
#include <wbemcomn.h>
#include "a51tools.h"
__int64 g_nCurrentTime = 1;
__int64 g_nReadFailures = 0;
__int64 g_nWriteFailures = 0;
//
// FILE_ATTRIBUTE_NOT_CONTENT_INDEXED is not actually supported on CreateFile,
// contrary to the docs. However, also contrary to the docs, it is inherited
// from the parent directory
//
#define A51_FILE_CREATION_FLAGS 0 //FILE_ATTRIBUTE_NOT_CONTENT_INDEXED
CTempMemoryManager g_Manager;
void* TempAlloc(DWORD dwLen)
{
return g_Manager.Allocate(dwLen);
}
void TempFree(void* p, DWORD dwLen)
{
g_Manager.Free(p, dwLen);
}
void* TempAlloc(CTempMemoryManager& Manager, DWORD dwLen)
{
return Manager.Allocate(dwLen);
}
void TempFree(CTempMemoryManager& Manager, void* p, DWORD dwLen)
{
Manager.Free(p, dwLen);
}
HANDLE g_hLastEvent = NULL;
CCritSec g_csEvents;
HANDLE A51GetNewEvent()
{
{
CInCritSec ics(&g_csEvents);
if(g_hLastEvent)
{
HANDLE h = g_hLastEvent;
g_hLastEvent = NULL;
return h;
}
}
return CreateEvent(NULL, TRUE, FALSE, NULL);
}
void A51ReturnEvent(HANDLE hEvent)
{
{
CInCritSec ics(&g_csEvents);
if(g_hLastEvent == NULL)
{
g_hLastEvent = hEvent;
return;
}
}
CloseHandle(hEvent);
}
HRESULT A51TranslateErrorCode(long lRes)
{
if (SUCCEEDED(lRes))
return WBEM_S_NO_ERROR;
switch(lRes)
{
case ERROR_FILE_NOT_FOUND:
case ERROR_PATH_NOT_FOUND:
return WBEM_E_NOT_FOUND;
case ERROR_OUTOFMEMORY:
return WBEM_E_OUT_OF_MEMORY;
case ERROR_NOT_ENOUGH_QUOTA:
case ERROR_DISK_FULL:
return WBEM_E_OUT_OF_DISK_SPACE;
default:
return WBEM_E_FAILED;
}
}
long __stdcall EnsureDirectory(LPCWSTR wszPath, LPSECURITY_ATTRIBUTES pSA)
{
if(!CreateDirectoryW(wszPath, NULL))
{
long lRes = GetLastError();
if(lRes != ERROR_ALREADY_EXISTS)
return lRes;
else
return ERROR_SUCCESS;
}
else
return ERROR_SUCCESS;
}
long __stdcall EnsureDirectoryRecursiveForFile(LPWSTR wszPath,
LPSECURITY_ATTRIBUTES pSA);
long __stdcall EnsureDirectoryForFile(LPCWSTR wszPath, LPSECURITY_ATTRIBUTES pSA)
{
//
// Make a copy, since we will be messing with it
//
CFileName wszNewPath;
if (wszNewPath == NULL)
return ERROR_OUTOFMEMORY;
wcscpy(wszNewPath, wszPath);
return EnsureDirectoryRecursiveForFile(wszNewPath, pSA);
}
long __stdcall EnsureDirectoryRecursiveForFile(LPWSTR wszPath,
LPSECURITY_ATTRIBUTES pSA)
{
long lRes;
//
// Find the last backslash and remove
//
WCHAR* pwcLastSlash = wcsrchr(wszPath, L'\\');
if(pwcLastSlash == NULL)
return ERROR_BAD_PATHNAME;
*pwcLastSlash = 0;
//
// Try to create it
//
if(!CreateDirectoryW(wszPath, pSA))
{
//
// Call ourselves recursively --- to create our parents
//
lRes = EnsureDirectoryRecursiveForFile(wszPath, pSA);
if(lRes != ERROR_SUCCESS)
{
*pwcLastSlash = L'\\';
return lRes;
}
//
// Try again
//
BOOL bRes = CreateDirectoryW(wszPath, pSA);
*pwcLastSlash = L'\\';
if(bRes)
return ERROR_SUCCESS;
else
return GetLastError();
}
else
{
*pwcLastSlash = L'\\';
return ERROR_SUCCESS;
}
}
inline WCHAR HexDigit(BYTE b)
{
if(b < 10)
return L'0' + b;
else
return L'A' + b - 10;
}
bool A51Hash(LPCWSTR wszName, LPWSTR wszHash)
{
//
// Have to upper-case everything
//
DWORD dwBufferSize = wcslen(wszName)*2+2;
LPWSTR wszBuffer = (WCHAR*)TempAlloc(dwBufferSize);
if (wszBuffer == NULL)
return false;
CTempFreeMe vdm(wszBuffer, dwBufferSize);
wbem_wcsupr(wszBuffer, wszName);
BYTE RawHash[16];
MD5::Transform((void*)wszBuffer, wcslen(wszBuffer)*2, RawHash);
WCHAR* pwc = wszHash;
for(int i = 0; i < 16; i++)
{
*(pwc++) = HexDigit(RawHash[i]/16);
*(pwc++) = HexDigit(RawHash[i]%16);
}
*pwc = 0;
return true;
}
long A51WriteFile(LPCWSTR wszFullPath, DWORD dwLen, BYTE* pBuffer)
{
long lRes;
A51TRACE(( "Create file %S\n", wszFullPath));
//
// Create the right file
//
HANDLE hFile = CreateFileW(wszFullPath, GENERIC_WRITE, 0,
NULL, CREATE_ALWAYS, A51_FILE_CREATION_FLAGS, NULL);
if(hFile == INVALID_HANDLE_VALUE)
{
lRes = GetLastError();
if(lRes == ERROR_PATH_NOT_FOUND)
{
lRes = EnsureDirectoryForFile(wszFullPath, NULL);
if(lRes != ERROR_SUCCESS)
return lRes;
hFile = CreateFileW(wszFullPath, GENERIC_WRITE, 0,
NULL, CREATE_ALWAYS, A51_FILE_CREATION_FLAGS, NULL);
if(hFile == INVALID_HANDLE_VALUE)
{
lRes = GetLastError();
_ASSERT(lRes != ERROR_SUCCESS, L"success error code from fail");
return lRes;
}
}
else
return lRes;
}
CCloseMe cm(hFile);
//
// Write it and close
//
if(!WriteFile(hFile, pBuffer, dwLen, &dwLen, NULL))
{
lRes = GetLastError();
_ASSERT(lRes != ERROR_SUCCESS, L"success error code from fail");
return lRes;
}
return ERROR_SUCCESS;
}
long A51DeleteFile(LPCWSTR wszFullPath)
{
A51TRACE(("Delete file %S\n", wszFullPath));
//
// Delete the right file
//
if(!DeleteFileW(wszFullPath))
return GetLastError();
return ERROR_SUCCESS;
}
long A51WriteToFileAsync(HANDLE hFile, long lStartingOffset, BYTE* pBuffer,
DWORD dwBufferLen, OVERLAPPED* pov)
{
pov->Offset = lStartingOffset;
pov->OffsetHigh = 0;
pov->hEvent = A51GetNewEvent();
BOOL bRes;
DWORD dwWritten;
while(!(bRes = WriteFile(hFile, pBuffer, dwBufferLen, &dwWritten, pov)) &&
(GetLastError() == ERROR_INVALID_USER_BUFFER ||
GetLastError() == ERROR_NOT_ENOUGH_MEMORY))
{
//
// Out of buffer-space --- wait a bit and retry
//
g_nWriteFailures++;
Sleep(16);
}
if(!bRes)
{
long lRes = GetLastError();
if(lRes == ERROR_IO_PENDING)
// perfect!
return ERROR_SUCCESS;
else
{
A51ReturnEvent(pov->hEvent);
pov->hEvent = NULL;
return lRes;
}
}
else
{
//
// Succeeded synchronously --- clean up and return
//
A51ReturnEvent(pov->hEvent);
pov->hEvent = NULL;
if(dwWritten != dwBufferLen)
return ERROR_OUTOFMEMORY;
else
return ERROR_SUCCESS;
}
}
long A51WriteToFileSync(HANDLE hFile, long lStartingOffset, BYTE* pBuffer,
DWORD dwBufferLen)
{
OVERLAPPED ov;
long lRes = A51WriteToFileAsync(hFile, lStartingOffset, pBuffer,
dwBufferLen, &ov);
if(lRes != ERROR_SUCCESS)
return lRes;
if(ov.hEvent)
{
CReturnMe rm(ov.hEvent);
//
// Wait for completion
//
DWORD dwWritten;
if(!GetOverlappedResult(hFile, &ov, &dwWritten, TRUE))
return GetLastError();
if(dwWritten != dwBufferLen)
return ERROR_OUTOFMEMORY;
}
return ERROR_SUCCESS;
}
long A51ReadFromFileAsync(HANDLE hFile, long lStartingOffset, BYTE* pBuffer,
DWORD dwBufferLen, OVERLAPPED* pov)
{
long lRes;
pov->Offset = lStartingOffset;
pov->OffsetHigh = 0;
pov->hEvent = A51GetNewEvent();
BOOL bRes;
DWORD dwRead;
while(!(bRes = ReadFile(hFile, pBuffer, dwBufferLen, &dwRead, pov)) &&
(GetLastError() == ERROR_INVALID_USER_BUFFER ||
GetLastError() == ERROR_NOT_ENOUGH_MEMORY))
{
//
// Out of buffer-space --- wait a bit and retry
//
g_nReadFailures++;
Sleep(16);
}
if(!bRes)
{
if(GetLastError() == ERROR_IO_PENDING)
// perfect!
return ERROR_SUCCESS;
else
{
lRes = GetLastError();
_ASSERT(lRes != ERROR_SUCCESS, L"Success returned on failure");
A51ReturnEvent(pov->hEvent);
pov->hEvent = NULL;
if(lRes == ERROR_HANDLE_EOF)
lRes = ERROR_SUCCESS;
return lRes;
}
}
else
{
//
// Succeeded synchronously --- clean up and return
//
A51ReturnEvent(pov->hEvent);
pov->hEvent = NULL;
return ERROR_SUCCESS;
}
}
long A51ReadFromFileSync(HANDLE hFile, long lStartingOffset, BYTE* pBuffer,
DWORD dwBufferLen)
{
OVERLAPPED ov;
long lRes = A51ReadFromFileAsync(hFile, lStartingOffset, pBuffer,
dwBufferLen, &ov);
if(lRes != ERROR_SUCCESS)
return lRes;
if(ov.hEvent)
{
CReturnMe rm(ov.hEvent);
//
// Wait for completion
//
DWORD dwRead = 0;
if(!GetOverlappedResult(hFile, &ov, &dwRead, TRUE))
{
lRes = GetLastError();
_ASSERT(lRes != ERROR_SUCCESS, L"Success returned on failure");
if(lRes == ERROR_HANDLE_EOF)
lRes = ERROR_SUCCESS;
return lRes;
}
}
return ERROR_SUCCESS;
}
long RemoveDirectoryRecursive(LPCWSTR wszDirectoryPath, bool bAbortOnFiles);
long A51RemoveDirectory(LPCWSTR wszFullPath, bool bAbortOnFiles)
{
long lRes = RemoveDirectoryRecursive(wszFullPath, bAbortOnFiles);
if(lRes == ERROR_PATH_NOT_FOUND || lRes == ERROR_FILE_NOT_FOUND)
return ERROR_FILE_NOT_FOUND;
else
return lRes;
}
long RemoveDirectoryRecursive(LPCWSTR wszDirectoryPath, bool bAbortOnFiles)
{
long lRes;
//
// Try removing it right away
//
if(!RemoveDirectoryW(wszDirectoryPath))
{
long lRes = GetLastError();
if(lRes == ERROR_PATH_NOT_FOUND || lRes == ERROR_FILE_NOT_FOUND)
return ERROR_FILE_NOT_FOUND;
else if(lRes != ERROR_DIR_NOT_EMPTY && lRes != ERROR_SHARING_VIOLATION)
return lRes;
}
else
return ERROR_SUCCESS;
//
// Not empty (or at least we are not sure it is empty) --- enumerate
// everything
//
CFileName wszMap;
if (wszMap == NULL)
return ERROR_OUTOFMEMORY;
wcscpy(wszMap, wszDirectoryPath);
wcscat(wszMap, L"\\*");
CFileName wszChild;
if (wszChild == NULL)
return ERROR_OUTOFMEMORY;
wcscpy(wszChild, wszDirectoryPath);
wcscat(wszChild, L"\\");
long lChildLen = wcslen(wszChild);
WIN32_FIND_DATAW fd;
HANDLE hSearch = FindFirstFileW(wszMap, &fd);
if(hSearch == INVALID_HANDLE_VALUE)
return ERROR_DIR_NOT_EMPTY;
do
{
if(fd.cFileName[0] == L'.')
continue;
if((fd.dwFileAttributes & FILE_ATTRIBUTE_DIRECTORY) == 0)
{
if(bAbortOnFiles)
{
// There is a file --- give up
FindClose(hSearch);
return ERROR_DIR_NOT_EMPTY;
}
else
{
wcscpy(wszChild + lChildLen, fd.cFileName);
if(!DeleteFileW(wszChild))
{
FindClose(hSearch);
return GetLastError();
}
}
}
else
{
wcscpy(wszChild + lChildLen, fd.cFileName);
lRes = RemoveDirectoryRecursive(wszChild, bAbortOnFiles);
if(lRes != ERROR_SUCCESS && lRes != ERROR_FILE_NOT_FOUND)
{
FindClose(hSearch);
return lRes;
}
}
}
while(FindNextFileW(hSearch, &fd));
FindClose(hSearch);
//
// And try again
//
if(!RemoveDirectoryW(wszDirectoryPath))
{
long lRes = GetLastError();
if(lRes == ERROR_PATH_NOT_FOUND || lRes == ERROR_FILE_NOT_FOUND)
return ERROR_FILE_NOT_FOUND;
else if(lRes == ERROR_SHARING_VIOLATION)
return ERROR_SUCCESS;
else
return lRes;
}
return ERROR_SUCCESS;
}
CRITICAL_SECTION g_csLog;
char* g_szText = NULL;
long g_lTextLen = 0;
WCHAR g_wszLogFilename[MAX_PATH] = L"";
void A51Trace(LPCSTR szFormat, ...)
{
if((g_wszLogFilename[0] == 0) || (g_szText == NULL))
{
InitializeCriticalSection(&g_csLog);
delete g_szText;
g_wszLogFilename[0] = 0;
g_szText = NULL;
HKEY hKey;
long lRes = RegOpenKeyExW(HKEY_LOCAL_MACHINE,
L"SOFTWARE\\Microsoft\\WBEM\\CIMOM",
0, KEY_READ, &hKey);
if(lRes)
{
wcscpy(g_wszLogFilename, L"c:\\a51.log");
}
else
{
DWORD dwLen = MAX_PATH;
CFileName wszTmp;
if (wszTmp == 0)
{
wcscpy(g_wszLogFilename, L"c:\\a51.log");
RegCloseKey(hKey);
}
else
{
lRes = RegQueryValueExW(hKey, L"Logging Directory", NULL, NULL,
(LPBYTE)(wchar_t*)wszTmp, &dwLen);
RegCloseKey(hKey);
if(lRes)
{
wcscpy(g_wszLogFilename, L"c:\\a51.log");
}
else
{
if (ExpandEnvironmentStringsW(wszTmp,g_wszLogFilename,MAX_PATH) == 0)
{
wcscpy(g_wszLogFilename, L"c:\\a51.log");
}
else
{
if (g_wszLogFilename[wcslen(g_wszLogFilename)] == L'\\')
{
wcscat(g_wszLogFilename, L"a51.log");
}
else
{
wcscat(g_wszLogFilename, L"\\a51.log");
}
}
}
}
}
g_szText = new char[3000000];
if (g_szText == NULL)
return;
}
EnterCriticalSection(&g_csLog);
char szBuffer[256];
va_list argptr;
va_start(argptr, szFormat);
vsprintf(szBuffer, szFormat, argptr);
long lLen = strlen(szBuffer);
if(g_lTextLen + lLen > 2900000)
A51TraceFlush();
strcpy(g_szText + g_lTextLen, szBuffer);
g_lTextLen += lLen;
LeaveCriticalSection(&g_csLog);
}
void A51TraceFlush()
{
FILE* fLog = NULL;
fLog = _wfopen(g_wszLogFilename, L"a");
if(fLog)
{
fwrite(g_szText, 1, g_lTextLen, fLog);
g_lTextLen = 0;
// fflush(fLog);
fclose(fLog);
}
}