#include <wbemcomn.h>
#include "a51tools.h"
#include "objheap.h"
#include "index.h"

#define ROSWELL_HEAPALLOC_TYPE_BUSY 0xA51A51A5

//*******************************************************************
//*******************************************************************
long CObjectHeap::Initialize(CPageSource  * pAbstractSource, 
                             WCHAR * wszBaseName,
                             DWORD dwBaseNameLen)
{
    if (m_bInit)
        return ERROR_SUCCESS;
        
    long lRes;
    
    lRes = m_Heap.Initialize(pAbstractSource);

    if(lRes != ERROR_SUCCESS)
        return lRes;

    lRes = m_Index.Initialize(dwBaseNameLen, wszBaseName, pAbstractSource);
    if(lRes != ERROR_SUCCESS)
        return lRes;

    m_bInit = TRUE;

    return lRes;
}

//*******************************************************************
//*******************************************************************
long CObjectHeap::Uninitialize(DWORD dwShutDownFlags)
{
    if (!m_bInit)
        return ERROR_SUCCESS;
        
    m_Index.Shutdown(dwShutDownFlags);

    m_Heap.Shutdown(dwShutDownFlags);

    m_bInit = FALSE;

    return ERROR_SUCCESS;
}

//*******************************************************************
//*******************************************************************
void CObjectHeap::InvalidateCache()
{
    m_Index.InvalidateCache();
    m_Heap.InvalidateCache();
}

//*******************************************************************
//*******************************************************************
long CObjectHeap::GetIndexFileName(LPCWSTR wszFilePath, LPWSTR wszIndexFileName)
{
    WIN32_FIND_DATAW wfd;

    long lRes = m_Index.FindFirst(wszFilePath, &wfd, NULL);
    if(lRes != ERROR_SUCCESS)
    {
        if(lRes == ERROR_PATH_NOT_FOUND)
            lRes = ERROR_FILE_NOT_FOUND;
        return lRes;
    }

    wcscpy(wszIndexFileName, wfd.cFileName);
    return ERROR_SUCCESS;
}

//*******************************************************************
//*******************************************************************
long CObjectHeap::GetFileInfo(LPCWSTR wszFilePath, TPage *pnPage, TOffset* pnOffset,
                            DWORD* pdwLength)
{
    CFileName wszIndexFileName;
    if(wszIndexFileName == NULL)
        return ERROR_OUTOFMEMORY;

    long lRes = GetIndexFileName(wszFilePath, wszIndexFileName);
    if(lRes != ERROR_SUCCESS)
        return lRes;

    return ParseInfoFromIndexFile(wszIndexFileName, pnPage, pnOffset, pdwLength);
}

//*******************************************************************
//*******************************************************************
long CObjectHeap::ParseInfoFromIndexFile(LPCWSTR wszIndexFileName, 
										TPage *pnPage,
                                        TOffset* pnOffset, 
										DWORD* pdwLength)
{
    WCHAR* pDot = wcschr(wszIndexFileName, L'.');
    if(pDot == NULL)
        return ERROR_INVALID_PARAMETER;

    WCHAR* pwc = pDot+1;
    *pnPage = 0;
    while(*pwc && *pwc != L'.')
    {
        *pnPage = (*pnPage * 10) + (*pwc - '0');
		pwc++;
    }

    if(*pwc != L'.')
        return ERROR_INVALID_PARAMETER;

    pwc++;

    *pnOffset = 0;
    while(*pwc && *pwc != L'.')
    {
        *pnOffset = (*pnOffset * 10) + (*pwc - '0');
		pwc++;
    }

    if(*pwc != L'.')
        return ERROR_INVALID_PARAMETER;

    pwc++;

    *pdwLength = 0;
    while(*pwc && *pwc != L'.')
    {
        *pdwLength = (*pdwLength * 10) + (*pwc - '0');
		pwc++;
    }

    return ERROR_SUCCESS;
}

//*******************************************************************
//*******************************************************************
long CObjectHeap::CreateIndexFile(LPCWSTR wszFilePath, 
								  TPage nPage, 
								  TOffset nOffset,
								  DWORD dwLength)
{
    //
    // Simply append the numbers to the file path
    //

    CFileName wszIndexFilePath;
    if(wszIndexFilePath == NULL)
        return ERROR_OUTOFMEMORY;

    swprintf(wszIndexFilePath, L"%s.%u.%u.%u",
        wszFilePath, nPage, nOffset, dwLength);

    return CreateZeroLengthFile(wszIndexFilePath);
}

