//
// Copyright (c) Microsoft Corporation.  All rights reserved.
//
//
// Use of this source code is subject to the terms of the Microsoft shared
// source or premium shared source license agreement under which you licensed
// this source code. If you did not accept the terms of the license agreement,
// you are not authorized to use this source code. For the terms of the license,
// please see the license agreement between you and Microsoft or, if applicable,
// see the SOURCE.RTF on your install media or the root of your tools installation.
// THE SOURCE CODE IS PROVIDED "AS IS", WITH NO WARRANTIES OR INDEMNITIES.
//
#ifndef __PARTITION_H__
#define __PARTITION_H__

#include <partdrv.h>
#include "detector.hpp"
#include <dlist.h>

// The lower 16bits of the MountFlags are used by filesys
#define MOUNTFLAGS_TYPE_HIDDEN      AFS_FLAG_HIDDEN     // 1  - Hidden file system
#define MOUNTFLAGS_TYPE_BOOTABLE    AFS_FLAG_BOOTABLE   // 2  - May contain system registry
#define MOUNTFLAGS_TYPE_ROOTFS      AFS_FLAG_ROOTFS     // 4  - Mount as root of file system, "\"
#define MOUNTFLAGS_TYPE_HIDEROM     AFS_FLAG_HIDEROM    // 8  - Hide ROM when mounting FS root; use with FLAG_ROOTFS
#define MOUNTFLAGS_TYPE_MOUNTROM    AFS_FLAG_MOUNTROM   // 16 - Mount the new filesystem as an additional ROM filesystem

#define PARTITION_STATE_DISABLED        0x00000002  //  Partition is Disabled
#define PARTITION_STATE_MEDIA_DETACHED  0x00000004  //  Media is detached from this partition
#define PARTITION_STATE_MOUNTED         0x00000008  //  Partition is Mounted

#ifndef _PREFAST_
#pragma warning (disable: 4068) // Disable pragma warnings
#endif

typedef DWORD (* POPENSTORE) (HANDLE hDisk, PDWORD pdwStoreId);
typedef void  (* PCLOSESTORE) (DWORD dwStoreId);
typedef DWORD (* PFORMATSTORE) (DWORD dwStoreId);
typedef DWORD (* PISSTOREFORMATTED) (DWORD dwStoreId);
typedef DWORD (* PGETSTOREINFO) (DWORD dwStoreId, PD_STOREINFO *info);
typedef DWORD (* PCREATEPARTITION) (DWORD dwStoreId, LPCTSTR partName, BYTE bPartType, SECTORNUM numSectors, BOOL bAuto);
typedef DWORD (* PDELETEPARTITION) (DWORD dwStoreId, LPCTSTR partName);
typedef DWORD (* PRENAMEPARTITION) (DWORD dwStoreId, LPCTSTR oldName, LPCTSTR newName);
typedef DWORD (* PSETPARTITIONATTRS) (DWORD dwStoreId, LPCTSTR partName, DWORD attrs);
typedef DWORD (* PSETPARTITIONSIZE) (DWORD dwStoreId, LPCTSTR partName, SECTORNUM numSectors);
typedef DWORD (* PFORMATPARTITION) (DWORD dwStoreId, LPCTSTR partName, BYTE bPartType, BOOL bAuto);
typedef DWORD (* PGETPARTITIONINFO) (DWORD dwStoreId, LPCTSTR partName, PPD_PARTINFO info);
typedef DWORD (* PFINDPARTITIONSTART) (DWORD dwStoreId, PDWORD pdwSearchId);
typedef DWORD (* PFINDPARTITIONNEXT) (DWORD dwSearchId, PPD_PARTINFO info);
typedef void  (* PFINDPARTITIONCLOSE) (DWORD dwSearchId);
typedef DWORD (* POPENPARTITION) (DWORD dwStoreId, LPCTSTR partName,PDWORD pdwPartitionId);
typedef void  (* PCLOSEPARTITION) (DWORD dwPartitionId);
typedef DWORD (* PREADPARTITION) (DWORD dwPartitionId, PBYTE pInBuf, DWORD nInBufSize, PBYTE pOutBuf, DWORD nOutBufSize, PDWORD pBytesReturned);
typedef DWORD (* PWRITEPARTITION) (DWORD dwPartitionId, PBYTE pInBuf, DWORD nInBufSize, PBYTE pOutBuf, DWORD nOutBufSize, PDWORD pBytesReturned);
typedef DWORD (* PPD_DEVICEIOCONTROL) (DWORD dwPartitionId, DWORD dwCode, LPVOID pInBuf, DWORD nInBufSize, LPVOID pOutBuf, DWORD nOutBufSize, PDWORD pBytesReturned);
typedef DWORD (* PMEDIACHANGEEVENT) (DWORD StoreId, STORAGE_MEDIA_CHANGE_EVENT_TYPE EventId, STORAGE_MEDIA_ATTACH_RESULT* pResult);
typedef DWORD (* PSTOREIOCONTROL) (DWORD StoreId, DWORD Ioctl, LPVOID pInBuf, DWORD nInBufSize, LPVOID pOutBuf, DWORD nOutBufSize, PDWORD pBytesReturned, LPOVERLAPPED pOverlapped);

