disk.c

来自「一个类似windows」· C语言 代码 · 共 2,090 行 · 第 1/5 页

  IrpStack = IoGetCurrentIrpStackLocation(Irp);

  /* FIXME: Unlock removable media upon shutdown */


  /* No retry */
  IrpStack->Parameters.Others.Argument4 = (PVOID)0;

  /* Send shutdown or flush request to the port driver */
  Srb->CdbLength = 0;
  if (IrpStack->MajorFunction == IRP_MJ_SHUTDOWN)
    Srb->Function = SRB_FUNCTION_SHUTDOWN;
  else
    Srb->Function = SRB_FUNCTION_FLUSH;

  /* Init completion routine */
  IoSetCompletionRoutine(Irp,
			 ScsiClassIoComplete,
			 Srb,
			 TRUE,
			 TRUE,
			 TRUE);

  /* Prepare next stack location for a call to the port driver */
  IrpStack = IoGetNextIrpStackLocation(Irp);
  IrpStack->MajorFunction = IRP_MJ_SCSI;
  IrpStack->Parameters.Scsi.Srb = Srb;
  Srb->OriginalRequest = Irp;

  /* Call port driver */
  return(IoCallDriver(DeviceExtension->PortDeviceObject, Irp));
}


/**********************************************************************
 * NAME							INTERNAL
 *	DiskClassUpdatePartitionDeviceObjects
 *
 * DESCRIPTION
 *	Deletes, modifies or creates partition device objects.
 *
 * RUN LEVEL
 *	PASSIVE_LEVEL
 *
 * ARGUMENTS
 *	DeviceObject
 *		Pointer to the device.
 *
 *	Irp
 *		Pointer to the IRP
 *
 * RETURN VALUE
 *	None
 */

