geometry.c

The Disk sample is used with Classpnp.sys as disk driver. The sample supports Plug and Play, Power M

                RtlInitUnicodeString(&unicodeString, buffer);

                InitializeObjectAttributes(&objectAttributes,
                                           &unicodeString,
                                           OBJ_CASE_INSENSITIVE,
                                           diskKey,
                                           NULL);

                status = ZwOpenKey(&targetKey, KEY_READ, &objectAttributes);

                if(!NT_SUCCESS(status)) {
                    DebugPrint((1, "DiskScanBusDetectInfo: Error %#08lx "
                                   "opening target key\n", 
                                status));
                    break;
                }

                status = DiskSaveBusDetectInfo(DriverObject, 
                                               targetKey,
                                               diskNumber);

                ZwClose(targetKey);
            }

            ZwClose(diskKey);
        }
        ZwClose(adapterKey);
    }
    return;
}


NTSTATUS
DiskSaveBusDetectInfo(
    IN PDRIVER_OBJECT DriverObject,
    IN HANDLE TargetKey,
    IN ULONG DiskNumber
    )
/*++

Routine Description:

    This routine will transfer the firmware/ntdetect reported information 
    in the specified target key into the appropriate entry in the 
    DetectInfoList.

Arguments:

    DriverObject - the object for this driver.
    
    TargetKey - the key for the disk being saved.    

    DiskNumber - the ordinal of the entry in the DiskPeripheral tree for this 
                 entry

Return Value:

    status    

--*/
{
    PDISK_DETECT_INFO diskInfo;

    UNICODE_STRING unicodeString;

    PKEY_VALUE_FULL_INFORMATION keyData;
    ULONG length;

    NTSTATUS status;

    PAGED_CODE();

    diskInfo = &(DetectInfoList[DiskNumber]);

    if(diskInfo->Initialized) {

        ASSERT(FALSE);
        DebugPrint((1, "DiskSaveBusDetectInfo: disk entry %#x already has a "
                        "signature of %#08lx and mbr checksum of %#08lx\n", 
                        DiskNumber,
                        diskInfo->Signature,
                        diskInfo->MbrCheckSum));
        return STATUS_UNSUCCESSFUL;
    }

    RtlInitUnicodeString(&unicodeString, L"Identifier");

    keyData = ExAllocatePoolWithTag(PagedPool,
                                    VALUE_BUFFER_SIZE,
                                    DISK_TAG_UPDATE_GEOM);

    if(keyData == NULL) {
        DebugPrint((1, "DiskSaveBusDetectInfo: Couldn't allocate space for "
                       "registry data\n"));
        return STATUS_INSUFFICIENT_RESOURCES;
    }

    //
    // Get disk peripheral identifier.
    //

    status = ZwQueryValueKey(TargetKey,
                             &unicodeString,
                             KeyValueFullInformation,
                             keyData,
                             VALUE_BUFFER_SIZE,
                             &length);

    if(!NT_SUCCESS(status)) {
        DebugPrint((1, "DiskSaveBusDetectInfo: Error %#08lx getting "
                       "Identifier\n",
                    status));
        ExFreePool(keyData);
        return status;

    } else if (keyData->DataLength < 9*sizeof(WCHAR)) {

        //
        // the data is too short to use (we subtract 9 chars in normal path)
        //
        DebugPrint((1, "DiskSaveBusDetectInfo: Saved data was invalid, "
                    "not enough data in registry!\n"));
        ExFreePool(keyData);
        return STATUS_UNSUCCESSFUL;
        
    } else {

        UNICODE_STRING identifier;
        ULONG value;

        //
        // Complete unicode string.
        //
    
        identifier.Buffer = (PWSTR) ((PUCHAR)keyData + keyData->DataOffset);
        identifier.Length = (USHORT) keyData->DataLength;
        identifier.MaximumLength = (USHORT) keyData->DataLength;
    
        //
        // Get the first value out of the identifier - this will be the MBR 
        // checksum.
        //
    
        status = RtlUnicodeStringToInteger(&identifier, 16, &value);
    
        if(!NT_SUCCESS(status)) {
            DebugPrint((1, "DiskSaveBusDetectInfo: Error %#08lx converting "
                           "identifier %wZ into MBR xsum\n", 
                           status,
                           &identifier));
            ExFreePool(keyData);
            return status;
        }
    
        diskInfo->MbrCheckSum = value;
    
        //
        // Shift the string over to get the disk signature
        //
    
        identifier.Buffer += 9;
        identifier.Length -= 9 * sizeof(WCHAR);
        identifier.MaximumLength -= 9 * sizeof(WCHAR);
    
        status = RtlUnicodeStringToInteger(&identifier, 16, &value);
    
        if(!NT_SUCCESS(status)) {
            DebugPrint((1, "DiskSaveBusDetectInfo: Error %#08lx converting "
                           "identifier %wZ into disk signature\n", 
                           status,
                           &identifier));
            ExFreePool(keyData);
            value = 0;
        }

        diskInfo->Signature = value;
    }

    //
    // Here is where we would save away the extended int13 data.
    //

    //
    // Mark this entry as initialized so we can make sure not to do it again.
    //

    diskInfo->Initialized = TRUE;


    return STATUS_SUCCESS;
}