class PartitionDriver_t
{
private:
    HMODULE             m_hPartDriver;
    POPENSTORE          m_pOpenStore;
    PCLOSESTORE         m_pCloseStore;
    PFORMATSTORE        m_pFormatStore;
    PISSTOREFORMATTED   m_pIsStoreFormatted;
    PGETSTOREINFO       m_pGetStoreInfo;
    PCREATEPARTITION    m_pCreatePartition;
    PDELETEPARTITION    m_pDeletePartition;
    PRENAMEPARTITION    m_pRenamePartition;
    PSETPARTITIONATTRS  m_pSetPartitionAttrs;
    PSETPARTITIONSIZE   m_pSetPartitionSize;
    PFORMATPARTITION    m_pFormatPartition;
    PGETPARTITIONINFO   m_pGetPartitionInfo;
    PFINDPARTITIONSTART m_pFindPartitionStart;
    PFINDPARTITIONNEXT  m_pFindPartitionNext;
    PFINDPARTITIONCLOSE m_pFindPartitionClose;
    POPENPARTITION      m_pOpenPartition;
    PCLOSEPARTITION     m_pClosePartition;
    PREADPARTITION      m_pReadPartition;
    PWRITEPARTITION     m_pWritePartition;
    PPD_DEVICEIOCONTROL m_pDeviceIoControl;
    PMEDIACHANGEEVENT   m_pMediaChangeEvent;
    PSTOREIOCONTROL     m_pStoreIoControl;
public:    
    PartitionDriver_t () :
        m_hPartDriver (NULL),
        m_pOpenStore (NULL),
        m_pCloseStore (NULL),
        m_pFormatStore (NULL),
        m_pIsStoreFormatted (NULL),
        m_pGetStoreInfo (NULL),
        m_pCreatePartition (NULL),
        m_pDeletePartition (NULL),
        m_pRenamePartition (NULL),
        m_pSetPartitionAttrs (NULL),
        m_pSetPartitionSize (NULL),
        m_pFormatPartition (NULL),
        m_pGetPartitionInfo (NULL),
        m_pFindPartitionStart (NULL),
        m_pFindPartitionNext (NULL),
        m_pFindPartitionClose (NULL),
        m_pOpenPartition (NULL),
        m_pClosePartition (NULL),
        m_pReadPartition (NULL),
        m_pWritePartition (NULL),
        m_pDeviceIoControl (NULL),
        m_pMediaChangeEvent (NULL),
        m_pStoreIoControl (NULL)
    {
    }
    virtual ~PartitionDriver_t ()
    {
        if (m_hPartDriver)
            FreeLibrary ( m_hPartDriver);
    }        
    DWORD LoadPartitionDriver (const WCHAR *szPartDriver);
    