static VOID
DiskClassUpdatePartitionDeviceObjects(IN PDEVICE_OBJECT DiskDeviceObject,
				      IN PIRP Irp)
{
  PDRIVE_LAYOUT_INFORMATION PartitionList;
  PPARTITION_INFORMATION PartitionEntry;
  PDEVICE_EXTENSION DeviceExtension;
  PDEVICE_EXTENSION DiskDeviceExtension;
  PDISK_DATA DiskData;
  ULONG PartitionCount;
  ULONG PartitionOrdinal;
  ULONG PartitionNumber;
  ULONG LastPartitionNumber;
  ULONG i;
  BOOLEAN Found;
  WCHAR NameBuffer[MAX_PATH];
  UNICODE_STRING DeviceName;
  PDEVICE_OBJECT DeviceObject;
  NTSTATUS Status;

  DPRINT("ScsiDiskUpdatePartitionDeviceObjects() called\n");

  /* Get partition list */
  PartitionList = Irp->AssociatedIrp.SystemBuffer;

  /* Round partition count up by 4 */
  PartitionCount = ((PartitionList->PartitionCount + 3) / 4) * 4;

  /* Remove the partition numbers from the partition list */
  for (i = 0; i < PartitionCount; i++)
    {
      PartitionList->PartitionEntry[i].PartitionNumber = 0;
    }

  DiskDeviceExtension = DiskDeviceObject->DeviceExtension;

  /* Traverse on-disk partition list */
  LastPartitionNumber = 0;
  DeviceExtension = DiskDeviceExtension;
  DiskData = (PDISK_DATA)(DeviceExtension + 1);
  while (TRUE)
    {
      DeviceExtension = DiskData->NextPartition;
      if (DeviceExtension == NULL)
	break;

      /* Get disk data */
      DiskData = (PDISK_DATA)(DeviceExtension + 1);

      /* Update last partition number */
      if (DiskData->PartitionNumber > LastPartitionNumber)
	LastPartitionNumber = DiskData->PartitionNumber;

      /* Ignore unused on-disk partitions */
      if (DeviceExtension->PartitionLength.QuadPart == 0ULL)
	continue;

      Found = FALSE;
      PartitionOrdinal = 0;
      for (i = 0; i < PartitionCount; i++)
	{
	  /* Get current partition entry */
	  PartitionEntry = &PartitionList->PartitionEntry[i];

	  /* Ignore empty (aka unused) or extended partitions */
	  if (PartitionEntry->PartitionType == PARTITION_ENTRY_UNUSED ||
	      IsContainerPartition (PartitionEntry->PartitionType))
	    continue;

	  PartitionOrdinal++;

	  /* Check for matching partition start offset and length */
	  if ((PartitionEntry->StartingOffset.QuadPart !=
	       DeviceExtension->StartingOffset.QuadPart) ||
	      (PartitionEntry->PartitionLength.QuadPart !=
	       DeviceExtension->PartitionLength.QuadPart))
	    continue;

	  DPRINT("Found matching partition entry for partition %lu\n",
		  DiskData->PartitionNumber);

	  /* Found matching partition */
	  Found = TRUE;

	  /* Update partition number in partition list */
	  PartitionEntry->PartitionNumber = DiskData->PartitionNumber;
	  break;
	}

      if (Found == TRUE)
	{
	  /* Get disk data for current partition */
	  DiskData = (PDISK_DATA)(DeviceExtension + 1);

	  /* Update partition type if partiton will be rewritten */
	  if (PartitionEntry->RewritePartition == TRUE)
	    DiskData->PartitionType = PartitionEntry->PartitionType;

	  /* Assign new partiton ordinal */
	  DiskData->PartitionOrdinal = PartitionOrdinal;

	  DPRINT("Partition ordinal %lu was assigned to partition %lu\n",
		 DiskData->PartitionOrdinal,
		 DiskData->PartitionNumber);
	}
      else
	{
	  /* Delete this partition */
	  DeviceExtension->PartitionLength.QuadPart = 0ULL;

	  DPRINT("Deleting partition %lu\n",
		 DiskData->PartitionNumber);
	}
    }

  /* Traverse partiton list and create new partiton devices */
  PartitionOrdinal = 0;
  for (i = 0; i < PartitionCount; i++)
    {
      /* Get current partition entry */
      PartitionEntry = &PartitionList->PartitionEntry[i];

      /* Ignore empty (aka unused) or extended partitions */
      if (PartitionEntry->PartitionType == PARTITION_ENTRY_UNUSED ||
	  IsContainerPartition (PartitionEntry->PartitionType))
	continue;

      PartitionOrdinal++;

      /* Ignore unchanged partition entries */
      if (PartitionEntry->RewritePartition == FALSE)
	continue;

      /* Check for an unused device object */
      PartitionNumber = 0;
      DeviceExtension = DiskDeviceExtension;
      DiskData = (PDISK_DATA)(DeviceExtension + 1);
      while (TRUE)
	{
	  DeviceExtension = DiskData->NextPartition;
	  if (DeviceExtension == NULL)
	    break;

	  /* Get partition disk data */
	  DiskData = (PDISK_DATA)(DeviceExtension + 1);

	  /* Found a free (unused) partition (device object) */
	  if (DeviceExtension->PartitionLength.QuadPart == 0ULL)
	    {
	      PartitionNumber = DiskData->PartitionNumber;
	      break;
	    }
	}

      if (PartitionNumber == 0)
	{
	  /* Create a new partition device object */
	  DPRINT("Create new partition device object\n");

	  /* Get new partiton number */
	  LastPartitionNumber++;
	  PartitionNumber = LastPartitionNumber;

	  /* Create partition device object */
	  swprintf(NameBuffer,
		   L"\\Device\\Harddisk%lu\\Partition%lu",
		   DiskDeviceExtension->DeviceNumber,
		   PartitionNumber);
	  RtlInitUnicodeString(&DeviceName,
			       NameBuffer);

	  Status = IoCreateDevice(DiskDeviceObject->DriverObject,
				  sizeof(DEVICE_EXTENSION) + sizeof(DISK_DATA),
				  &DeviceName,
				  FILE_DEVICE_DISK,
				  0,
				  FALSE,
				  &DeviceObject);
	  if (!NT_SUCCESS(Status))
	    {
	      DPRINT("IoCreateDevice() failed (Status %lx)\n", Status);
	      continue;
	    }

	  DeviceObject->Flags |= DO_DIRECT_IO;
	  DeviceObject->StackSize = DiskDeviceObject->StackSize;
	  DeviceObject->Characteristics = DiskDeviceObject->Characteristics;
	  DeviceObject->AlignmentRequirement = DiskDeviceObject->AlignmentRequirement;

	  /* Initialize device extension */
	  DeviceExtension = DeviceObject->DeviceExtension;
	  RtlCopyMemory(DeviceExtension,
			DiskDeviceObject->DeviceExtension,
			sizeof(DEVICE_EXTENSION));
	  DeviceExtension->DeviceObject = DeviceObject;

	  /* Initialize lookaside list for SRBs */
	  ScsiClassInitializeSrbLookasideList(DeviceExtension,
					      8);

	  /* Link current partition device extension to previous disk data */
	  DiskData->NextPartition = DeviceExtension;
	  DiskData = (PDISK_DATA)(DeviceExtension + 1);
	  DiskData->NextPartition = NULL;
	}
      else
	{
	  /* Reuse an existing partition device object */
	  DPRINT("Reuse an exisiting partition device object\n");
	  DiskData = (PDISK_DATA)(DeviceExtension + 1);
	}

      /* Update partition data and device extension */
      DiskData->PartitionNumber = PartitionNumber;
      DiskData->PartitionOrdinal = PartitionOrdinal;
      DiskData->PartitionType = PartitionEntry->PartitionType;
      DiskData->BootIndicator = PartitionEntry->BootIndicator;
      DiskData->HiddenSectors = PartitionEntry->HiddenSectors;
      DeviceExtension->StartingOffset = PartitionEntry->StartingOffset;
      DeviceExtension->PartitionLength = PartitionEntry->PartitionLength;

      /* Update partition number in the partition list */
      PartitionEntry->PartitionNumber = PartitionNumber;

      DPRINT("Partition ordinal %lu was assigned to partition %lu\n",
	     DiskData->PartitionOrdinal,
	     DiskData->PartitionNumber);
    }

  DPRINT("ScsiDiskUpdatePartitionDeviceObjects() done\n");
}


