]> git.proxmox.com Git - mirror_edk2.git/blobdiff - MdeModulePkg/Universal/Disk/RamDiskDxe/RamDiskProtocol.c
MdeModulePkg/RamDiskDxe: Restrict on RAM disk size (CVE-2018-12180)
[mirror_edk2.git] / MdeModulePkg / Universal / Disk / RamDiskDxe / RamDiskProtocol.c
index de117212329a9a68ee2583a8fa6739ffebc3c7f7..e8250d5c1bcc6e68a30f6b5fba935cb03c11fef0 100644 (file)
@@ -1,7 +1,7 @@
 /** @file\r
   The realization of EFI_RAM_DISK_PROTOCOL.\r
 \r
-  Copyright (c) 2016, Intel Corporation. All rights reserved.<BR>\r
+  Copyright (c) 2016 - 2019, Intel Corporation. All rights reserved.<BR>\r
   (C) Copyright 2016 Hewlett Packard Enterprise Development LP<BR>\r
   This program and the accompanying materials\r
   are licensed and made available under the terms and conditions of the BSD License\r
@@ -31,6 +31,9 @@ MEDIA_RAM_DISK_DEVICE_PATH  mRamDiskDeviceNodeTemplate = {
   }\r
 };\r
 \r
+BOOLEAN  mRamDiskSsdtTableKeyValid = FALSE;\r
+UINTN    mRamDiskSsdtTableKey;\r
+\r
 \r
 /**\r
   Initialize the RAM disk device node.\r
@@ -58,6 +61,504 @@ RamDiskInitDeviceNode (
 }\r
 \r
 \r
+/**\r
+  Initialize and publish NVDIMM root device SSDT in ACPI table.\r
+\r
+  @retval EFI_SUCCESS        The NVDIMM root device SSDT is published.\r
+  @retval Others             The NVDIMM root device SSDT is not published.\r
+\r
+**/\r
+EFI_STATUS\r
+RamDiskPublishSsdt (\r
+  VOID\r
+  )\r
+{\r
+  EFI_STATUS                     Status;\r
+  EFI_ACPI_DESCRIPTION_HEADER    *Table;\r
+  UINTN                          SectionInstance;\r
+  UINTN                          TableSize;\r
+\r
+  Status          = EFI_SUCCESS;\r
+  SectionInstance = 0;\r
+\r
+  //\r
+  // Scan all the EFI raw section instances in FV to find the NVDIMM root\r
+  // device SSDT.\r
+  //\r
+  while (TRUE) {\r
+    Status = GetSectionFromFv (\r
+               &gEfiCallerIdGuid,\r
+               EFI_SECTION_RAW,\r
+               SectionInstance,\r
+               (VOID **) &Table,\r
+               &TableSize\r
+               );\r
+    if (EFI_ERROR (Status)) {\r
+      break;\r
+    }\r
+\r
+    if (Table->OemTableId == SIGNATURE_64 ('R', 'a', 'm', 'D', 'i', 's', 'k', ' ')) {\r
+      Status = mAcpiTableProtocol->InstallAcpiTable (\r
+                                     mAcpiTableProtocol,\r
+                                     Table,\r
+                                     TableSize,\r
+                                     &mRamDiskSsdtTableKey\r
+                                     );\r
+      ASSERT_EFI_ERROR (Status);\r
+\r
+      if (!EFI_ERROR (Status)) {\r
+        mRamDiskSsdtTableKeyValid = TRUE;\r
+      }\r
+\r
+      FreePool (Table);\r
+      return Status;\r
+    } else {\r
+      FreePool (Table);\r
+      SectionInstance++;\r
+    }\r
+  }\r
+\r
+  return Status;\r
+}\r
+\r
+\r
+/**\r
+  Publish the RAM disk NVDIMM Firmware Interface Table (NFIT) to the ACPI\r
+  table.\r
+\r
+  @param[in] PrivateData          Points to RAM disk private data.\r
+\r
+  @retval EFI_SUCCESS             The RAM disk NFIT has been published.\r
+  @retval others                  The RAM disk NFIT has not been published.\r
+\r
+**/\r
+EFI_STATUS\r
+RamDiskPublishNfit (\r
+  IN RAM_DISK_PRIVATE_DATA        *PrivateData\r
+  )\r
+{\r
+  EFI_STATUS                                    Status;\r
+  EFI_MEMORY_DESCRIPTOR                         *MemoryMap;\r
+  EFI_MEMORY_DESCRIPTOR                         *MemoryMapEntry;\r
+  EFI_MEMORY_DESCRIPTOR                         *MemoryMapEnd;\r
+  UINTN                                         TableIndex;\r
+  VOID                                          *TableHeader;\r
+  EFI_ACPI_TABLE_VERSION                        TableVersion;\r
+  UINTN                                         TableKey;\r
+  EFI_ACPI_DESCRIPTION_HEADER                   *NfitHeader;\r
+  EFI_ACPI_6_1_NFIT_SYSTEM_PHYSICAL_ADDRESS_RANGE_STRUCTURE\r
+                                                *SpaRange;\r
+  VOID                                          *Nfit;\r
+  UINT32                                        NfitLen;\r
+  UINTN                                         MemoryMapSize;\r
+  UINTN                                         MapKey;\r
+  UINTN                                         DescriptorSize;\r
+  UINT32                                        DescriptorVersion;\r
+  UINT64                                        CurrentData;\r
+  UINT8                                         Checksum;\r
+  BOOLEAN                                       MemoryFound;\r
+\r
+  //\r
+  // Get the EFI memory map.\r
+  //\r
+  MemoryMapSize = 0;\r
+  MemoryMap     = NULL;\r
+  MemoryFound   = FALSE;\r
+\r
+  Status = gBS->GetMemoryMap (\r
+                  &MemoryMapSize,\r
+                  MemoryMap,\r
+                  &MapKey,\r
+                  &DescriptorSize,\r
+                  &DescriptorVersion\r
+                  );\r
+  ASSERT (Status == EFI_BUFFER_TOO_SMALL);\r
+  do {\r
+    MemoryMap = (EFI_MEMORY_DESCRIPTOR *) AllocatePool (MemoryMapSize);\r
+    ASSERT (MemoryMap != NULL);\r
+    Status = gBS->GetMemoryMap (\r
+                    &MemoryMapSize,\r
+                    MemoryMap,\r
+                    &MapKey,\r
+                    &DescriptorSize,\r
+                    &DescriptorVersion\r
+                    );\r
+    if (EFI_ERROR (Status)) {\r
+      FreePool (MemoryMap);\r
+    }\r
+  } while (Status == EFI_BUFFER_TOO_SMALL);\r
+  ASSERT_EFI_ERROR (Status);\r
+\r
+  MemoryMapEntry = MemoryMap;\r
+  MemoryMapEnd   = (EFI_MEMORY_DESCRIPTOR *) ((UINT8 *) MemoryMap + MemoryMapSize);\r
+  while ((UINTN) MemoryMapEntry < (UINTN) MemoryMapEnd) {\r
+    if ((MemoryMapEntry->Type == EfiReservedMemoryType) &&\r
+        (MemoryMapEntry->PhysicalStart <= PrivateData->StartingAddr) &&\r
+        (MemoryMapEntry->PhysicalStart +\r
+         MultU64x32 (MemoryMapEntry->NumberOfPages, EFI_PAGE_SIZE)\r
+         >= PrivateData->StartingAddr + PrivateData->Size)) {\r
+      MemoryFound = TRUE;\r
+      DEBUG ((\r
+        EFI_D_INFO,\r
+        "RamDiskPublishNfit: RAM disk with reserved meomry type, will publish to NFIT.\n"\r
+        ));\r
+      break;\r
+    }\r
+    MemoryMapEntry = NEXT_MEMORY_DESCRIPTOR (MemoryMapEntry, DescriptorSize);\r
+  }\r
+  FreePool (MemoryMap);\r
+\r
+  if (!MemoryFound) {\r
+    return EFI_NOT_FOUND;\r
+  }\r
+\r
+  //\r
+  // Determine whether there is a NFIT already in the ACPI table.\r
+  //\r
+  Status      = EFI_SUCCESS;\r
+  TableIndex  = 0;\r
+  TableKey    = 0;\r
+  TableHeader = NULL;\r
+\r
+  while (!EFI_ERROR (Status)) {\r
+    Status = mAcpiSdtProtocol->GetAcpiTable (\r
+                                 TableIndex,\r
+                                 (EFI_ACPI_SDT_HEADER **)&TableHeader,\r
+                                 &TableVersion,\r
+                                 &TableKey\r
+                                 );\r
+    if (!EFI_ERROR (Status)) {\r
+      TableIndex++;\r
+\r
+      if (((EFI_ACPI_SDT_HEADER *)TableHeader)->Signature ==\r
+          EFI_ACPI_6_1_NVDIMM_FIRMWARE_INTERFACE_TABLE_STRUCTURE_SIGNATURE) {\r
+        break;\r
+      }\r
+    }\r
+  }\r
+\r
+  if (!EFI_ERROR (Status)) {\r
+    //\r
+    // A NFIT is already in the ACPI table.\r
+    //\r
+    DEBUG ((\r
+      EFI_D_INFO,\r
+      "RamDiskPublishNfit: A NFIT is already exist in the ACPI Table.\n"\r
+      ));\r
+\r
+    NfitHeader = (EFI_ACPI_DESCRIPTION_HEADER *)TableHeader;\r
+    NfitLen    = NfitHeader->Length + sizeof (EFI_ACPI_6_1_NFIT_SYSTEM_PHYSICAL_ADDRESS_RANGE_STRUCTURE);\r
+    Nfit       = AllocateZeroPool (NfitLen);\r
+    if (Nfit == NULL) {\r
+      return EFI_OUT_OF_RESOURCES;\r
+    }\r
+    CopyMem (Nfit, TableHeader, NfitHeader->Length);\r
+\r
+    //\r
+    // Update the NFIT head pointer.\r
+    //\r
+    NfitHeader = (EFI_ACPI_DESCRIPTION_HEADER *)Nfit;\r
+\r
+    //\r
+    // Uninstall the origin NFIT from the ACPI table.\r
+    //\r
+    Status = mAcpiTableProtocol->UninstallAcpiTable (\r
+                                   mAcpiTableProtocol,\r
+                                   TableKey\r
+                                   );\r
+    ASSERT_EFI_ERROR (Status);\r
+\r
+    if (EFI_ERROR (Status)) {\r
+      FreePool (Nfit);\r
+      return Status;\r
+    }\r
+\r
+    //\r
+    // Append the System Physical Address (SPA) Range Structure at the end\r
+    // of the origin NFIT.\r
+    //\r
+    SpaRange   = (EFI_ACPI_6_1_NFIT_SYSTEM_PHYSICAL_ADDRESS_RANGE_STRUCTURE *)\r
+                 ((UINT8 *)Nfit + NfitHeader->Length);\r
+\r
+    //\r
+    // Update the length field of the NFIT\r
+    //\r
+    NfitHeader->Length   = NfitLen;\r
+\r
+    //\r
+    // The checksum will be updated after the new contents are appended.\r
+    //\r
+    NfitHeader->Checksum = 0;\r
+  } else {\r
+    //\r
+    // Assumption is made that if no NFIT is in the ACPI table, there is no\r
+    // NVDIMM root device in the \SB scope.\r
+    // Therefore, a NVDIMM root device will be reported via Secondary System\r
+    // Description Table (SSDT).\r
+    //\r
+    Status = RamDiskPublishSsdt ();\r
+    if (EFI_ERROR (Status)) {\r
+      return Status;\r
+    }\r
+\r
+    //\r
+    // No NFIT is in the ACPI table, we will create one here.\r
+    //\r
+    DEBUG ((\r
+      EFI_D_INFO,\r
+      "RamDiskPublishNfit: No NFIT is in the ACPI Table, will create one.\n"\r
+      ));\r
+\r
+    NfitLen = sizeof (EFI_ACPI_6_1_NVDIMM_FIRMWARE_INTERFACE_TABLE) +\r
+              sizeof (EFI_ACPI_6_1_NFIT_SYSTEM_PHYSICAL_ADDRESS_RANGE_STRUCTURE);\r
+    Nfit    = AllocateZeroPool (NfitLen);\r
+    if (Nfit == NULL) {\r
+      return EFI_OUT_OF_RESOURCES;\r
+    }\r
+\r
+    SpaRange = (EFI_ACPI_6_1_NFIT_SYSTEM_PHYSICAL_ADDRESS_RANGE_STRUCTURE *)\r
+               ((UINT8 *)Nfit + sizeof (EFI_ACPI_6_1_NVDIMM_FIRMWARE_INTERFACE_TABLE));\r
+\r
+    NfitHeader                  = (EFI_ACPI_DESCRIPTION_HEADER *)Nfit;\r
+    NfitHeader->Signature       = EFI_ACPI_6_1_NVDIMM_FIRMWARE_INTERFACE_TABLE_STRUCTURE_SIGNATURE;\r
+    NfitHeader->Length          = NfitLen;\r
+    NfitHeader->Revision        = EFI_ACPI_6_1_NVDIMM_FIRMWARE_INTERFACE_TABLE_REVISION;\r
+    NfitHeader->Checksum        = 0;\r
+    NfitHeader->OemRevision     = PcdGet32 (PcdAcpiDefaultOemRevision);\r
+    NfitHeader->CreatorId       = PcdGet32 (PcdAcpiDefaultCreatorId);\r
+    NfitHeader->CreatorRevision = PcdGet32 (PcdAcpiDefaultCreatorRevision);\r
+    CurrentData                 = PcdGet64 (PcdAcpiDefaultOemTableId);\r
+    CopyMem (NfitHeader->OemId, PcdGetPtr (PcdAcpiDefaultOemId), sizeof (NfitHeader->OemId));\r
+    CopyMem (&NfitHeader->OemTableId, &CurrentData, sizeof (UINT64));\r
+  }\r
+\r
+  //\r
+  // Fill in the content of the SPA Range Structure.\r
+  //\r
+  SpaRange->Type   = EFI_ACPI_6_1_NFIT_SYSTEM_PHYSICAL_ADDRESS_RANGE_STRUCTURE_TYPE;\r
+  SpaRange->Length = sizeof (EFI_ACPI_6_1_NFIT_SYSTEM_PHYSICAL_ADDRESS_RANGE_STRUCTURE);\r
+  SpaRange->SystemPhysicalAddressRangeBase   = PrivateData->StartingAddr;\r
+  SpaRange->SystemPhysicalAddressRangeLength = PrivateData->Size;\r
+  CopyGuid (&SpaRange->AddressRangeTypeGUID, &PrivateData->TypeGuid);\r
+\r
+  Checksum             = CalculateCheckSum8((UINT8 *)Nfit, NfitHeader->Length);\r
+  NfitHeader->Checksum = Checksum;\r
+\r
+  //\r
+  // Publish the NFIT to the ACPI table.\r
+  // Note, since the NFIT might be modified by other driver, therefore, we\r
+  // do not track the returning TableKey from the InstallAcpiTable().\r
+  //\r
+  Status = mAcpiTableProtocol->InstallAcpiTable (\r
+                                 mAcpiTableProtocol,\r
+                                 Nfit,\r
+                                 NfitHeader->Length,\r
+                                 &TableKey\r
+                                 );\r
+  ASSERT_EFI_ERROR (Status);\r
+\r
+  FreePool (Nfit);\r
+\r
+  if (EFI_ERROR (Status)) {\r
+    return Status;\r
+  }\r
+\r
+  PrivateData->InNfit = TRUE;\r
+\r
+  return EFI_SUCCESS;\r
+}\r
+\r
+\r
+/**\r
+  Unpublish the RAM disk NVDIMM Firmware Interface Table (NFIT) from the\r
+  ACPI table.\r
+\r
+  @param[in] PrivateData          Points to RAM disk private data.\r
+\r
+  @retval EFI_SUCCESS             The RAM disk NFIT has been unpublished.\r
+  @retval others                  The RAM disk NFIT has not been unpublished.\r
+\r
+**/\r
+EFI_STATUS\r
+RamDiskUnpublishNfit (\r
+  IN RAM_DISK_PRIVATE_DATA        *PrivateData\r
+  )\r
+{\r
+  EFI_STATUS                                    Status;\r
+  UINTN                                         TableIndex;\r
+  VOID                                          *TableHeader;\r
+  EFI_ACPI_TABLE_VERSION                        TableVersion;\r
+  UINTN                                         TableKey;\r
+  EFI_ACPI_DESCRIPTION_HEADER                   *NewNfitHeader;\r
+  EFI_ACPI_6_1_NFIT_SYSTEM_PHYSICAL_ADDRESS_RANGE_STRUCTURE\r
+                                                *SpaRange;\r
+  VOID                                          *NewNfit;\r
+  VOID                                          *NewNfitPtr;\r
+  EFI_ACPI_6_1_NFIT_STRUCTURE_HEADER            *NfitStructHeader;\r
+  UINT32                                        NewNfitLen;\r
+  UINT32                                        RemainLen;\r
+  UINT8                                         Checksum;\r
+\r
+  //\r
+  // Find the NFIT in the ACPI table.\r
+  //\r
+  Status      = EFI_SUCCESS;\r
+  TableIndex  = 0;\r
+  TableKey    = 0;\r
+  TableHeader = NULL;\r
+\r
+  while (!EFI_ERROR (Status)) {\r
+    Status = mAcpiSdtProtocol->GetAcpiTable (\r
+                                 TableIndex,\r
+                                 (EFI_ACPI_SDT_HEADER **)&TableHeader,\r
+                                 &TableVersion,\r
+                                 &TableKey\r
+                                 );\r
+    if (!EFI_ERROR (Status)) {\r
+      TableIndex++;\r
+\r
+      if (((EFI_ACPI_SDT_HEADER *)TableHeader)->Signature ==\r
+          EFI_ACPI_6_1_NVDIMM_FIRMWARE_INTERFACE_TABLE_STRUCTURE_SIGNATURE) {\r
+        break;\r
+      }\r
+    }\r
+  }\r
+\r
+  if (EFI_ERROR (Status)) {\r
+    //\r
+    // No NFIT is found in the ACPI table.\r
+    //\r
+    return EFI_NOT_FOUND;\r
+  }\r
+\r
+  NewNfitLen    = ((EFI_ACPI_DESCRIPTION_HEADER *)TableHeader)->Length -\r
+                  sizeof (EFI_ACPI_6_1_NFIT_SYSTEM_PHYSICAL_ADDRESS_RANGE_STRUCTURE);\r
+\r
+  //\r
+  // After removing this RAM disk from the NFIT, if no other structure is in\r
+  // the NFIT, we just remove the NFIT and the SSDT which is used to report\r
+  // the NVDIMM root device.\r
+  //\r
+  if (NewNfitLen == sizeof (EFI_ACPI_6_1_NVDIMM_FIRMWARE_INTERFACE_TABLE)) {\r
+    //\r
+    // Remove the NFIT.\r
+    //\r
+    Status = mAcpiTableProtocol->UninstallAcpiTable (\r
+                                   mAcpiTableProtocol,\r
+                                   TableKey\r
+                                   );\r
+    ASSERT_EFI_ERROR (Status);\r
+    if (EFI_ERROR (Status)) {\r
+      return Status;\r
+    }\r
+\r
+    //\r
+    // Remove the SSDT which is used by RamDiskDxe driver to report the NVDIMM\r
+    // root device.\r
+    // We do not care the return status since this SSDT might already be\r
+    // uninstalled by other drivers to update the information of the NVDIMM\r
+    // root device.\r
+    //\r
+    if (mRamDiskSsdtTableKeyValid) {\r
+      mRamDiskSsdtTableKeyValid = FALSE;\r
+\r
+      mAcpiTableProtocol->UninstallAcpiTable (\r
+                            mAcpiTableProtocol,\r
+                            mRamDiskSsdtTableKey\r
+                            );\r
+    }\r
+\r
+    return EFI_SUCCESS;\r
+  }\r
+\r
+  NewNfit = AllocateZeroPool (NewNfitLen);\r
+  if (NewNfit == NULL) {\r
+    return EFI_OUT_OF_RESOURCES;\r
+  }\r
+\r
+  //\r
+  // Get a copy of the old NFIT header content.\r
+  //\r
+  CopyMem (NewNfit, TableHeader, sizeof (EFI_ACPI_6_1_NVDIMM_FIRMWARE_INTERFACE_TABLE));\r
+  NewNfitHeader           = (EFI_ACPI_DESCRIPTION_HEADER *)NewNfit;\r
+  NewNfitHeader->Length   = NewNfitLen;\r
+  NewNfitHeader->Checksum = 0;\r
+\r
+  //\r
+  // Copy the content of required NFIT structures.\r
+  //\r
+  NewNfitPtr       = (UINT8 *)NewNfit + sizeof (EFI_ACPI_6_1_NVDIMM_FIRMWARE_INTERFACE_TABLE);\r
+  RemainLen        = NewNfitLen - sizeof (EFI_ACPI_6_1_NVDIMM_FIRMWARE_INTERFACE_TABLE);\r
+  NfitStructHeader = (EFI_ACPI_6_1_NFIT_STRUCTURE_HEADER *)\r
+                     ((UINT8 *)TableHeader + sizeof (EFI_ACPI_6_1_NVDIMM_FIRMWARE_INTERFACE_TABLE));\r
+  while (RemainLen > 0) {\r
+    if ((NfitStructHeader->Type == EFI_ACPI_6_1_NFIT_SYSTEM_PHYSICAL_ADDRESS_RANGE_STRUCTURE_TYPE) &&\r
+        (NfitStructHeader->Length == sizeof (EFI_ACPI_6_1_NFIT_SYSTEM_PHYSICAL_ADDRESS_RANGE_STRUCTURE))) {\r
+      SpaRange = (EFI_ACPI_6_1_NFIT_SYSTEM_PHYSICAL_ADDRESS_RANGE_STRUCTURE *)NfitStructHeader;\r
+\r
+      if ((SpaRange->SystemPhysicalAddressRangeBase == PrivateData->StartingAddr) &&\r
+          (SpaRange->SystemPhysicalAddressRangeLength == PrivateData->Size) &&\r
+          (CompareGuid (&SpaRange->AddressRangeTypeGUID, &PrivateData->TypeGuid))) {\r
+        //\r
+        // Skip the SPA Range Structure for the RAM disk to be unpublished\r
+        // from NFIT.\r
+        //\r
+        NfitStructHeader = (EFI_ACPI_6_1_NFIT_STRUCTURE_HEADER *)\r
+                           ((UINT8 *)NfitStructHeader + NfitStructHeader->Length);\r
+        continue;\r
+      }\r
+    }\r
+\r
+    //\r
+    // Copy the content of origin NFIT.\r
+    //\r
+    CopyMem (NewNfitPtr, NfitStructHeader, NfitStructHeader->Length);\r
+    NewNfitPtr = (UINT8 *)NewNfitPtr + NfitStructHeader->Length;\r
+\r
+    //\r
+    // Move to the header of next NFIT structure.\r
+    //\r
+    RemainLen       -= NfitStructHeader->Length;\r
+    NfitStructHeader = (EFI_ACPI_6_1_NFIT_STRUCTURE_HEADER *)\r
+                       ((UINT8 *)NfitStructHeader + NfitStructHeader->Length);\r
+  }\r
+\r
+  Checksum                = CalculateCheckSum8((UINT8 *)NewNfit, NewNfitHeader->Length);\r
+  NewNfitHeader->Checksum = Checksum;\r
+\r
+  Status = mAcpiTableProtocol->UninstallAcpiTable (\r
+                                 mAcpiTableProtocol,\r
+                                 TableKey\r
+                                 );\r
+  ASSERT_EFI_ERROR (Status);\r
+\r
+  if (EFI_ERROR (Status)) {\r
+    FreePool (NewNfit);\r
+    return Status;\r
+  }\r
+\r
+  //\r
+  // Publish the NFIT to the ACPI table.\r
+  // Note, since the NFIT might be modified by other driver, therefore, we\r
+  // do not track the returning TableKey from the InstallAcpiTable().\r
+  //\r
+  Status = mAcpiTableProtocol->InstallAcpiTable (\r
+                                 mAcpiTableProtocol,\r
+                                 NewNfit,\r
+                                 NewNfitLen,\r
+                                 &TableKey\r
+                                 );\r
+  ASSERT_EFI_ERROR (Status);\r
+\r
+  FreePool (NewNfit);\r
+  if (EFI_ERROR (Status)) {\r
+    return Status;\r
+  }\r
+\r
+  return EFI_SUCCESS;\r
+}\r
+\r
+\r
 /**\r
   Register a RAM disk with specified address, size and type.\r
 \r
@@ -112,7 +613,8 @@ RamDiskRegister (
   //\r
   // Add check to prevent data read across the memory boundary\r
   //\r
-  if (RamDiskBase + RamDiskSize > ((UINTN) -1) - RAM_DISK_BLOCK_SIZE + 1) {\r
+  if ((RamDiskSize > MAX_UINTN) ||\r
+      (RamDiskBase > MAX_UINTN - RamDiskSize + 1)) {\r
     return EFI_INVALID_PARAMETER;\r
   }\r
 \r
@@ -211,12 +713,15 @@ RamDiskRegister (
   // Insert the newly created one to the registered RAM disk list\r
   //\r
   InsertTailList (&RegisteredRamDisks, &PrivateData->ThisInstance);\r
-  ListEntryNum++;\r
 \r
   gBS->ConnectController (PrivateData->Handle, NULL, NULL, TRUE);\r
 \r
   FreePool (RamDiskDevNode);\r
 \r
+  if ((mAcpiTableProtocol != NULL) && (mAcpiSdtProtocol != NULL)) {\r
+    RamDiskPublishNfit (PrivateData);\r
+  }\r
+\r
   return EFI_SUCCESS;\r
 \r
 ErrorExit:\r
@@ -308,6 +813,13 @@ RamDiskUnregister (
       if ((StartingAddr == PrivateData->StartingAddr) &&\r
           (EndingAddr == PrivateData->StartingAddr + PrivateData->Size - 1) &&\r
           (CompareGuid (&RamDiskDevNode->TypeGuid, &PrivateData->TypeGuid))) {\r
+        //\r
+        // Remove the content for this RAM disk in NFIT.\r
+        //\r
+        if (PrivateData->InNfit) {\r
+          RamDiskUnpublishNfit (PrivateData);\r
+        }\r
+\r
         //\r
         // Uninstall the EFI_DEVICE_PATH_PROTOCOL & EFI_BLOCK_IO(2)_PROTOCOL\r
         //\r
@@ -335,7 +847,6 @@ RamDiskUnregister (
 \r
         FreePool (PrivateData->DevicePath);\r
         FreePool (PrivateData);\r
-        ListEntryNum--;\r
         Found = TRUE;\r
 \r
         break;\r