    inline DWORD OpenStore (HANDLE hDisk, PDWORD pdwStoreId)
    {
        __try {
            return m_pOpenStore (hDisk,pdwStoreId);
        } __except (ReportFault (GetExceptionInformation (), 0), EXCEPTION_EXECUTE_HANDLER) {
            return ERROR_GEN_FAILURE;
        }    
    }
    inline void CloseStore (DWORD dwStoreId)
    {
#pragma prefast (disable: 322, "Recover from part driver exception")
        __try {
            m_pCloseStore (dwStoreId);
        } __except (ReportFault (GetExceptionInformation (), 0), EXCEPTION_EXECUTE_HANDLER) {
            DEBUGCHK (0);
        }    
#pragma prefast (pop)        
    }    

    /* stores */
    inline DWORD FormatStore (DWORD dwStoreId)
    {
        __try {
            return m_pFormatStore (dwStoreId);
        } __except (ReportFault (GetExceptionInformation (), 0), EXCEPTION_EXECUTE_HANDLER) {
            return ERROR_GEN_FAILURE;
        }    
    }
    
    inline DWORD IsStoreFormatted (DWORD dwStoreId)
    {
        __try {
            return m_pIsStoreFormatted ( dwStoreId);
        } __except (ReportFault (GetExceptionInformation (), 0), EXCEPTION_EXECUTE_HANDLER) {
            return ERROR_GEN_FAILURE;
        }    
    }
    
    inline DWORD GetStoreInfo (DWORD dwStoreId, PD_STOREINFO *pInfo)
    {
        __try {
            return m_pGetStoreInfo (dwStoreId, pInfo);
        } __except (ReportFault (GetExceptionInformation (), 0), EXCEPTION_EXECUTE_HANDLER) {
            return ERROR_GEN_FAILURE;
        }    
    }

    /* partitions management */
    inline DWORD CreatePartition (DWORD dwStoreId, LPCTSTR partName, BYTE bPartType, SECTORNUM numSectors, BOOL bAuto)
    {
        __try {
            return m_pCreatePartition (dwStoreId, partName, bPartType, numSectors, bAuto);
        } __except (ReportFault (GetExceptionInformation (), 0), EXCEPTION_EXECUTE_HANDLER) {
            return ERROR_GEN_FAILURE;
        }    
    }
    inline DWORD DeletePartition (DWORD dwStoreId, LPCTSTR partName)
    {
        __try {
            return m_pDeletePartition (dwStoreId, partName);
        } __except (ReportFault (GetExceptionInformation (), 0), EXCEPTION_EXECUTE_HANDLER) {
            return ERROR_GEN_FAILURE;
        }    
    }
    
    inline DWORD RenamePartition (DWORD dwStoreId, LPCTSTR oldName, LPCTSTR newName)
    {
        __try {
            return m_pRenamePartition (dwStoreId, oldName, newName);
        } __except (ReportFault (GetExceptionInformation (), 0), EXCEPTION_EXECUTE_HANDLER) {
            return ERROR_GEN_FAILURE;
        }    
    }
    inline DWORD SetPartitionAttrs (DWORD dwStoreId, LPCTSTR partName, DWORD attrs)
    {
        __try {
            return m_pSetPartitionAttrs (dwStoreId, partName, attrs);
        } __except (ReportFault (GetExceptionInformation (), 0), EXCEPTION_EXECUTE_HANDLER) {
            return ERROR_GEN_FAILURE;
        }    
    }    
    inline DWORD SetPartitionSize (DWORD dwStoreId, LPCTSTR partName, SECTORNUM snNumSectors)
    {
        __try {
            return m_pSetPartitionSize (dwStoreId, partName, snNumSectors);
        } __except (ReportFault (GetExceptionInformation (), 0), EXCEPTION_EXECUTE_HANDLER) {
            return ERROR_GEN_FAILURE;
        }    
    }
    