//*******************************************************************
//*******************************************************************
long CObjectHeap::DeleteIndexFile(LPCWSTR wszFilePath, LPCWSTR wszIndexFileName)
{
    //
    // Construct the full path to the index file by concatenating the directory 
    // of the original file with the name
    //

    CFileName wszIndexFilePath;
    if(wszIndexFilePath == NULL)
        return ERROR_OUTOFMEMORY;

    WCHAR* pwcLastSlash = wcsrchr(wszFilePath, L'\\');
    if(pwcLastSlash == NULL)
        return ERROR_INVALID_PARAMETER;

    int nPrefixLen = (pwcLastSlash - wszFilePath + 1);
    memcpy(wszIndexFilePath, wszFilePath, nPrefixLen * sizeof(WCHAR));

    wcscpy(wszIndexFilePath + nPrefixLen, wszIndexFileName);
	return DeleteZeroLengthFile(wszIndexFilePath);
}

    

//*******************************************************************
//*******************************************************************
long CObjectHeap::CreateZeroLengthFile(LPCWSTR wszFilePath)
{
	return m_Index.Create(wszFilePath);
}

//*******************************************************************
//*******************************************************************
long CObjectHeap::DeleteZeroLengthFile(LPCWSTR wszFilePath)
{
    return m_Index.Delete(wszFilePath);
}
    
//*******************************************************************
//*******************************************************************
long CObjectHeap::WriteAllocation(DWORD dwDataLength, BYTE* pData, TPage *pnPage, TOffset *pnOffset)
{
    return m_Heap.WriteNewBuffer(dwDataLength, pData, pnPage, pnOffset);
}
//*******************************************************************
//*******************************************************************
long CObjectHeap::WriteExistingAllocation(TPage nOldPage, TOffset nOldOffset, DWORD dwDataLength, BYTE *pBuffer, DWORD *pnNewPage, DWORD *pnNewOffset)
{
    return m_Heap.WriteExistingBuffer(dwDataLength, pBuffer, nOldPage, nOldOffset, pnNewPage, pnNewOffset);
}

//*******************************************************************
//*******************************************************************
long CObjectHeap::ReadAllocation(TPage nPage, TOffset nOffset, DWORD dwDataLength, BYTE* pBuffer)
{
    //
    // Prepare a buffer with the complete allocation
    //

    BYTE* pAllocation;
	DWORD dwReadLength;

    long lRes = m_Heap.ReadBuffer(nPage, nOffset, &pAllocation, &dwReadLength);
    if(lRes != ERROR_SUCCESS)
        return lRes;

    memcpy(pBuffer, pAllocation, dwDataLength);

	delete [] pAllocation;

    return ERROR_SUCCESS;
}
    