DISK_GEOMETRY_SOURCE
DiskUpdateGeometry(
    IN PFUNCTIONAL_DEVICE_EXTENSION FdoExtension
    )
/*++

Routine Description:

    This routine checks the DetectInfoList saved away during disk driver init
    to see if any geometry information was reported for this drive.  If the 
    geometry data exists (determined by matching non-zero signatures or 
    non-zero MBR checksums) then it will be saved in the RealGeometry member
    of the disk data block.

    ClassReadDriveCapacity MUST be called after calling this routine to update
    the cylinder count based on the size of the disk and the presence of any 
    disk management software.

Arguments:

    DeviceExtension - Supplies a pointer to the device information for disk.

Return Value:

    Inidicates whether the "RealGeometry" in the data block is now valid.

--*/

{
    PDISK_DATA diskData = FdoExtension->CommonExtension.DriverData;

    ULONG i;
    PDISK_DETECT_INFO diskInfo;

    BOOLEAN found = FALSE;

    NTSTATUS status;

    PAGED_CODE();


    ASSERT(FdoExtension->CommonExtension.IsFdo);
    ASSERT((FdoExtension->DeviceObject->Characteristics & FILE_REMOVABLE_MEDIA) == 0);

    //
    // If we've already set a non-default geometry for this drive then there's 
    // no need to try and update again.
    //

    if(diskData->GeometrySource != DiskGeometryUnknown) {
        return diskData->GeometrySource;
    }

    //
    // Scan through the saved detect info to see if we can find a match 
    // for this device.
    //

    for(i = 0; i < DetectInfoCount; i++) {

        ASSERT(DetectInfoList != NULL);

        diskInfo = &(DetectInfoList[i]);

        if((diskData->Mbr.Signature != 0) && 
           (diskData->Mbr.Signature == diskInfo->Signature)) { 
            DebugPrint((1, "DiskUpdateGeometry: found match for signature "
                           "%#08lx\n",
                        diskData->Mbr.Signature));
            found = TRUE;
            break;
        } else if((diskData->Mbr.Signature == 0) && 
                  (diskData->Mbr.MbrCheckSum != 0) &&
                  (diskData->Mbr.MbrCheckSum == diskInfo->MbrCheckSum)) {
            DebugPrint((1, "DiskUpdateGeometry: found match for xsum %#08lx\n",
                        diskData->Mbr.MbrCheckSum));
            found = TRUE;
            break;
        }
    }

    if(found) {

        ULONG cylinders;
        ULONG sectorsPerTrack;
        ULONG tracksPerCylinder;

        ULONG sectors;
        ULONG length;

        //
        // Point to the array of drive parameters.
        //
        
        cylinders = diskInfo->DriveParameters.MaxCylinders + 1;
        sectorsPerTrack = diskInfo->DriveParameters.SectorsPerTrack;
        tracksPerCylinder = diskInfo->DriveParameters.MaxHeads + 1;

        //
        // Since the BIOS may not report the full drive, recalculate the drive
        // size based on the volume size and the BIOS values for tracks per
        // cylinder and sectors per track..
        //
    
        length = tracksPerCylinder * sectorsPerTrack;
    
        if (length == 0) {
    
            //
            // The BIOS information is bogus.
            //
    
            DebugPrint((1, "DiskUpdateGeometry: H (%d) or S(%d) is zero\n",
                        tracksPerCylinder, sectorsPerTrack));
            return FALSE;
        }
    
        //
        // since we are copying the structure RealGeometry here, we should
        // really initialize all the fields, especially since a zero'd
        // BytesPerSector field would cause a trap in xHalReadPartitionTable()
        //

        diskData->RealGeometry = FdoExtension->DiskGeometry;
        
        //
        // Save the geometry information away in the disk data block and 
        // set the bit indicating that we found a valid one.
        //

        diskData->RealGeometry.SectorsPerTrack = sectorsPerTrack;
        diskData->RealGeometry.TracksPerCylinder = tracksPerCylinder;
        diskData->RealGeometry.Cylinders.QuadPart = (LONGLONG)cylinders;

        DebugPrint((1, "DiskUpdateGeometry: BIOS spt %#x, #heads %#x, "
                       "#cylinders %#x\n",
                   sectorsPerTrack, tracksPerCylinder, cylinders));

        diskData->GeometrySource = DiskGeometryFromBios;
        diskInfo->Device = FdoExtension->DeviceObject;

    } else {
    
       DebugPrint((1, "DiskUpdateGeometry: no match found for signature %#08lx\n", diskData->Mbr.Signature));
    }

    if(diskData->GeometrySource == DiskGeometryUnknown) {

        //
        // We couldn't find a geometry from the BIOS.  Check with the port 
        // driver and see if it can provide one.
        //

        status = DiskGetPortGeometry(FdoExtension, &(diskData->RealGeometry));

        if(NT_SUCCESS(status)) {

            //
            // Check the geometry to make sure it's valid.
            //

            if((diskData->RealGeometry.TracksPerCylinder * 
                diskData->RealGeometry.SectorsPerTrack) != 0) {

                diskData->GeometrySource = DiskGeometryFromPort;
                DebugPrint((1, "DiskUpdateGeometry: using Port geometry for disk %#p\n", FdoExtension));
                
                if (diskData->RealGeometry.BytesPerSector == 0) {
                    
                    DebugPrint((0, "DiskDriverReinit: Port driver failed to "
                                "set BytesPerSector in the RealGeometry\n"));
                    diskData->RealGeometry.BytesPerSector =
                        FdoExtension->DiskGeometry.BytesPerSector;
                    if (diskData->RealGeometry.BytesPerSector == 0) {
                        ASSERT(!"BytesPerSector is still zero!");
                    }

                }
            }
        }
    }

    //
    // If we came up with a "real" geometry for this drive then set it in the 
    // device extension.
    //

    if(diskData->GeometrySource != DiskGeometryUnknown) {
        FdoExtension->DiskGeometry = diskData->RealGeometry;

        //
        // Increment the count of used geometry entries.
        //

        InterlockedIncrement(&DetectInfoUsedCount);
    }

    return diskData->GeometrySource;
}