    inline DWORD FormatPartition (DWORD dwStoreId, LPCTSTR partName, BYTE bPartType, BOOL bAuto)
    {
        __try {
            return m_pFormatPartition (dwStoreId, partName, bPartType, bAuto);
        } __except (ReportFault (GetExceptionInformation (), 0), EXCEPTION_EXECUTE_HANDLER) {
            return ERROR_GEN_FAILURE;
        }    
    }
    
    inline DWORD GetPartitionInfo (DWORD dwStoreId, LPCTSTR partName, PD_PARTINFO *pInfo)
    {
        __try {
            return m_pGetPartitionInfo (dwStoreId, partName, pInfo);
        } __except (ReportFault (GetExceptionInformation (), 0), EXCEPTION_EXECUTE_HANDLER) {
            return ERROR_GEN_FAILURE;
        }    
    }
    
    inline DWORD FindPartitionStart (DWORD dwStoreId,PDWORD pdwSearchId)
    {
        __try {
            return m_pFindPartitionStart (dwStoreId, pdwSearchId);
        } __except (ReportFault (GetExceptionInformation (), 0), EXCEPTION_EXECUTE_HANDLER) {
            return ERROR_GEN_FAILURE;
        }    
    }
    
    inline DWORD FindPartitionNext (DWORD dwSearchId, PD_PARTINFO *pInfo)
    {
        __try {
            return m_pFindPartitionNext (dwSearchId, pInfo);
        } __except (ReportFault (GetExceptionInformation (), 0), EXCEPTION_EXECUTE_HANDLER) {
            return ERROR_GEN_FAILURE;
        }    
    }
    
    inline void FindPartitionClose (DWORD dwSearchId)
    {
#pragma prefast (disable: 322, "Recover from part driver exception")
        __try {
            m_pFindPartitionClose (dwSearchId);
        } __except (ReportFault (GetExceptionInformation (), 0), EXCEPTION_EXECUTE_HANDLER) {
        }    
#pragma prefast (pop)        
    }
    

    /* partition I/O */
    inline DWORD OpenPartition (DWORD dwStoreId, LPCTSTR partName, PDWORD pdwPartitionId)
    {
        __try {
            return m_pOpenPartition (dwStoreId, partName, pdwPartitionId);
        } __except (ReportFault (GetExceptionInformation (), 0), EXCEPTION_EXECUTE_HANDLER) {
            return ERROR_GEN_FAILURE;
        }    
    }
    
    inline void ClosePartition (DWORD dwPartitionId)
    {
#pragma prefast (disable: 322, "Recover from part driver exception")
        __try {
            m_pClosePartition (dwPartitionId);
        } __except (ReportFault (GetExceptionInformation (), 0), EXCEPTION_EXECUTE_HANDLER) {
            DEBUGCHK (0);
        }    
#pragma prefast (pop)        
    }
    inline DWORD DeviceIoControl (DWORD dwPartitionId, DWORD dwCode, LPVOID pInBuf, DWORD nInBufSize, LPVOID pOutBuf, DWORD nOutBufSize, PDWORD pBytesReturned)
    {
        __try {
            return m_pDeviceIoControl (dwPartitionId, dwCode, pInBuf, nInBufSize, pOutBuf, nOutBufSize, pBytesReturned);
        } __except (ReportFault (GetExceptionInformation (), 0), EXCEPTION_EXECUTE_HANDLER) {
            return ERROR_GEN_FAILURE;
        }    
    }

