drivevolume.cpp

ReactOS是一些高手根据Windows XP的内核编写出的类XP。内核实现机理和API函数调用几乎相同。甚至可以兼容XP的程序。喜欢研究系统内核的人可以看一看。

#include "DriveVolume.h"


DriveVolume::DriveVolume ()
{
    Handle = INVALID_HANDLE_VALUE;
    BitmapDetail = NULL;
    return;
}


DriveVolume::~DriveVolume ()
{
    Close ();
    Directories.clear ();
    Files.clear ();
    return;
}


void DriveVolume::Close (void)
{
    if (Handle != INVALID_HANDLE_VALUE)
    {
        CloseHandle (Handle);
        Handle = INVALID_HANDLE_VALUE;
    }

    if (BitmapDetail != NULL)
    {
        free (BitmapDetail);
        BitmapDetail = NULL;
    }

    return;
}


// "Name" should be the drive letter followed by a colon. ie, "c:"
// It's a string to allow for further expansion (ie, defragging over the network?)
// or some other baloney reason
bool DriveVolume::Open (wstring Name)
{
    wchar_t FileName[100];
    bool ReturnVal;

    swprintf (FileName, L"\\\\.\\%s", Name.c_str());
    RootPath = Name.c_str();
    RootPath += L"\\";

    Handle = CreateFile 
    (
        FileName,
        MAXIMUM_ALLOWED,                          // access
        FILE_SHARE_READ | FILE_SHARE_WRITE,       // share type
        NULL,                                     // security descriptor
        OPEN_EXISTING,                            // open type
        NULL,                                     // attributes (none)
        NULL                                      // template
    );

    if (Handle == INVALID_HANDLE_VALUE)
        ReturnVal = false;
    else
    {
        wchar_t  VolName[64];
        DWORD VolSN;
        DWORD VolMaxFileLen;
        DWORD FSFlags;
        wchar_t  FSName[64];
        BOOL  Result;

        ReturnVal = true;
        Result = GetVolumeInformation
        (
            RootPath.c_str(),
            VolName,
            sizeof (VolName),
            &VolSN,
            &VolMaxFileLen,
            &FSFlags,
            FSName,
            sizeof (FSName)
        );

        if (Result)
        {
            wchar_t SerialText[10];

            VolInfo.FileSystem = FSName;
            VolInfo.MaxNameLen = VolMaxFileLen;
            VolInfo.Name       = VolName;

            swprintf (SerialText, L"%x-%x", (VolSN & 0xffff0000) >> 16, 
                VolSN & 0x0000ffff);

            wcsupr (SerialText);
            VolInfo.Serial     = SerialText;
        }
        else
        {
            VolInfo.FileSystem = L"(Unknown)";
            VolInfo.MaxNameLen = 255;
            VolInfo.Name       = L"(Unknown)";
            VolInfo.Serial     = L"(Unknown)";
        }
    }

    return (ReturnVal);
}


bool DriveVolume::ObtainInfo (void)
{
    BOOL Result;
    DWORD BytesGot;
    uint64 nan;

    BytesGot = 0;
    ZeroMemory (&Geometry, sizeof (Geometry));
    Result = DeviceIoControl
    (
        Handle,
        IOCTL_DISK_GET_DRIVE_GEOMETRY,
        NULL,
        0,
        &Geometry,
        sizeof (Geometry),
        &BytesGot,
        NULL
    );

    // Call failed? Aww :(
    if (!Result)
        return (false);

    // Get cluster size
    DWORD SectorsPerCluster;
    DWORD BytesPerSector;
    DWORD FreeClusters;
    DWORD TotalClusters;

    Result = GetDiskFreeSpace
    (
        RootPath.c_str(),
        &SectorsPerCluster,
        &BytesPerSector,
        &FreeClusters,
        &TotalClusters
    );

    // Failed? Weird.
    if (!Result)
        return (false);

    VolInfo.ClusterSize = SectorsPerCluster * BytesPerSector;        

    Result = GetDiskFreeSpaceEx
    (
        RootPath.c_str(),
        (PULARGE_INTEGER)&nan,
        (PULARGE_INTEGER)&VolInfo.TotalBytes,
        (PULARGE_INTEGER)&VolInfo.FreeBytes
    );

    return (true);
}