//*******************************************************************
//*******************************************************************
long CObjectHeap::WriteObject(LPCWSTR wszFilePath1, LPCWSTR wszFilePath2, DWORD dwBufferLen, BYTE* pBuffer)
{
    if (!m_bInit)
        return -1;

    long lRes;

    if(dwBufferLen == 0)
    {
        //
        // We do not use the heap for 0-length files, we create them directly
        //

        return CreateZeroLengthFile(wszFilePath1);
    }

    //
    // Now, check if this file already exists
    //

    CFileName wszIndexFileName1;
    if(wszIndexFileName1 == NULL)
        return ERROR_OUTOFMEMORY;

    lRes = GetIndexFileName(wszFilePath1, wszIndexFileName1);
    if(lRes != ERROR_FILE_NOT_FOUND && lRes != ERROR_SUCCESS)
        return lRes;

    if(lRes == ERROR_SUCCESS)
    {
        //
        // Already there. 
        //

		TPage nOldPage;
        TOffset nOldOffset;
        DWORD dwOldLength;
		TPage nNewPage;
		TOffset nNewOffset;

        lRes = ParseInfoFromIndexFile(wszIndexFileName1, &nOldPage, &nOldOffset, &dwOldLength);
        if(lRes != ERROR_SUCCESS)
            return lRes;


        //
        // Enough space in place --- just write the data and update the
        // length
        //

        lRes = WriteExistingAllocation(nOldPage, nOldOffset, dwBufferLen, pBuffer, &nNewPage, &nNewOffset);
        if(lRes != ERROR_SUCCESS)
            return lRes;

        if((dwOldLength != dwBufferLen) || (nOldPage != nNewPage) || (nOldOffset != nNewOffset))
        {

            //
            // One of the bits of the path has changed so we need to re-create the index
            //

            lRes = DeleteIndexFile(wszFilePath1, wszIndexFileName1);
            if(lRes != ERROR_SUCCESS)
                return lRes;

            lRes = CreateIndexFile(wszFilePath1, nNewPage, nNewOffset, dwBufferLen);
            if(lRes != ERROR_SUCCESS)
                return lRes;

			if (wszFilePath2)
			{
				CFileName wszIndexFileName2;
				if(wszIndexFileName2 == NULL)
					return ERROR_OUTOFMEMORY;
				lRes = GetIndexFileName(wszFilePath2, wszIndexFileName2);
				if(lRes != ERROR_SUCCESS)
					return lRes;

				lRes = DeleteIndexFile(wszFilePath2, wszIndexFileName2);
				if(lRes != ERROR_SUCCESS)
					return lRes;

				lRes = CreateIndexFile(wszFilePath2, nNewPage, nNewOffset, dwBufferLen);
				if(lRes != ERROR_SUCCESS)
					return lRes;
			}
        }

        return ERROR_SUCCESS;
    }

    //
    // it wasn't there
    //

	TPage nPage;
    TOffset nOffset;

    lRes = WriteAllocation(dwBufferLen, pBuffer, &nPage, &nOffset);
    if(lRes != ERROR_SUCCESS)
        return lRes;

    lRes = CreateIndexFile(wszFilePath1, nPage, nOffset, dwBufferLen);
    if(lRes != ERROR_SUCCESS)
        return lRes;

	if (wszFilePath2)
		lRes = CreateIndexFile(wszFilePath2, nPage, nOffset, dwBufferLen);
    
    return lRes;
}

//***********************************************************************
//***********************************************************************
long CObjectHeap::WriteLink(LPCWSTR wszLinkPath)
{
    if (!m_bInit)
        return -1;

    return CreateZeroLengthFile(wszLinkPath);
}
//*******************************************************************
//*******************************************************************
long CObjectHeap::ReadObject(LPCWSTR wszFilePath, DWORD* pdwLength, BYTE** ppBuffer)
{
    if (!m_bInit)
        return -1;    

    long lRes;

    //
    // Find the file
    //

	TPage nPage;
    TOffset nOffset;
    lRes = GetFileInfo(wszFilePath, &nPage, &nOffset, pdwLength);
    if(lRes != ERROR_SUCCESS)
        return lRes;

    //
    // Read the allocation
    //

    *ppBuffer = (BYTE*)TempAlloc(*pdwLength);
    if(*ppBuffer == NULL)
        return ERROR_OUTOFMEMORY;

    lRes = ReadAllocation(nPage, nOffset, *pdwLength, *ppBuffer);
    if(lRes != ERROR_SUCCESS)
    {
        TempFree(*ppBuffer);
        return lRes;
    }

    return ERROR_SUCCESS;
}