/**********************************************************************
 * NAME							INTERNAL
 *	ScsiDiskSearchForDisk
 *
 * DESCRIPTION
 *	Searches the hardware tree for the given disk.
 *
 * RUN LEVEL
 *	PASSIVE_LEVEL
 *
 * ARGUMENTS
 *	DeviceExtension
 *		Disk device extension.
 *
 *	BusKey
 *		Handle to the hardware bus key.
 *
 *	DetectedDiskNumber
 *		Returned disk number.
 *
 * RETURN VALUE
 *	TRUE: Disk was found.
 *	FALSE: Search failed.
 */

static BOOLEAN
ScsiDiskSearchForDisk(IN PDEVICE_EXTENSION DeviceExtension,
		      IN HANDLE BusKey,
		      OUT PULONG DetectedDiskNumber)
{
  PKEY_VALUE_FULL_INFORMATION ValueData;
  OBJECT_ATTRIBUTES ObjectAttributes;
  PDISK_DATA DiskData;
  UNICODE_STRING IdentifierString;
  UNICODE_STRING NameString;
  HANDLE BusInstanceKey;
  HANDLE ControllerKey;
  HANDLE DiskKey;
  HANDLE DiskInstanceKey;
  ULONG BusNumber;
  ULONG ControllerNumber;
  ULONG DiskNumber;
  ULONG Length;
  WCHAR Buffer[32];
  BOOLEAN DiskFound;
  NTSTATUS Status;

  DPRINT("ScsiDiskSearchForDiskData() called\n");

  DiskFound = FALSE;

  /* Enumerate buses */
  for (BusNumber = 0; ; BusNumber++)
    {
      /* Open bus instance subkey */
      swprintf(Buffer,
	       L"%lu",
	       BusNumber);

      RtlInitUnicodeString(&NameString,
			   Buffer);

      InitializeObjectAttributes(&ObjectAttributes,
				 &NameString,
				 OBJ_CASE_INSENSITIVE,
				 BusKey,
				 NULL);

      Status = ZwOpenKey(&BusInstanceKey,
			 KEY_READ,
			 &ObjectAttributes);
      if (!NT_SUCCESS(Status))
	{
	  break;
	}

      /* Open 'DiskController' subkey */
      RtlInitUnicodeString(&NameString,
			   L"DiskController");

      InitializeObjectAttributes(&ObjectAttributes,
				 &NameString,
				 OBJ_CASE_INSENSITIVE,
				 BusInstanceKey,
				 NULL);

      Status = ZwOpenKey(&ControllerKey,
			 KEY_READ,
			 &ObjectAttributes);
      if (!NT_SUCCESS(Status))
	{
	  ZwClose(BusInstanceKey);
	  continue;
	}

      /* Enumerate controllers */
      for (ControllerNumber = 0; ; ControllerNumber++)
	{
	  /* Open 'DiskPeripheral' subkey */
	  swprintf(Buffer,
		   L"%lu\\DiskPeripheral",
		   ControllerNumber);

	  RtlInitUnicodeString(&NameString,
			       Buffer);

	  InitializeObjectAttributes(&ObjectAttributes,
				     &NameString,
				     OBJ_CASE_INSENSITIVE,
				     ControllerKey,
				     NULL);

	  Status = ZwOpenKey(&DiskKey,
			     KEY_READ,
			     &ObjectAttributes);
	  if (!NT_SUCCESS(Status))
	    {
	      break;
	    }

	  /* Enumerate disks */
	  for (DiskNumber = 0; ; DiskNumber++)
	    {
	      /* Open disk instance subkey */
	      swprintf(Buffer,
		       L"%lu",
		       DiskNumber);

	      RtlInitUnicodeString(&NameString,