    inline DWORD MediaChangeEvent (DWORD StoreId, STORAGE_MEDIA_CHANGE_EVENT_TYPE EventId, STORAGE_MEDIA_ATTACH_RESULT* pResult)
    {
        if (!m_pMediaChangeEvent) {
            // The partition driver does not export the "PD_ MediaChangeEvent" 
            // function.
            return ERROR_NOT_SUPPORTED;
        }

        __try {
            return m_pMediaChangeEvent (StoreId, EventId, pResult);
        } __except (ReportFault (GetExceptionInformation (), 0), EXCEPTION_EXECUTE_HANDLER) {
            return ERROR_GEN_FAILURE;
        }
    }

    inline DWORD StoreIoControl (DWORD StoreId, HANDLE hDisk, DWORD Ioctl, LPVOID pInBuf, DWORD nInBufSize, LPVOID pOutBuf, DWORD nOutBufSize, PDWORD pBytesReturned, LPOVERLAPPED pOverlapped)
    {
        BOOL fRet;

        if (m_pStoreIoControl) {

            // The partition driver exports a PD_StoreIoControl function, so pass
            // it the IOCTL.

            fRet = m_pStoreIoControl (StoreId, Ioctl, pInBuf, nInBufSize, pOutBuf, nOutBufSize, pBytesReturned, pOverlapped);

#ifdef UNDER_CE
        } else if ( (pInBuf || pOutBuf) &&
            (!pInBuf || !::IsKModeAddr (reinterpret_cast<DWORD> (pInBuf))) && 
            (!pOutBuf || !::IsKModeAddr (reinterpret_cast<DWORD> (pOutBuf)))) {

            // All top level buffers exist in user VM. In this case, the IOCTL is
            // being forwarded, so continue the forwarding. The existing caller
            // process will be the caller process as viewed by the driver.

            fRet = ::ForwardDeviceIoControl (hDisk, Ioctl, pInBuf, nInBufSize, pOutBuf, nOutBufSize, pBytesReturned, pOverlapped);
#endif // UNDER_CE

        } else {

            // At least one of the top level buffers is in kernel VM, so perform
            // a normal DeviceIoControl. This process (kernel) will be the caller
            // process as viewed by the driver.

            fRet = ::DeviceIoControl (hDisk, Ioctl, pInBuf, nInBufSize, pOutBuf, nOutBufSize, pBytesReturned, pOverlapped);
        }

        return fRet;
    }

};

class PartitionDisk_t : public MountableDisk_t
{
private:
    DWORD m_dwAttrs;
    PD_PARTINFO m_pi;
    TCHAR m_szFriendlyName[MAX_PATH];
    TCHAR m_szFileSys[FILESYSNAMESIZE];
    
    HANDLE m_hPartition;
    DetectorState_t m_DetectorState;
    BOOL m_fFormatOnMount;
    DWORD m_dwState;

    LONG m_RefCount;
    TCHAR m_szFolderName[FOLDERNAMESIZE];
    TCHAR m_szPartitionName[PARTITIONNAMESIZE];

    PartitionDriver_t *m_pPartDriver; // Just a reference... StoreDisk_t takes care of the lifetime of this object !!!

public:

    DWORD m_dwPartitionId;
    DWORD m_dwStoreId;

    StoreDisk_t *m_pStore;

    DLIST m_link;   //Link to the next PartitionDisk

    PartitionDisk_t (StoreDisk_t *pStore, DWORD dwStoreId, DWORD dwPartitionId, PartitionDriver_t *pPartDriver, const WCHAR *szFolderName, BlockDevice_t* /* pBlockDevice */);
    PartitionDisk_t ();
    virtual ~PartitionDisk_t ();

public:

    inline LONG AddRef ()
    {
        LONG RefCount;
        RefCount = InterlockedIncrement (&m_RefCount);
        return RefCount;
    }

    inline LONG Release ()
    {
        LONG RefCount;
        RefCount = InterlockedDecrement (&m_RefCount);
        if (RefCount == 0)
        {
            delete this;
        }
        return RefCount;
    }

    inline DWORD RefCount ()         { return m_RefCount; }