//*******************************************************************
//*******************************************************************
long CObjectHeap::DeleteLink(LPCWSTR wszFilePath)
{
    if (!m_bInit)
        return -1;    

    //
    // Find the index file
    //

    CFileName wszIndexFileName;
    if(wszIndexFileName == NULL)
        return ERROR_OUTOFMEMORY;

    long lRes = GetIndexFileName(wszFilePath, wszIndexFileName);
    if(lRes != ERROR_SUCCESS)
	{
        return lRes;
	}

	//If we have any index information at the end of the path we need to add that to
	//the main path and delete that.  The delete requires an accurate path.
	CFileName wszActualFileName;
    if(wszActualFileName == NULL)
        return ERROR_OUTOFMEMORY;

	wcscpy(wszActualFileName, wszFilePath);
	wchar_t *wszDot = wcschr(wszIndexFileName, L'.');
	if (wszDot != NULL)
		wcscat(wszActualFileName, wszDot);

	return DeleteZeroLengthFile(wszActualFileName);
}
//*******************************************************************
//*******************************************************************
long CObjectHeap::DeleteObject(LPCWSTR wszFilePath)
{
    if (!m_bInit)
        return -1;    

    //
    // Find the index file
    //

    CFileName wszIndexFileName;
    if(wszIndexFileName == NULL)
        return ERROR_OUTOFMEMORY;

    long lRes = GetIndexFileName(wszFilePath, wszIndexFileName);
    if(lRes != ERROR_SUCCESS)
        return lRes;

    //
    // Delete the allocation
    //

	TPage nPage;
    TOffset nOffset;
    DWORD dwLength;
    lRes = ParseInfoFromIndexFile(wszIndexFileName, &nPage, &nOffset, &dwLength);
	if(lRes == ERROR_INVALID_PARAMETER)
	{
		OutputDebugString(L"WinMgmt: Deleting an object that does not have details of where object is!\n");
		DebugBreak();
		return ERROR_INVALID_OPERATION;
    }
    
    lRes = m_Heap.DeleteBuffer(nPage, nOffset);
    if(lRes != ERROR_SUCCESS)
        return lRes;
    
    //
    // Delete the index itself
    //

    lRes = DeleteIndexFile(wszFilePath, wszIndexFileName);
    if(lRes != ERROR_SUCCESS)
        return lRes;

    return ERROR_SUCCESS;
}

long CObjectHeap::ObjectEnumerationBegin(const wchar_t *wszSearchPrefix, void **ppHandle)
{
	return IndexEnumerationBegin(wszSearchPrefix, ppHandle);
}

long CObjectHeap::ObjectEnumerationEnd(void *pHandle)
{
	return IndexEnumerationEnd(pHandle);
}

long CObjectHeap::ObjectEnumerationNext(void *pHandle, CFileName &wszFileName, BYTE **ppBlob, DWORD *pdwSize)
{
    if (!m_bInit)
        return -1;  
	
	long lRes = m_Index.IndexEnumerationNext(pHandle, wszFileName);

	if (lRes == ERROR_SUCCESS)
	{
		//We need to retrieve the object from the heap!
		TPage nPage;
		TOffset nOffset;
		DWORD dwLength;
		lRes = ParseInfoFromIndexFile(wszFileName, &nPage, &nOffset, &dwLength);
		if(lRes == ERROR_INVALID_PARAMETER)
			lRes = ERROR_SUCCESS;	//This is a plain enumeration, no blobs associated with it
		else
		{
			//Remove extra stuff from end of string...
			for (int nCount = 0, nIndex = wcslen(wszFileName); nCount != 3; nIndex --)
			{
				if (wszFileName[nIndex-1] == L'.')
				{
					if (++nCount == 3)
						wszFileName[nIndex-1] = L'\0';
				}

			}
			DWORD dwSize = 0;
			lRes = m_Heap.ReadBuffer(nPage, nOffset, ppBlob, &dwSize);
			*pdwSize = dwLength;
		}
	}
	return lRes;
}

long CObjectHeap::ObjectEnumerationFree(void *pHandle, BYTE *pBlob)
{
    if (!m_bInit)
        return -1; 

	delete [] pBlob;
	
	return ERROR_SUCCESS;
}

long CObjectHeap::IndexEnumerationBegin(const wchar_t *wszSearchPrefix, void **ppHandle)
{
    if (!m_bInit)
        return -1;
	return m_Index.IndexEnumerationBegin(wszSearchPrefix, ppHandle);
}
long CObjectHeap::IndexEnumerationEnd(void *pHandle)
{
    if (!m_bInit)
        return -1;    
	return m_Index.IndexEnumerationEnd(pHandle);
}
long CObjectHeap::IndexEnumerationNext(void *pHandle, CFileName &wszFileName)
{
    if (!m_bInit)
        return -1;  
	
	long lRes = m_Index.IndexEnumerationNext(pHandle, wszFileName);

	if (lRes == ERROR_SUCCESS)
	{
		wchar_t *pDot = wcschr(wszFileName, L'.');
		if (pDot)
			*pDot = L'\0';
	}

	return lRes;
}

long CObjectHeap::FlushCaches()
{
	long lRes = m_Index.FlushCaches();
	if (lRes == 0)
		lRes = m_Heap.FlushCaches();
	return lRes;
}