// Get bitmap, several clusters at a time ...
#define CLUSTERS 4096
bool DriveVolume::GetBitmap (void)
{
    STARTING_LCN_INPUT_BUFFER StartingLCN;
    VOLUME_BITMAP_BUFFER *Bitmap = NULL;
    uint32 BitmapSize;
    DWORD BytesReturned;
    BOOL Result;

    StartingLCN.StartingLcn.QuadPart = 0;

    // Allocate buffer
    // Call FSCTL_GET_VOLUME_BITMAP once with a very small buffer
    // This will leave the total number of clusters in Bitmap->BitmapSize and we can
    // then correctly allocate based off that
    // I suppose this won't work if your drive has only 40 clusters on it or so :)
    BitmapSize = sizeof (VOLUME_BITMAP_BUFFER) + 4;
    Bitmap = (VOLUME_BITMAP_BUFFER *) malloc (BitmapSize);

    Result = DeviceIoControl
    (
        Handle,
        FSCTL_GET_VOLUME_BITMAP,
        &StartingLCN,
        sizeof (StartingLCN),
        Bitmap,
        BitmapSize,
        &BytesReturned,
        NULL
    );

    // Bad result?
    if (Result == FALSE  &&  GetLastError () != ERROR_MORE_DATA)
    {
        //wprintf ("\nDeviceIoControl returned false, GetLastError() was not ERROR_MORE_DATA\n");
        free (Bitmap);
        return (false);
    }

    // Otherwise, we're good
    BitmapSize = sizeof (VOLUME_BITMAP_BUFFER) + (Bitmap->BitmapSize.QuadPart / 8) + 1;
    Bitmap = (VOLUME_BITMAP_BUFFER *) realloc (Bitmap, BitmapSize);
    Result = DeviceIoControl
    (
        Handle,
        FSCTL_GET_VOLUME_BITMAP,
        &StartingLCN,
        sizeof (StartingLCN),
        Bitmap,
        BitmapSize,
        &BytesReturned,
        NULL
    );

    DWORD LastError = GetLastError ();

    if (Result == FALSE)
    {
        wprintf (L"\nCouldn't properly read volume bitmap\n");
        free (Bitmap);
        return (false);
    }

    // Convert to a L'quick use' bitmap
    //const int BitShift[] = { 1, 2, 4, 8, 16, 32, 64, 128 };

    VolInfo.ClusterCount = Bitmap->BitmapSize.QuadPart;

    if (BitmapDetail != NULL)
        free (BitmapDetail);

    BitmapDetail = (uint32 *) malloc (sizeof(uint32) * (1 + (VolInfo.ClusterCount / 32)));
    memcpy (BitmapDetail, Bitmap->Buffer, sizeof(uint32) * (1 + (VolInfo.ClusterCount / 32)));

    /*
    BitmapDetail = (Cluster *) malloc (VolInfo.ClusterCount * sizeof (Cluster));
    for (uint64 i = 0; i < VolInfo.ClusterCount; i++)
    {
        if (Bitmap->Buffer[i / 8] & BitShift[i % 8])
            BitmapDetail[i].Allocated = true;
        else
            BitmapDetail[i].Allocated = false;
    }
    */

    free (Bitmap);
    return (true);
}


bool DriveVolume::IsClusterUsed (uint64 Cluster)
{
    return ((BitmapDetail[Cluster / 32] & (1 << (Cluster % 32))) ? true : false);
    //return (BitmapDetail[Cluster].Allocated);
}


void DriveVolume::SetClusterUsed (uint64 Cluster, bool Used)
{
    if (Used)
        BitmapDetail[Cluster / 32] |= (1 << (Cluster % 32));
    else
        BitmapDetail[Cluster / 32] &= ~(1 << (Cluster % 32));

    return;
}


typedef struct
{
    DriveVolume *Volume;
    double       *Percent;
    bool        *QuitMonitor;
    uint64       ClusterCount;
    uint64       ClusterProgress;
} BuildDBInfo;


bool DriveVolume::BuildFileList (bool &QuitMonitor, double &Percent)
{
    BuildDBInfo Info;

    Files.clear ();
    Directories.clear ();
    Directories.push_back (RootPath);

    Info.Volume = this;
    Info.QuitMonitor = &QuitMonitor;
    Info.ClusterCount = (GetVolumeInfo().TotalBytes - GetVolumeInfo().FreeBytes) / (uint64)GetVolumeInfo().ClusterSize;
    Info.ClusterProgress = 0;
    Info.Percent = &Percent;

    ScanDirectory (RootPath, BuildDBCallback, &Info);

    if (QuitMonitor == true)
    {
        Directories.resize (0);
        Files.resize (0);
    }

    return (true);
}


// UserData = pointer to BuildDBInfo instance
bool BuildDBCallback (FileInfo &Info, HANDLE &FileHandle, void *UserData)
{
    BuildDBInfo *DBInfo = (BuildDBInfo *) UserData;
    DriveVolume *Vol = DBInfo->Volume;

    Vol->Files.push_back (Info);

    if (*(DBInfo->QuitMonitor) == true)
        return (false);

    DBInfo->ClusterProgress += (uint64)Info.Clusters;
    *(DBInfo->Percent) = 
        ((double)DBInfo->ClusterProgress / (double)DBInfo->ClusterCount) * 100.0f;

    return (true);
}


wstring &DriveVolume::GetDBDir (uint32 Indice)
{
    return (Directories[Indice]);
}


uint32 DriveVolume::GetDBDirCount (void)
{
    return (Directories.size());
}


FileInfo &DriveVolume::GetDBFile (uint32 Indice)
{
    return (Files[Indice]);
}


uint32 DriveVolume::GetDBFileCount (void)
{
    return (Files.size());
}


uint32 DriveVolume::RemoveDBFile (uint32 Indice)
{
    vector<FileInfo>::iterator it;

    it = Files.begin() + Indice;
    Files.erase (it);
    return (GetDBFileCount());
}


bool DriveVolume::ScanDirectory (wstring DirPrefix, ScanCallback Callback, void *UserData)
{
    WIN32_FIND_DATA FindData;
    HANDLE          FindHandle;
    wstring          SearchString;
    uint32          DirIndice;

    DirIndice = Directories.size() - 1;

    SearchString = DirPrefix;
    SearchString += L"*.*";
    ZeroMemory (&FindData, sizeof (FindData));
    FindHandle = FindFirstFile (SearchString.c_str(), &FindData);

    if (FindHandle == INVALID_HANDLE_VALUE)
        return (false);

    do
    {
        FileInfo Info;
        HANDLE   Handle;
        bool     CallbackResult;

        Handle = INVALID_HANDLE_VALUE;

        // First copy over the easy stuff.
        Info.Name = FindData.cFileName;

        // DonLL't ever include '.L' and '..'
        if (Info.Name == L"."  ||  Info.Name == L"..")
            continue;

        //Info.FullName = DirPrefix + Info.Name;
        Info.Size      = (uint64)FindData.nFileSizeLow + ((uint64)FindData.nFileSizeHigh << (uint64)32);
        Info.DirIndice = DirIndice;

        Info.Attributes.Archive    = (FindData.dwFileAttributes & FILE_ATTRIBUTE_ARCHIVE)       ? 1 : 0;