    inline BOOL IsMediaDetached ()   { return m_dwState & PARTITION_STATE_MEDIA_DETACHED; }
    inline void SetMediaDetached ()  { m_dwState |= PARTITION_STATE_MEDIA_DETACHED; }
    inline void SetMediaAttached ()  { m_dwState &= ~PARTITION_STATE_MEDIA_DETACHED; }
    
    inline BOOL IsMounted ()         { return m_dwState & PARTITION_STATE_MOUNTED; }
    inline void SetMounted ()        { m_dwState |= PARTITION_STATE_MOUNTED; }
    inline void SetUnmounted ()      { m_dwState &= ~PARTITION_STATE_MOUNTED; }

    inline BOOL IsDisabled ()        { return m_dwState & PARTITION_STATE_DISABLED; }
    inline void SetDisabled ()       { m_dwState |= PARTITION_STATE_DISABLED; }
    inline void SetEnabled ()        { m_dwState &= ~PARTITION_STATE_DISABLED; }

    inline const WCHAR* const GetName ()  { return m_szPartitionName; }

public:
    BOOL Init (LPCTSTR szPartitionName, LPCTSTR RootRegKey);

    BOOL Mount (HANDLE hPartition);
    BOOL Unmount ();
    void Enable ();
    void Disable ();
    DWORD CreateHandle (HANDLE *phHandle, HANDLE hProc);


    inline BOOL SetPartitionAttributes (DWORD dwAttrs)
    {
        if (ERROR_SUCCESS == m_pPartDriver->SetPartitionAttrs (m_dwStoreId, m_szPartitionName, dwAttrs))
        {
            m_dwAttrs = dwAttrs;
            return TRUE;
        }
        return FALSE;
    }
    inline BOOL SetPartitionSize (SECTORNUM snNumSectors)
    {
        if (ERROR_SUCCESS == m_pPartDriver->SetPartitionSize (m_dwStoreId, m_szPartitionName, snNumSectors)) {
            // reload partition info
            m_pPartDriver->GetPartitionInfo (m_dwStoreId, m_szPartitionName, &m_pi);
            return TRUE;
        } else {
            return FALSE;
        }
    }
    inline VOID SetFormatOnMount (BOOL fFormatOnMount)
    {
        m_fFormatOnMount = fFormatOnMount;
    }

    BOOL GetPartitionInfo (PARTINFO *pInfo);
    BOOL RenamePartition (LPCTSTR szNewName);
    BOOL FormatPartition (BYTE bPartType, BOOL bAuto);    
    
    virtual void SignalActivity ();

    virtual LRESULT GetName (WCHAR* pDiskName, DWORD NameChars);

    virtual LRESULT DiskIoControl (DWORD IoControlCode, void* pInBuf, 
            DWORD InBufBytes, void* pOutBuf, DWORD OutBufBytes, 
            DWORD* pBytesReturned, OVERLAPPED* pOverlapped);

    virtual LRESULT GetVolumeInfo (CE_VOLUME_INFO* pInfo);

    virtual HANDLE GetDeviceHandle ();    

    LRESULT RunAllDetectors (HKEY hFsKey,
            __out_ecount (FileSystemNameChars) WCHAR* pFileSystemName, 
            size_t FileSystemNameChars);

    LRESULT DetectFileSystem (__out_ecount (FileSystemNameChars) WCHAR* pFileSystemName,
        size_t FileSystemNameChars);

    BOOL ComparePartitions (PartitionDisk_t *pPartition);

    virtual BOOL IsAvailable ();


#ifdef UNDER_CE
    virtual BOOL DoIOControl (DWORD IoControlCode, void* pInBuf, DWORD InBufSize, 
                    void* pOutBuf, DWORD OutBufSize, DWORD* pBytesReturned, OVERLAPPED* pOverlapped);
#endif

};

#endif //__PARTITION_H__