NTSTATUS
DiskUpdateRemovableGeometry (
    IN PFUNCTIONAL_DEVICE_EXTENSION FdoExtension
    )

/*++

Routine Description:

    This routine updates the geometry of the disk.  It will query the port 
    driver to see if it can provide any geometry info.  If not it will use 
    the current head & sector count.
    
    Based on these values & the capacity of the drive as reported by 
    ClassReadDriveCapacity it will determine a new cylinder count for the 
    device.

Arguments:

    Fdo - Supplies the functional device object whos size needs to be updated.

Return Value:

    Returns the status of the opertion.

--*/
{
    PCOMMON_DEVICE_EXTENSION commonExtension = &(FdoExtension->CommonExtension);
    PDISK_DATA diskData = commonExtension->DriverData;
    PDISK_GEOMETRY geometry = &(diskData->RealGeometry);

    NTSTATUS status;

    PAGED_CODE();

    ASSERT_FDO(commonExtension->DeviceObject);
    if (FdoExtension->DeviceDescriptor) {
        ASSERT(FdoExtension->DeviceDescriptor->RemovableMedia);
    }
    ASSERT(TEST_FLAG(FdoExtension->DeviceObject->Characteristics,
                     FILE_REMOVABLE_MEDIA));

    //
    // Attempt to determine the disk geometry.  First we'll check with the 
    // port driver to see what it suggests for a value.
    //

    status = DiskGetPortGeometry(FdoExtension, geometry);

    if(NT_SUCCESS(status) && 
       ((geometry->TracksPerCylinder * geometry->SectorsPerTrack) != 0)) {

        FdoExtension->DiskGeometry = (*geometry);
    }

    return status;
}


NTSTATUS
DiskGetPortGeometry(
    IN PFUNCTIONAL_DEVICE_EXTENSION FdoExtension,
    OUT PDISK_GEOMETRY Geometry
    )
/*++

Routine Description:

    This routine will query the port driver for disk geometry.  Some port 
    drivers (in particular IDEPORT) may be able to provide geometry for the 
    device.
    
Arguments:

    FdoExtension - the device object for the disk.
    
    Geometry - a structure to save the geometry information into (if any is 
               available)
               
Return Value:
    
    STATUS_SUCCESS if geometry information can be provided or 
    error status indicating why it can't.
    
--*/
{
    PCOMMON_DEVICE_EXTENSION commonExtension = &(FdoExtension->CommonExtension);
    PIRP irp;
    PIO_STACK_LOCATION irpStack;
    KEVENT event;

    NTSTATUS status;

    PAGED_CODE();

    //
    // Build an irp to send IOCTL_DISK_GET_DRIVE_GEOMETRY to the lower driver.