MdeModulePkg/Universal/CapsulePei/Common/CapsuleCoalesce.c

   1 /** @file
   2   The logic to process capsule.
   3
   4   Caution: This module requires additional review when modified.
   5   This driver will have external input - capsule image.
   6   This external input must be validated carefully to avoid security issue like
   7   buffer overflow, integer overflow.
   8
   9   CapsuleDataCoalesce() will do basic validation before coalesce capsule data
  10   into memory.
  11
  12 (C) Copyright 2014 Hewlett-Packard Development Company, L.P.<BR>
  13 Copyright (c) 2011 - 2018, Intel Corporation. All rights reserved.<BR>
  14 SPDX-License-Identifier: BSD-2-Clause-Patent
  15
  16 **/
  17
  18 #include <Uefi.h>
  19 #include <PiPei.h>
  20
  21 #include <Guid/CapsuleVendor.h>
  22
  23 #include <Library/BaseMemoryLib.h>
  24 #include <Library/DebugLib.h>
  25 #include <Library/PrintLib.h>
  26 #include <Library/BaseLib.h>
  27
  28 #include "CommonHeader.h"
  29
  30 #define MIN_COALESCE_ADDR  (1024 * 1024)
  31
  32 /**
  33   Given a pointer to the capsule block list, info on the available system
  34   memory, and the size of a buffer, find a free block of memory where a
  35   buffer of the given size can be copied to safely.
  36
  37   @param BlockList   Pointer to head of capsule block descriptors
  38   @param MemBase     Pointer to the base of memory in which we want to find free space
  39   @param MemSize     The size of the block of memory pointed to by MemBase
  40   @param DataSize    How big a free block we want to find
  41
  42   @return A pointer to a memory block of at least DataSize that lies somewhere
  43           between MemBase and (MemBase + MemSize). The memory pointed to does not
  44           contain any of the capsule block descriptors or capsule blocks pointed to
  45           by the BlockList.
  46
  47 **/
  48 UINT8 *
  49 FindFreeMem (
  50   EFI_CAPSULE_BLOCK_DESCRIPTOR  *BlockList,
  51   UINT8                         *MemBase,
  52   UINTN                         MemSize,
  53   UINTN                         DataSize
  54   );
  55
  56 /**
  57   The capsule block descriptors may be fragmented and spread all over memory.
  58   To simplify the coalescing of capsule blocks, first coalesce all the
  59   capsule block descriptors low in memory.
  60
  61   The descriptors passed in can be fragmented throughout memory. Here
  62   they are relocated into memory to turn them into a contiguous (null
  63   terminated) array.
  64
  65   @param PeiServices    pointer to PEI services table
  66   @param BlockList      pointer to the capsule block descriptors
  67   @param NumDescriptors number of capsule data block descriptors, whose Length is non-zero.
  68   @param MemBase        base of system memory in which we can work
  69   @param MemSize        size of the system memory pointed to by MemBase
  70
  71   @retval NULL    could not relocate the descriptors
  72   @retval Pointer to the base of the successfully-relocated block descriptors.
  73
  74 **/
  75 EFI_CAPSULE_BLOCK_DESCRIPTOR *
  76 RelocateBlockDescriptors (
  77   IN EFI_PEI_SERVICES              **PeiServices,
  78   IN EFI_CAPSULE_BLOCK_DESCRIPTOR  *BlockList,
  79   IN UINTN                         NumDescriptors,
  80   IN UINT8                         *MemBase,
  81   IN UINTN                         MemSize
  82   );
  83
  84 /**
  85   Check every capsule header.
  86
  87   @param CapsuleHeader   The pointer to EFI_CAPSULE_HEADER
  88
  89   @retval FALSE  Capsule is OK
  90   @retval TRUE   Capsule is corrupted
  91
  92 **/
  93 BOOLEAN
  94 IsCapsuleCorrupted (
  95   IN EFI_CAPSULE_HEADER  *CapsuleHeader
  96   );
  97
  98 /**
  99   Determine if two buffers overlap in memory.
 100
 101   @param Buff1   pointer to first buffer
 102   @param Size1   size of Buff1
 103   @param Buff2   pointer to second buffer
 104   @param Size2   size of Buff2
 105
 106   @retval TRUE    Buffers overlap in memory.
 107   @retval FALSE   Buffer doesn't overlap.
 108
 109 **/
 110 BOOLEAN
 111 IsOverlapped (
 112   UINT8  *Buff1,
 113   UINTN  Size1,
 114   UINT8  *Buff2,
 115   UINTN  Size2
 116   );
 117
 118 /**
 119   Given a pointer to a capsule block descriptor, traverse the list to figure
 120   out how many legitimate descriptors there are, and how big the capsule it
 121   refers to is.
 122
 123   @param Desc            Pointer to the capsule block descriptors
 124   @param NumDescriptors  Optional pointer to where to return the number of capsule data descriptors, whose Length is non-zero.
 125   @param CapsuleSize     Optional pointer to where to return the capsule image size
 126   @param CapsuleNumber   Optional pointer to where to return the number of capsule
 127
 128   @retval EFI_NOT_FOUND   No descriptors containing data in the list
 129   @retval EFI_SUCCESS     Return data is valid
 130
 131 **/
 132 EFI_STATUS
 133 GetCapsuleInfo (
 134   IN EFI_CAPSULE_BLOCK_DESCRIPTOR  *Desc,
 135   IN OUT UINTN                     *NumDescriptors OPTIONAL,
 136   IN OUT UINTN                     *CapsuleSize OPTIONAL,
 137   IN OUT UINTN                     *CapsuleNumber OPTIONAL
 138   );
 139
 140 /**
 141   Given a pointer to the capsule block list, info on the available system
 142   memory, and the size of a buffer, find a free block of memory where a
 143   buffer of the given size can be copied to safely.
 144
 145   @param BlockList   Pointer to head of capsule block descriptors
 146   @param MemBase     Pointer to the base of memory in which we want to find free space
 147   @param MemSize     The size of the block of memory pointed to by MemBase
 148   @param DataSize    How big a free block we want to find
 149
 150   @return A pointer to a memory block of at least DataSize that lies somewhere
 151           between MemBase and (MemBase + MemSize). The memory pointed to does not
 152           contain any of the capsule block descriptors or capsule blocks pointed to
 153           by the BlockList.
 154
 155 **/
 156 UINT8 *
 157 FindFreeMem (
 158   EFI_CAPSULE_BLOCK_DESCRIPTOR  *BlockList,
 159   UINT8                         *MemBase,
 160   UINTN                         MemSize,
 161   UINTN                         DataSize
 162   )
 163 {
 164   UINTN                         Size;
 165   EFI_CAPSULE_BLOCK_DESCRIPTOR  *CurrDesc;
 166   EFI_CAPSULE_BLOCK_DESCRIPTOR  *TempDesc;
 167   UINT8                         *MemEnd;
 168   BOOLEAN                       Failed;
 169
 170   //
 171   // Need at least enough to copy the data to at the end of the buffer, so
 172   // say the end is less the data size for easy comparisons here.
 173   //
 174   MemEnd   = MemBase + MemSize - DataSize;
 175   CurrDesc = BlockList;
 176   //
 177   // Go through all the descriptor blocks and see if any obstruct the range
 178   //
 179   while (CurrDesc != NULL) {
 180     //
 181     // Get the size of this block list and see if it's in the way
 182     //
 183     Failed   = FALSE;
 184     TempDesc = CurrDesc;
 185     Size     = sizeof (EFI_CAPSULE_BLOCK_DESCRIPTOR);
 186     while (TempDesc->Length != 0) {
 187       Size += sizeof (EFI_CAPSULE_BLOCK_DESCRIPTOR);
 188       TempDesc++;
 189     }
 190
 191     if (IsOverlapped (MemBase, DataSize, (UINT8 *)CurrDesc, Size)) {
 192       //
 193       // Set our new base to the end of this block list and start all over
 194       //
 195       MemBase  = (UINT8 *)CurrDesc + Size;
 196       CurrDesc = BlockList;
 197       if (MemBase > MemEnd) {
 198         return NULL;
 199       }
 200
 201       Failed = TRUE;
 202     }
 203
 204     //
 205     // Now go through all the blocks and make sure none are in the way
 206     //
 207     while ((CurrDesc->Length != 0) && (!Failed)) {
 208       if (IsOverlapped (MemBase, DataSize, (UINT8 *)(UINTN)CurrDesc->Union.DataBlock, (UINTN)CurrDesc->Length)) {
 209         //
 210         // Set our new base to the end of this block and start all over
 211         //
 212         Failed   = TRUE;
 213         MemBase  = (UINT8 *)((UINTN)CurrDesc->Union.DataBlock) + CurrDesc->Length;
 214         CurrDesc = BlockList;
 215         if (MemBase > MemEnd) {
 216           return NULL;
 217         }
 218       }
 219
 220       CurrDesc++;
 221     }
 222
 223     //
 224     // Normal continuation -- jump to next block descriptor list
 225     //
 226     if (!Failed) {
 227       CurrDesc = (EFI_CAPSULE_BLOCK_DESCRIPTOR  *)(UINTN)CurrDesc->Union.ContinuationPointer;
 228     }
 229   }
 230
 231   return MemBase;
 232 }
 233
 234 /**
 235   Validate capsule by MemoryResource.
 236
 237   @param MemoryResource  Pointer to the buffer of memory resource descriptor.
 238   @param Address         Address to be validated.
 239   @param Size            Size to be validated.
 240
 241   @retval TRUE  No memory resource descriptor reported in HOB list before capsule Coalesce,
 242                 or it is valid in one MemoryResource.
 243           FALSE It is not in any MemoryResource.
 244
 245 **/
 246 BOOLEAN
 247 ValidateCapsuleByMemoryResource (
 248   IN MEMORY_RESOURCE_DESCRIPTOR  *MemoryResource,
 249   IN EFI_PHYSICAL_ADDRESS        Address,
 250   IN UINT64                      Size
 251   )
 252 {
 253   UINTN  Index;
 254
 255   //
 256   // Sanity Check
 257   //
 258   if (Size > MAX_ADDRESS) {
 259     DEBUG ((DEBUG_ERROR, "ERROR: Size(0x%lx) > MAX_ADDRESS\n", Size));
 260     return FALSE;
 261   }
 262
 263   //
 264   // Sanity Check
 265   //
 266   if (Address > (MAX_ADDRESS - Size)) {
 267     DEBUG ((DEBUG_ERROR, "ERROR: Address(0x%lx) > (MAX_ADDRESS - Size(0x%lx))\n", Address, Size));
 268     return FALSE;
 269   }
 270
 271   if (MemoryResource == NULL) {
 272     //
 273     // No memory resource descriptor reported in HOB list before capsule Coalesce.
 274     //
 275     return TRUE;
 276   }
 277
 278   for (Index = 0; MemoryResource[Index].ResourceLength != 0; Index++) {
 279     if ((Address >= MemoryResource[Index].PhysicalStart) &&
 280         ((Address + Size) <= (MemoryResource[Index].PhysicalStart + MemoryResource[Index].ResourceLength)))
 281     {
 282       DEBUG ((
 283         DEBUG_INFO,
 284         "Address(0x%lx) Size(0x%lx) in MemoryResource[0x%x] - Start(0x%lx) Length(0x%lx)\n",
 285         Address,
 286         Size,
 287         Index,
 288         MemoryResource[Index].PhysicalStart,
 289         MemoryResource[Index].ResourceLength
 290         ));
 291       return TRUE;
 292     }
 293   }
 294
 295   DEBUG ((DEBUG_ERROR, "ERROR: Address(0x%lx) Size(0x%lx) not in any MemoryResource\n", Address, Size));
 296   return FALSE;
 297 }
 298
 299 /**
 300   Check the integrity of the capsule descriptors.
 301
 302   @param BlockList       Pointer to the capsule descriptors
 303   @param MemoryResource  Pointer to the buffer of memory resource descriptor.
 304
 305   @retval NULL           BlockList is not valid.
 306   @retval LastBlockDesc  Last one Block in BlockList
 307
 308 **/
 309 EFI_CAPSULE_BLOCK_DESCRIPTOR *
 310 ValidateCapsuleIntegrity (
 311   IN EFI_CAPSULE_BLOCK_DESCRIPTOR  *BlockList,
 312   IN MEMORY_RESOURCE_DESCRIPTOR    *MemoryResource
 313   )
 314 {
 315   EFI_CAPSULE_HEADER            *CapsuleHeader;
 316   UINT64                        CapsuleSize;
 317   UINTN                         CapsuleCount;
 318   EFI_CAPSULE_BLOCK_DESCRIPTOR  *Ptr;
 319
 320   DEBUG ((DEBUG_INFO, "ValidateCapsuleIntegrity\n"));
 321
 322   //
 323   // Go through the list to look for inconsistencies. Check for:
 324   //   * misaligned block descriptors.
 325   //   * The first capsule header guid
 326   //   * The first capsule header flag
 327   //   * The first capsule header HeaderSize
 328   //   * Below check will be done in ValidateCapsuleByMemoryResource()
 329   //     Length > MAX_ADDRESS
 330   //     Ptr + sizeof (EFI_CAPSULE_BLOCK_DESCRIPTOR) > MAX_ADDRESS
 331   //     DataBlock + Length > MAX_ADDRESS
 332   //
 333   CapsuleSize  = 0;
 334   CapsuleCount = 0;
 335   Ptr          = BlockList;
 336
 337   if (!ValidateCapsuleByMemoryResource (MemoryResource, (EFI_PHYSICAL_ADDRESS)(UINTN)Ptr, sizeof (EFI_CAPSULE_BLOCK_DESCRIPTOR))) {
 338     return NULL;
 339   }
 340
 341   DEBUG ((DEBUG_INFO, "Ptr - 0x%p\n", Ptr));
 342   DEBUG ((DEBUG_INFO, "Ptr->Length - 0x%lx\n", Ptr->Length));
 343   DEBUG ((DEBUG_INFO, "Ptr->Union - 0x%lx\n", Ptr->Union.ContinuationPointer));
 344   while ((Ptr->Length != 0) || (Ptr->Union.ContinuationPointer != (EFI_PHYSICAL_ADDRESS)(UINTN)NULL)) {
 345     //
 346     // Make sure the descriptor is aligned at UINT64 in memory
 347     //
 348     if ((UINTN)Ptr & (sizeof (UINT64) - 1)) {
 349       DEBUG ((DEBUG_ERROR, "ERROR: BlockList address failed alignment check\n"));
 350       return NULL;
 351     }
 352
 353     if (Ptr->Length == 0) {
 354       //
 355       // Descriptor points to another list of block descriptors somewhere
 356       // else.
 357       //
 358       Ptr = (EFI_CAPSULE_BLOCK_DESCRIPTOR  *)(UINTN)Ptr->Union.ContinuationPointer;
 359       if (!ValidateCapsuleByMemoryResource (MemoryResource, (EFI_PHYSICAL_ADDRESS)(UINTN)Ptr, sizeof (EFI_CAPSULE_BLOCK_DESCRIPTOR))) {
 360         return NULL;
 361       }
 362
 363       DEBUG ((DEBUG_INFO, "Ptr(C) - 0x%p\n", Ptr));
 364       DEBUG ((DEBUG_INFO, "Ptr->Length - 0x%lx\n", Ptr->Length));
 365       DEBUG ((DEBUG_INFO, "Ptr->Union - 0x%lx\n", Ptr->Union.ContinuationPointer));
 366     } else {
 367       if (!ValidateCapsuleByMemoryResource (MemoryResource, Ptr->Union.DataBlock, Ptr->Length)) {
 368         return NULL;
 369       }
 370
 371       //
 372       // To enhance the reliability of check-up, the first capsule's header is checked here.
 373       // More reliabilities check-up will do later.
 374       //
 375       if (CapsuleSize == 0) {
 376         //
 377         // Move to the first capsule to check its header.
 378         //
 379         CapsuleHeader = (EFI_CAPSULE_HEADER *)((UINTN)Ptr->Union.DataBlock);
 380         //
 381         // Sanity check
 382         //
 383         if (Ptr->Length < sizeof (EFI_CAPSULE_HEADER)) {
 384           DEBUG ((DEBUG_ERROR, "ERROR: Ptr->Length(0x%lx) < sizeof(EFI_CAPSULE_HEADER)\n", Ptr->Length));
 385           return NULL;
 386         }
 387
 388         //
 389         // Make sure HeaderSize field is valid
 390         //
 391         if (CapsuleHeader->HeaderSize > CapsuleHeader->CapsuleImageSize) {
 392           DEBUG ((DEBUG_ERROR, "ERROR: CapsuleHeader->HeaderSize(0x%x) > CapsuleHeader->CapsuleImageSize(0x%x)\n", CapsuleHeader->HeaderSize, CapsuleHeader->CapsuleImageSize));
 393           return NULL;
 394         }
 395
 396         if (IsCapsuleCorrupted (CapsuleHeader)) {
 397           return NULL;
 398         }
 399
 400         CapsuleCount++;
 401         CapsuleSize = CapsuleHeader->CapsuleImageSize;
 402       }
 403
 404       if (CapsuleSize >= Ptr->Length) {
 405         CapsuleSize = CapsuleSize - Ptr->Length;
 406       } else {
 407         DEBUG ((DEBUG_ERROR, "ERROR: CapsuleSize(0x%lx) < Ptr->Length(0x%lx)\n", CapsuleSize, Ptr->Length));
 408         //
 409         // Sanity check
 410         //
 411         return NULL;
 412       }
 413
 414       //
 415       // Move to next BLOCK descriptor
 416       //
 417       Ptr++;
 418       if (!ValidateCapsuleByMemoryResource (MemoryResource, (EFI_PHYSICAL_ADDRESS)(UINTN)Ptr, sizeof (EFI_CAPSULE_BLOCK_DESCRIPTOR))) {
 419         return NULL;
 420       }
 421
 422       DEBUG ((DEBUG_INFO, "Ptr(B) - 0x%p\n", Ptr));
 423       DEBUG ((DEBUG_INFO, "Ptr->Length - 0x%lx\n", Ptr->Length));
 424       DEBUG ((DEBUG_INFO, "Ptr->Union - 0x%lx\n", Ptr->Union.ContinuationPointer));
 425     }
 426   }
 427
 428   if (CapsuleCount == 0) {
 429     //
 430     // No any capsule is found in BlockList
 431     //
 432     DEBUG ((DEBUG_ERROR, "ERROR: CapsuleCount(0x%x) == 0\n", CapsuleCount));
 433     return NULL;
 434   }
 435
 436   if (CapsuleSize != 0) {
 437     //
 438     // Capsule data is incomplete.
 439     //
 440     DEBUG ((DEBUG_ERROR, "ERROR: CapsuleSize(0x%lx) != 0\n", CapsuleSize));
 441     return NULL;
 442   }
 443
 444   return Ptr;
 445 }
 446
 447 /**
 448   The capsule block descriptors may be fragmented and spread all over memory.
 449   To simplify the coalescing of capsule blocks, first coalesce all the
 450   capsule block descriptors low in memory.
 451
 452   The descriptors passed in can be fragmented throughout memory. Here
 453   they are relocated into memory to turn them into a contiguous (null
 454   terminated) array.
 455
 456   @param PeiServices    pointer to PEI services table
 457   @param BlockList      pointer to the capsule block descriptors
 458   @param NumDescriptors number of capsule data block descriptors, whose Length is non-zero.
 459   @param MemBase        base of system memory in which we can work
 460   @param MemSize        size of the system memory pointed to by MemBase
 461
 462   @retval NULL    could not relocate the descriptors
 463   @retval Pointer to the base of the successfully-relocated block descriptors.
 464
 465 **/
 466 EFI_CAPSULE_BLOCK_DESCRIPTOR  *
 467 RelocateBlockDescriptors (
 468   IN EFI_PEI_SERVICES              **PeiServices,
 469   IN EFI_CAPSULE_BLOCK_DESCRIPTOR  *BlockList,
 470   IN UINTN                         NumDescriptors,
 471   IN UINT8                         *MemBase,
 472   IN UINTN                         MemSize
 473   )
 474 {
 475   EFI_CAPSULE_BLOCK_DESCRIPTOR  *NewBlockList;
 476   EFI_CAPSULE_BLOCK_DESCRIPTOR  *CurrBlockDescHead;
 477   EFI_CAPSULE_BLOCK_DESCRIPTOR  *TempBlockDesc;
 478   EFI_CAPSULE_BLOCK_DESCRIPTOR  *PrevBlockDescTail;
 479   UINTN                         BufferSize;
 480   UINT8                         *RelocBuffer;
 481   UINTN                         BlockListSize;
 482
 483   //
 484   // Get the info on the blocks and descriptors. Since we're going to move
 485   // the descriptors low in memory, adjust the base/size values accordingly here.
 486   // NumDescriptors is the number of legit data descriptors, so add one for
 487   // a terminator. (Already done by caller, no check is needed.)
 488   //
 489
 490   BufferSize   = NumDescriptors * sizeof (EFI_CAPSULE_BLOCK_DESCRIPTOR);
 491   NewBlockList = (EFI_CAPSULE_BLOCK_DESCRIPTOR *)MemBase;
 492   if (MemSize < BufferSize) {
 493     return NULL;
 494   }
 495
 496   MemSize -= BufferSize;
 497   MemBase += BufferSize;
 498   //
 499   // Go through all the blocks and make sure none are in the way
 500   //
 501   TempBlockDesc = BlockList;
 502   while (TempBlockDesc->Union.ContinuationPointer != (EFI_PHYSICAL_ADDRESS)(UINTN)NULL) {
 503     if (TempBlockDesc->Length == 0) {
 504       //
 505       // Next block of descriptors
 506       //
 507       TempBlockDesc = (EFI_CAPSULE_BLOCK_DESCRIPTOR  *)(UINTN)TempBlockDesc->Union.ContinuationPointer;
 508     } else {
 509       //
 510       // If the capsule data pointed to by this descriptor is in the way,
 511       // move it.
 512       //
 513       if (IsOverlapped (
 514             (UINT8 *)NewBlockList,
 515             BufferSize,
 516             (UINT8 *)(UINTN)TempBlockDesc->Union.DataBlock,
 517             (UINTN)TempBlockDesc->Length
 518             ))
 519       {
 520         //
 521         // Relocate the block
 522         //
 523         RelocBuffer = FindFreeMem (BlockList, MemBase, MemSize, (UINTN)TempBlockDesc->Length);
 524         if (RelocBuffer == NULL) {
 525           return NULL;
 526         }
 527
 528         CopyMem ((VOID *)RelocBuffer, (VOID *)(UINTN)TempBlockDesc->Union.DataBlock, (UINTN)TempBlockDesc->Length);
 529         DEBUG ((DEBUG_INFO, "Capsule relocate descriptors from/to/size  0x%lX 0x%lX 0x%lX\n", TempBlockDesc->Union.DataBlock, (UINT64)(UINTN)RelocBuffer, TempBlockDesc->Length));
 530         TempBlockDesc->Union.DataBlock = (EFI_PHYSICAL_ADDRESS)(UINTN)RelocBuffer;
 531       }
 532
 533       TempBlockDesc++;
 534     }
 535   }
 536
 537   //
 538   // Now go through all the block descriptors to make sure that they're not
 539   // in the memory region we want to copy them to.
 540   //
 541   CurrBlockDescHead = BlockList;
 542   PrevBlockDescTail = NULL;
 543   while ((CurrBlockDescHead != NULL) && (CurrBlockDescHead->Union.ContinuationPointer != (EFI_PHYSICAL_ADDRESS)(UINTN)NULL)) {
 544     //
 545     // Get the size of this list then see if it overlaps our low region
 546     //
 547     TempBlockDesc = CurrBlockDescHead;
 548     BlockListSize = sizeof (EFI_CAPSULE_BLOCK_DESCRIPTOR);
 549     while (TempBlockDesc->Length != 0) {
 550       BlockListSize += sizeof (EFI_CAPSULE_BLOCK_DESCRIPTOR);
 551       TempBlockDesc++;
 552     }
 553
 554     if (IsOverlapped (
 555           (UINT8 *)NewBlockList,
 556           BufferSize,
 557           (UINT8 *)CurrBlockDescHead,
 558           BlockListSize
 559           ))
 560     {
 561       //
 562       // Overlaps, so move it out of the way
 563       //
 564       RelocBuffer = FindFreeMem (BlockList, MemBase, MemSize, BlockListSize);
 565       if (RelocBuffer == NULL) {
 566         return NULL;
 567       }
 568
 569       CopyMem ((VOID *)RelocBuffer, (VOID *)CurrBlockDescHead, BlockListSize);
 570       DEBUG ((DEBUG_INFO, "Capsule reloc descriptor block #2\n"));
 571       //
 572       // Point the previous block's next point to this copied version. If
 573       // the tail pointer is null, then this is the first descriptor block.
 574       //
 575       if (PrevBlockDescTail == NULL) {
 576         BlockList = (EFI_CAPSULE_BLOCK_DESCRIPTOR  *)RelocBuffer;
 577       } else {
 578         PrevBlockDescTail->Union.DataBlock = (EFI_PHYSICAL_ADDRESS)(UINTN)RelocBuffer;
 579       }
 580     }
 581
 582     //
 583     // Save our new tail and jump to the next block list
 584     //
 585     PrevBlockDescTail = TempBlockDesc;
 586     CurrBlockDescHead = (EFI_CAPSULE_BLOCK_DESCRIPTOR  *)(UINTN)TempBlockDesc->Union.ContinuationPointer;
 587   }
 588
 589   //
 590   // Cleared out low memory. Now copy the descriptors down there.
 591   //
 592   TempBlockDesc     = BlockList;
 593   CurrBlockDescHead = NewBlockList;
 594   while ((TempBlockDesc != NULL) && (TempBlockDesc->Union.ContinuationPointer != (EFI_PHYSICAL_ADDRESS)(UINTN)NULL)) {
 595     if (TempBlockDesc->Length != 0) {
 596       CurrBlockDescHead->Union.DataBlock = TempBlockDesc->Union.DataBlock;
 597       CurrBlockDescHead->Length          = TempBlockDesc->Length;
 598       CurrBlockDescHead++;
 599       TempBlockDesc++;
 600     } else {
 601       TempBlockDesc = (EFI_CAPSULE_BLOCK_DESCRIPTOR  *)(UINTN)TempBlockDesc->Union.ContinuationPointer;
 602     }
 603   }
 604
 605   //
 606   // Null terminate
 607   //
 608   CurrBlockDescHead->Union.ContinuationPointer = (EFI_PHYSICAL_ADDRESS)(UINTN)NULL;
 609   CurrBlockDescHead->Length                    = 0;
 610   return NewBlockList;
 611 }
 612
 613 /**
 614   Determine if two buffers overlap in memory.
 615
 616   @param Buff1   pointer to first buffer
 617   @param Size1   size of Buff1
 618   @param Buff2   pointer to second buffer
 619   @param Size2   size of Buff2
 620
 621   @retval TRUE    Buffers overlap in memory.
 622   @retval FALSE   Buffer doesn't overlap.
 623
 624 **/
 625 BOOLEAN
 626 IsOverlapped (
 627   UINT8  *Buff1,
 628   UINTN  Size1,
 629   UINT8  *Buff2,
 630   UINTN  Size2
 631   )
 632 {
 633   //
 634   // If buff1's end is less than the start of buff2, then it's ok.
 635   // Also, if buff1's start is beyond buff2's end, then it's ok.
 636   //
 637   if (((Buff1 + Size1) <= Buff2) || (Buff1 >= (Buff2 + Size2))) {
 638     return FALSE;
 639   }
 640
 641   return TRUE;
 642 }
 643
 644 /**
 645   Given a pointer to a capsule block descriptor, traverse the list to figure
 646   out how many legitimate descriptors there are, and how big the capsule it
 647   refers to is.
 648
 649   @param Desc            Pointer to the capsule block descriptors
 650   @param NumDescriptors  Optional pointer to where to return the number of capsule data descriptors, whose Length is non-zero.
 651   @param CapsuleSize     Optional pointer to where to return the capsule image size
 652   @param CapsuleNumber   Optional pointer to where to return the number of capsule
 653
 654   @retval EFI_NOT_FOUND   No descriptors containing data in the list
 655   @retval EFI_SUCCESS     Return data is valid
 656
 657 **/
 658 EFI_STATUS
 659 GetCapsuleInfo (
 660   IN EFI_CAPSULE_BLOCK_DESCRIPTOR  *Desc,
 661   IN OUT UINTN                     *NumDescriptors OPTIONAL,
 662   IN OUT UINTN                     *CapsuleSize OPTIONAL,
 663   IN OUT UINTN                     *CapsuleNumber OPTIONAL
 664   )
 665 {
 666   UINTN               Count;
 667   UINTN               Size;
 668   UINTN               Number;
 669   UINTN               ThisCapsuleImageSize;
 670   EFI_CAPSULE_HEADER  *CapsuleHeader;
 671
 672   DEBUG ((DEBUG_INFO, "GetCapsuleInfo enter\n"));
 673
 674   ASSERT (Desc != NULL);
 675
 676   Count                = 0;
 677   Size                 = 0;
 678   Number               = 0;
 679   ThisCapsuleImageSize = 0;
 680
 681   while (Desc->Union.ContinuationPointer != (EFI_PHYSICAL_ADDRESS)(UINTN)NULL) {
 682     if (Desc->Length == 0) {
 683       //
 684       // Descriptor points to another list of block descriptors somewhere
 685       //
 686       Desc = (EFI_CAPSULE_BLOCK_DESCRIPTOR  *)(UINTN)Desc->Union.ContinuationPointer;
 687     } else {
 688       //
 689       // Sanity Check
 690       // It is needed, because ValidateCapsuleIntegrity() only validate one individual capsule Size.
 691       // While here we need check all capsules size.
 692       //
 693       if (Desc->Length >= (MAX_ADDRESS - Size)) {
 694         DEBUG ((DEBUG_ERROR, "ERROR: Desc->Length(0x%lx) >= (MAX_ADDRESS - Size(0x%x))\n", Desc->Length, Size));
 695         return EFI_OUT_OF_RESOURCES;
 696       }
 697
 698       Size += (UINTN)Desc->Length;
 699       Count++;
 700
 701       //
 702       // See if this is first capsule's header
 703       //
 704       if (ThisCapsuleImageSize == 0) {
 705         CapsuleHeader = (EFI_CAPSULE_HEADER *)((UINTN)Desc->Union.DataBlock);
 706         //
 707         // This has been checked in ValidateCapsuleIntegrity()
 708         //
 709         Number++;
 710         ThisCapsuleImageSize = CapsuleHeader->CapsuleImageSize;
 711       }
 712
 713       //
 714       // This has been checked in ValidateCapsuleIntegrity()
 715       //
 716       ASSERT (ThisCapsuleImageSize >= Desc->Length);
 717       ThisCapsuleImageSize = (UINTN)(ThisCapsuleImageSize - Desc->Length);
 718
 719       //
 720       // Move to next
 721       //
 722       Desc++;
 723     }
 724   }
 725
 726   //
 727   // If no descriptors, then fail
 728   //
 729   if (Count == 0) {
 730     DEBUG ((DEBUG_ERROR, "ERROR: Count == 0\n"));
 731     return EFI_NOT_FOUND;
 732   }
 733
 734   //
 735   // checked in ValidateCapsuleIntegrity()
 736   //
 737   ASSERT (ThisCapsuleImageSize == 0);
 738
 739   if (NumDescriptors != NULL) {
 740     *NumDescriptors = Count;
 741   }
 742
 743   if (CapsuleSize != NULL) {
 744     *CapsuleSize = Size;
 745   }
 746
 747   if (CapsuleNumber != NULL) {
 748     *CapsuleNumber = Number;
 749   }
 750
 751   return EFI_SUCCESS;
 752 }
 753
 754 /**
 755   Check every capsule header.
 756
 757   @param CapsuleHeader   The pointer to EFI_CAPSULE_HEADER
 758
 759   @retval FALSE  Capsule is OK
 760   @retval TRUE   Capsule is corrupted
 761
 762 **/
 763 BOOLEAN
 764 IsCapsuleCorrupted (
 765   IN EFI_CAPSULE_HEADER  *CapsuleHeader
 766   )
 767 {
 768   //
 769   // A capsule to be updated across a system reset should contain CAPSULE_FLAGS_PERSIST_ACROSS_RESET.
 770   //
 771   if ((CapsuleHeader->Flags & CAPSULE_FLAGS_PERSIST_ACROSS_RESET) == 0) {
 772     return TRUE;
 773   }
 774
 775   //
 776   // Make sure the flags combination is supported by the platform.
 777   //
 778   if ((CapsuleHeader->Flags & (CAPSULE_FLAGS_PERSIST_ACROSS_RESET | CAPSULE_FLAGS_POPULATE_SYSTEM_TABLE)) == CAPSULE_FLAGS_POPULATE_SYSTEM_TABLE) {
 779     return TRUE;
 780   }
 781
 782   if ((CapsuleHeader->Flags & (CAPSULE_FLAGS_PERSIST_ACROSS_RESET | CAPSULE_FLAGS_INITIATE_RESET)) == CAPSULE_FLAGS_INITIATE_RESET) {
 783     return TRUE;
 784   }
 785
 786   return FALSE;
 787 }
 788
 789 /**
 790   Try to verify the integrity of a capsule test pattern before the
 791   capsule gets coalesced. This can be useful in narrowing down
 792   where capsule data corruption occurs.
 793
 794   The test pattern mode fills in memory with a counting UINT32 value.
 795   If the capsule is not divided up in a multiple of 4-byte blocks, then
 796   things get messy doing the check. Therefore there are some cases
 797   here where we just give up and skip the pre-coalesce check.
 798
 799   @param PeiServices  PEI services table
 800   @param Desc         Pointer to capsule descriptors
 801 **/
 802 VOID
 803 CapsuleTestPatternPreCoalesce (
 804   IN EFI_PEI_SERVICES              **PeiServices,
 805   IN EFI_CAPSULE_BLOCK_DESCRIPTOR  *Desc
 806   )
 807 {
 808   UINT32  *TestPtr;
 809   UINT32  TestCounter;
 810   UINT32  TestSize;
 811
 812   DEBUG ((DEBUG_INFO, "CapsuleTestPatternPreCoalesce\n"));
 813
 814   //
 815   // Find first data descriptor
 816   //
 817   while ((Desc->Length == 0) && (Desc->Union.ContinuationPointer != (EFI_PHYSICAL_ADDRESS)(UINTN)NULL)) {
 818     Desc = (EFI_CAPSULE_BLOCK_DESCRIPTOR  *)(UINTN)Desc->Union.ContinuationPointer;
 819   }
 820
 821   if (Desc->Union.ContinuationPointer == 0) {
 822     return;
 823   }
 824
 825   //
 826   // First one better be long enough to at least hold the test signature
 827   //
 828   if (Desc->Length < sizeof (UINT32)) {
 829     DEBUG ((DEBUG_INFO, "Capsule test pattern pre-coalesce punted #1\n"));
 830     return;
 831   }
 832
 833   TestPtr = (UINT32 *)(UINTN)Desc->Union.DataBlock;
 834   //
 835   // 0x54534554 "TEST"
 836   //
 837   if (*TestPtr != 0x54534554) {
 838     return;
 839   }
 840
 841   TestCounter = 0;
 842   TestSize    = (UINT32)Desc->Length - 2 * sizeof (UINT32);
 843   //
 844   // Skip over the signature and the size fields in the pattern data header
 845   //
 846   TestPtr += 2;
 847   while (1) {
 848     if ((TestSize & 0x03) != 0) {
 849       DEBUG ((DEBUG_INFO, "Capsule test pattern pre-coalesce punted #2\n"));
 850       return;
 851     }
 852
 853     while (TestSize > 0) {
 854       if (*TestPtr != TestCounter) {
 855         DEBUG ((DEBUG_INFO, "Capsule test pattern pre-coalesce failed data corruption check\n"));
 856         return;
 857       }
 858
 859       TestSize -= sizeof (UINT32);
 860       TestCounter++;
 861       TestPtr++;
 862     }
 863
 864     Desc++;
 865     while ((Desc->Length == 0) && (Desc->Union.ContinuationPointer != (EFI_PHYSICAL_ADDRESS)(UINTN)NULL)) {
 866       Desc = (EFI_CAPSULE_BLOCK_DESCRIPTOR  *)(UINTN)Desc->Union.ContinuationPointer;
 867     }
 868
 869     if (Desc->Union.ContinuationPointer == (EFI_PHYSICAL_ADDRESS)(UINTN)NULL) {
 870       return;
 871     }
 872
 873     TestSize = (UINT32)Desc->Length;
 874     TestPtr  = (UINT32 *)(UINTN)Desc->Union.DataBlock;
 875   }
 876 }
 877
 878 /**
 879   Checks for the presence of capsule descriptors.
 880   Get capsule descriptors from variable CapsuleUpdateData, CapsuleUpdateData1, CapsuleUpdateData2...
 881
 882   @param BlockListBuffer            Pointer to the buffer of capsule descriptors variables
 883   @param MemoryResource             Pointer to the buffer of memory resource descriptor.
 884   @param BlockDescriptorList        Pointer to the capsule descriptors list
 885
 886   @retval EFI_SUCCESS               a valid capsule is present
 887   @retval EFI_NOT_FOUND             if a valid capsule is not present
 888 **/
 889 EFI_STATUS
 890 BuildCapsuleDescriptors (
 891   IN  EFI_PHYSICAL_ADDRESS          *BlockListBuffer,
 892   IN  MEMORY_RESOURCE_DESCRIPTOR    *MemoryResource,
 893   OUT EFI_CAPSULE_BLOCK_DESCRIPTOR  **BlockDescriptorList
 894   )
 895 {
 896   UINTN                         Index;
 897   EFI_CAPSULE_BLOCK_DESCRIPTOR  *LastBlock;
 898   EFI_CAPSULE_BLOCK_DESCRIPTOR  *TempBlock;
 899   EFI_CAPSULE_BLOCK_DESCRIPTOR  *HeadBlock;
 900
 901   DEBUG ((DEBUG_INFO, "BuildCapsuleDescriptors enter\n"));
 902
 903   LastBlock = NULL;
 904   HeadBlock = NULL;
 905   TempBlock = NULL;
 906   Index     = 0;
 907
 908   while (BlockListBuffer[Index] != 0) {
 909     //
 910     // Test integrity of descriptors.
 911     //
 912     if (BlockListBuffer[Index] < MAX_ADDRESS) {
 913       TempBlock = ValidateCapsuleIntegrity ((EFI_CAPSULE_BLOCK_DESCRIPTOR *)(UINTN)BlockListBuffer[Index], MemoryResource);
 914       if (TempBlock != NULL) {
 915         if (LastBlock == NULL) {
 916           LastBlock = TempBlock;
 917
 918           //
 919           // Return the base of the block descriptors
 920           //
 921           HeadBlock = (EFI_CAPSULE_BLOCK_DESCRIPTOR *)(UINTN)BlockListBuffer[Index];
 922         } else {
 923           //
 924           // Combine the different BlockList into single BlockList.
 925           //
 926           LastBlock->Union.DataBlock = (EFI_PHYSICAL_ADDRESS)(UINTN)BlockListBuffer[Index];
 927           LastBlock->Length          = 0;
 928           LastBlock                  = TempBlock;
 929         }
 930       }
 931     } else {
 932       DEBUG ((DEBUG_ERROR, "ERROR: BlockListBuffer[Index](0x%lx) < MAX_ADDRESS\n", BlockListBuffer[Index]));
 933     }
 934
 935     Index++;
 936   }
 937
 938   if (HeadBlock != NULL) {
 939     *BlockDescriptorList = HeadBlock;
 940     return EFI_SUCCESS;
 941   }
 942
 943   return EFI_NOT_FOUND;
 944 }
 945
 946 /**
 947   The function to coalesce a fragmented capsule in memory.
 948
 949   Memory Map for coalesced capsule:
 950   MemBase +   ---->+---------------------------+<-----------+
 951   MemSize          | ------------------------- |            |
 952                    | |  Capsule [Num-1]      | |            |
 953                    | ------------------------- |            |
 954                    | |  ................     | |            |
 955                    | ------------------------- |            |
 956                    | |  Capsule [1]          | |            |
 957                    | ------------------------- |            |
 958                    | |  Capsule [0]          | |            |
 959                    | ------------------------- |            |
 960                    |    Capsule Image          |            |
 961 CapsuleImageBase-->+---------------------------+
 962                    | ------------------------- |            |
 963                    | |  CapsuleOffset[Num-1] | |            |
 964                    | ------------------------- |            |
 965                    | |  ................     | |        CapsuleSize
 966                    | ------------------------- |            |
 967                    | |  CapsuleOffset[1]     | |            |
 968                    | ------------------------- |            |
 969                    | |  CapsuleOffset[0]     | |            |
 970                    |---------------------------|            |
 971                    | |  CapsuleNumber        | |            |
 972                    | ------------------------- |            |
 973                    | |  CapsuleAllImageSize  | |            |
 974                    | ------------------------- |            |
 975                    |    PrivateData            |            |
 976      DestPtr  ---->+---------------------------+<-----------+
 977                    |                           |            |
 978                    |     FreeMem               |        FreeMemSize
 979                    |                           |            |
 980    FreeMemBase --->+---------------------------+<-----------+
 981                    |    Terminator             |
 982                    +---------------------------+
 983                    |    BlockDescriptor n      |
 984                    +---------------------------+
 985                    |    .................      |
 986                    +---------------------------+
 987                    |    BlockDescriptor 1      |
 988                    +---------------------------+
 989                    |    BlockDescriptor 0      |
 990                    +---------------------------+
 991                    |    PrivateDataDesc 0      |
 992       MemBase ---->+---------------------------+<----- BlockList
 993
 994   Caution: This function may receive untrusted input.
 995   The capsule data is external input, so this routine will do basic validation before
 996   coalesce capsule data into memory.
 997
 998   @param PeiServices        General purpose services available to every PEIM.
 999   @param BlockListBuffer    Pointer to the buffer of Capsule Descriptor Variables.
1000   @param MemoryResource     Pointer to the buffer of memory resource descriptor.
1001   @param MemoryBase         Pointer to the base of a block of memory that we can walk
1002                             all over while trying to coalesce our buffers.
1003                             On output, this variable will hold the base address of
1004                             a coalesced capsule.
1005   @param MemorySize         Size of the memory region pointed to by MemoryBase.
1006                             On output, this variable will contain the size of the
1007                             coalesced capsule.
1008
1009   @retval EFI_NOT_FOUND     If we could not find the capsule descriptors.
1010
1011   @retval EFI_BUFFER_TOO_SMALL
1012                             If we could not coalesce the capsule in the memory
1013                             region provided to us.
1014
1015   @retval EFI_SUCCESS       Processed the capsule successfully.
1016 **/
1017 EFI_STATUS
1018 EFIAPI
1019 CapsuleDataCoalesce (
1020   IN EFI_PEI_SERVICES            **PeiServices,
1021   IN EFI_PHYSICAL_ADDRESS        *BlockListBuffer,
1022   IN MEMORY_RESOURCE_DESCRIPTOR  *MemoryResource,
1023   IN OUT VOID                    **MemoryBase,
1024   IN OUT UINTN                   *MemorySize
1025   )
1026 {
1027   VOID                           *NewCapsuleBase;
1028   VOID                           *CapsuleImageBase;
1029   UINTN                          CapsuleIndex;
1030   UINT8                          *FreeMemBase;
1031   UINT8                          *DestPtr;
1032   UINTN                          DestLength;
1033   UINT8                          *RelocPtr;
1034   UINTN                          CapsuleTimes;
1035   UINT64                         SizeLeft;
1036   UINT64                         CapsuleImageSize;
1037   UINTN                          CapsuleSize;
1038   UINTN                          CapsuleNumber;
1039   UINTN                          DescriptorsSize;
1040   UINTN                          FreeMemSize;
1041   UINTN                          NumDescriptors;
1042   BOOLEAN                        CapsuleBeginFlag;
1043   EFI_STATUS                     Status;
1044   EFI_CAPSULE_HEADER             *CapsuleHeader;
1045   EFI_CAPSULE_PEIM_PRIVATE_DATA  PrivateData;
1046   EFI_CAPSULE_PEIM_PRIVATE_DATA  *PrivateDataPtr;
1047   EFI_CAPSULE_BLOCK_DESCRIPTOR   *BlockList;
1048   EFI_CAPSULE_BLOCK_DESCRIPTOR   *CurrentBlockDesc;
1049   EFI_CAPSULE_BLOCK_DESCRIPTOR   *TempBlockDesc;
1050   EFI_CAPSULE_BLOCK_DESCRIPTOR   PrivateDataDesc[2];
1051
1052   DEBUG ((DEBUG_INFO, "CapsuleDataCoalesce enter\n"));
1053
1054   CapsuleIndex     = 0;
1055   SizeLeft         = 0;
1056   CapsuleTimes     = 0;
1057   CapsuleImageSize = 0;
1058   PrivateDataPtr   = NULL;
1059   CapsuleHeader    = NULL;
1060   CapsuleBeginFlag = TRUE;
1061   CapsuleSize      = 0;
1062   NumDescriptors   = 0;
1063
1064   //
1065   // Build capsule descriptors list
1066   //
1067   Status = BuildCapsuleDescriptors (BlockListBuffer, MemoryResource, &BlockList);
1068   if (EFI_ERROR (Status)) {
1069     return Status;
1070   }
1071
1072   DEBUG_CODE (
1073     CapsuleTestPatternPreCoalesce (PeiServices, BlockList);
1074     );
1075
1076   //
1077   // Get the size of our descriptors and the capsule size. GetCapsuleInfo()
1078   // returns the number of descriptors that actually point to data, so add
1079   // one for a terminator. Do that below.
1080   //
1081   Status = GetCapsuleInfo (BlockList, &NumDescriptors, &CapsuleSize, &CapsuleNumber);
1082   if (EFI_ERROR (Status)) {
1083     return Status;
1084   }
1085
1086   DEBUG ((DEBUG_INFO, "CapsuleSize - 0x%x\n", CapsuleSize));
1087   DEBUG ((DEBUG_INFO, "CapsuleNumber - 0x%x\n", CapsuleNumber));
1088   DEBUG ((DEBUG_INFO, "NumDescriptors - 0x%x\n", NumDescriptors));
1089   if ((CapsuleSize == 0) || (NumDescriptors == 0) || (CapsuleNumber == 0)) {
1090     return EFI_NOT_FOUND;
1091   }
1092
1093   if (CapsuleNumber - 1 >= (MAX_ADDRESS - (sizeof (EFI_CAPSULE_PEIM_PRIVATE_DATA)  + sizeof (UINT64))) / sizeof (UINT64)) {
1094     DEBUG ((DEBUG_ERROR, "ERROR: CapsuleNumber - 0x%x\n", CapsuleNumber));
1095     return EFI_BUFFER_TOO_SMALL;
1096   }
1097
1098   //
1099   // Initialize our local copy of private data. When we're done, we'll create a
1100   // descriptor for it as well so that it can be put into free memory without
1101   // trashing anything.
1102   //
1103   PrivateData.Signature           = EFI_CAPSULE_PEIM_PRIVATE_DATA_SIGNATURE;
1104   PrivateData.CapsuleAllImageSize = (UINT64)CapsuleSize;
1105   PrivateData.CapsuleNumber       = (UINT64)CapsuleNumber;
1106   PrivateData.CapsuleOffset[0]    = 0;
1107   //
1108   // NOTE: Only data in sizeof (EFI_CAPSULE_PEIM_PRIVATE_DATA) is valid, CapsuleOffset field is uninitialized at this moment.
1109   // The code sets partial length here for Descriptor.Length check, but later it will use full length to reserve those PrivateData region.
1110   //
1111   PrivateDataDesc[0].Union.DataBlock = (EFI_PHYSICAL_ADDRESS)(UINTN)&PrivateData;
1112   PrivateDataDesc[0].Length          = sizeof (EFI_CAPSULE_PEIM_PRIVATE_DATA);
1113   PrivateDataDesc[1].Union.DataBlock = (EFI_PHYSICAL_ADDRESS)(UINTN)BlockList;
1114   PrivateDataDesc[1].Length          = 0;
1115   //
1116   // Add PrivateDataDesc[0] in beginning, as it is new descriptor. PrivateDataDesc[1] is NOT needed.
1117   // In addition, one NULL terminator is added in the end. See RelocateBlockDescriptors().
1118   //
1119   NumDescriptors += 2;
1120   //
1121   // Sanity check
1122   //
1123   if (CapsuleSize >= (MAX_ADDRESS - (sizeof (EFI_CAPSULE_PEIM_PRIVATE_DATA) + (CapsuleNumber - 1) * sizeof (UINT64) + sizeof (UINT64)))) {
1124     DEBUG ((DEBUG_ERROR, "ERROR: CapsuleSize - 0x%x\n", CapsuleSize));
1125     return EFI_BUFFER_TOO_SMALL;
1126   }
1127
1128   //
1129   // Need add sizeof(UINT64) for PrivateData alignment
1130   //
1131   CapsuleSize += sizeof (EFI_CAPSULE_PEIM_PRIVATE_DATA) + (CapsuleNumber - 1) * sizeof (UINT64) + sizeof (UINT64);
1132   BlockList    = PrivateDataDesc;
1133   //
1134   // Sanity check
1135   //
1136   if (NumDescriptors >= (MAX_ADDRESS / sizeof (EFI_CAPSULE_BLOCK_DESCRIPTOR))) {
1137     DEBUG ((DEBUG_ERROR, "ERROR: NumDescriptors - 0x%x\n", NumDescriptors));
1138     return EFI_BUFFER_TOO_SMALL;
1139   }
1140
1141   DescriptorsSize = NumDescriptors * sizeof (EFI_CAPSULE_BLOCK_DESCRIPTOR);
1142   //
1143   // Sanity check
1144   //
1145   if (DescriptorsSize >= (MAX_ADDRESS - CapsuleSize)) {
1146     DEBUG ((DEBUG_ERROR, "ERROR: DescriptorsSize - 0x%lx, CapsuleSize - 0x%lx\n", (UINT64)DescriptorsSize, (UINT64)CapsuleSize));
1147     return EFI_BUFFER_TOO_SMALL;
1148   }
1149
1150   //
1151   // Don't go below some min address. If the base is below it,
1152   // then move it up and adjust the size accordingly.
1153   //
1154   DEBUG ((DEBUG_INFO, "Capsule Memory range from 0x%8X to 0x%8X\n", (UINTN)*MemoryBase, (UINTN)*MemoryBase + *MemorySize));
1155   if ((UINTN)*MemoryBase < (UINTN)MIN_COALESCE_ADDR) {
1156     if (((UINTN)*MemoryBase + *MemorySize) < (UINTN)MIN_COALESCE_ADDR) {
1157       DEBUG ((DEBUG_ERROR, "ERROR: *MemoryBase + *MemorySize - 0x%x\n", (UINTN)*MemoryBase + *MemorySize));
1158       return EFI_BUFFER_TOO_SMALL;
1159     } else {
1160       *MemorySize = *MemorySize - ((UINTN)MIN_COALESCE_ADDR - (UINTN)*MemoryBase);
1161       *MemoryBase = (VOID *)(UINTN)MIN_COALESCE_ADDR;
1162     }
1163   }
1164
1165   if (*MemorySize <= (CapsuleSize + DescriptorsSize)) {
1166     DEBUG ((DEBUG_ERROR, "ERROR: CapsuleSize + DescriptorsSize - 0x%x\n", CapsuleSize + DescriptorsSize));
1167     return EFI_BUFFER_TOO_SMALL;
1168   }
1169
1170   FreeMemBase = *MemoryBase;
1171   FreeMemSize = *MemorySize;
1172   DEBUG ((DEBUG_INFO, "Capsule Free Memory from 0x%8X to 0x%8X\n", (UINTN)FreeMemBase, (UINTN)FreeMemBase + FreeMemSize));
1173
1174   //
1175   // Relocate all the block descriptors to low memory to make further
1176   // processing easier.
1177   //
1178   BlockList = RelocateBlockDescriptors (PeiServices, BlockList, NumDescriptors, FreeMemBase, FreeMemSize);
1179   if (BlockList == NULL) {
1180     //
1181     // Not enough room to relocate the descriptors
1182     //
1183     return EFI_BUFFER_TOO_SMALL;
1184   }
1185
1186   //
1187   // Take the top of memory for the capsule. UINT64 align up.
1188   //
1189   DestPtr          = FreeMemBase + FreeMemSize - CapsuleSize;
1190   DestPtr          = (UINT8 *)(((UINTN)DestPtr + sizeof (UINT64) - 1) & ~(sizeof (UINT64) - 1));
1191   FreeMemBase      = (UINT8 *)BlockList + DescriptorsSize;
1192   FreeMemSize      = (UINTN)DestPtr - (UINTN)FreeMemBase;
1193   NewCapsuleBase   = (VOID *)DestPtr;
1194   CapsuleImageBase = (UINT8 *)NewCapsuleBase + sizeof (EFI_CAPSULE_PEIM_PRIVATE_DATA) + (CapsuleNumber - 1) * sizeof (UINT64);
1195
1196   PrivateDataPtr = (EFI_CAPSULE_PEIM_PRIVATE_DATA *)NewCapsuleBase;
1197
1198   //
1199   // Move all the blocks to the top (high) of memory.
1200   // Relocate all the obstructing blocks. Note that the block descriptors
1201   // were coalesced when they were relocated, so we can just ++ the pointer.
1202   //
1203   CurrentBlockDesc = BlockList;
1204   while ((CurrentBlockDesc->Length != 0) || (CurrentBlockDesc->Union.ContinuationPointer != (EFI_PHYSICAL_ADDRESS)(UINTN)NULL)) {
1205     if (CapsuleTimes == 0) {
1206       //
1207       // The first entry is the block descriptor for EFI_CAPSULE_PEIM_PRIVATE_DATA.
1208       // CapsuleOffset field is uninitialized at this time. No need copy it, but need to reserve for future use.
1209       //
1210       ASSERT (CurrentBlockDesc->Union.DataBlock == (UINT64)(UINTN)&PrivateData);
1211       DestLength = sizeof (EFI_CAPSULE_PEIM_PRIVATE_DATA) + (CapsuleNumber - 1) * sizeof (UINT64);
1212     } else {
1213       DestLength = (UINTN)CurrentBlockDesc->Length;
1214     }
1215
1216     //
1217     // See if any of the remaining capsule blocks are in the way
1218     //
1219     TempBlockDesc = CurrentBlockDesc;
1220     while (TempBlockDesc->Length != 0) {
1221       //
1222       // Is this block in the way of where we want to copy the current descriptor to?
1223       //
1224       if (IsOverlapped (
1225             (UINT8 *)DestPtr,
1226             (UINTN)DestLength,
1227             (UINT8 *)(UINTN)TempBlockDesc->Union.DataBlock,
1228             (UINTN)TempBlockDesc->Length
1229             ))
1230       {
1231         //
1232         // Relocate the block
1233         //
1234         RelocPtr = FindFreeMem (BlockList, FreeMemBase, FreeMemSize, (UINTN)TempBlockDesc->Length);
1235         if (RelocPtr == NULL) {
1236           return EFI_BUFFER_TOO_SMALL;
1237         }
1238
1239         CopyMem ((VOID *)RelocPtr, (VOID *)(UINTN)TempBlockDesc->Union.DataBlock, (UINTN)TempBlockDesc->Length);
1240         DEBUG ((
1241           DEBUG_INFO,
1242           "Capsule reloc data block from 0x%8X to 0x%8X with size 0x%8X\n",
1243           (UINTN)TempBlockDesc->Union.DataBlock,
1244           (UINTN)RelocPtr,
1245           (UINTN)TempBlockDesc->Length
1246           ));
1247
1248         TempBlockDesc->Union.DataBlock = (EFI_PHYSICAL_ADDRESS)(UINTN)RelocPtr;
1249       }
1250
1251       //
1252       // Next descriptor
1253       //
1254       TempBlockDesc++;
1255     }
1256
1257     //
1258     // Ok, we made it through. Copy the block.
1259     // we just support greping one capsule from the lists of block descs list.
1260     //
1261     CapsuleTimes++;
1262     //
1263     // Skip the first block descriptor that filled with EFI_CAPSULE_PEIM_PRIVATE_DATA
1264     //
1265     if (CapsuleTimes > 1) {
1266       //
1267       // For every capsule entry point, check its header to determine whether to relocate it.
1268       // If it is invalid, skip it and move on to the next capsule. If it is valid, relocate it.
1269       //
1270       if (CapsuleBeginFlag) {
1271         CapsuleBeginFlag = FALSE;
1272         CapsuleHeader    = (EFI_CAPSULE_HEADER *)(UINTN)CurrentBlockDesc->Union.DataBlock;
1273         SizeLeft         = CapsuleHeader->CapsuleImageSize;
1274
1275         //
1276         // No more check here is needed, because IsCapsuleCorrupted() already in ValidateCapsuleIntegrity()
1277         //
1278         ASSERT (CapsuleIndex < CapsuleNumber);
1279
1280         //
1281         // Relocate this capsule
1282         //
1283         CapsuleImageSize += SizeLeft;
1284         //
1285         // Cache the begin offset of this capsule
1286         //
1287         ASSERT (PrivateDataPtr->Signature == EFI_CAPSULE_PEIM_PRIVATE_DATA_SIGNATURE);
1288         ASSERT ((UINTN)DestPtr >= (UINTN)CapsuleImageBase);
1289         PrivateDataPtr->CapsuleOffset[CapsuleIndex++] = (UINTN)DestPtr - (UINTN)CapsuleImageBase;
1290       }
1291
1292       //
1293       // Below ASSERT is checked in ValidateCapsuleIntegrity()
1294       //
1295       ASSERT (CurrentBlockDesc->Length <= SizeLeft);
1296
1297       CopyMem ((VOID *)DestPtr, (VOID *)(UINTN)(CurrentBlockDesc->Union.DataBlock), (UINTN)CurrentBlockDesc->Length);
1298       DEBUG ((
1299         DEBUG_INFO,
1300         "Capsule coalesce block no.0x%lX from 0x%lX to 0x%lX with size 0x%lX\n",
1301         (UINT64)CapsuleTimes,
1302         CurrentBlockDesc->Union.DataBlock,
1303         (UINT64)(UINTN)DestPtr,
1304         CurrentBlockDesc->Length
1305         ));
1306       DestPtr  += CurrentBlockDesc->Length;
1307       SizeLeft -= CurrentBlockDesc->Length;
1308
1309       if (SizeLeft == 0) {
1310         //
1311         // Here is the end of the current capsule image.
1312         //
1313         CapsuleBeginFlag = TRUE;
1314       }
1315     } else {
1316       //
1317       // The first entry is the block descriptor for EFI_CAPSULE_PEIM_PRIVATE_DATA.
1318       // CapsuleOffset field is uninitialized at this time. No need copy it, but need to reserve for future use.
1319       //
1320       ASSERT (CurrentBlockDesc->Length == sizeof (EFI_CAPSULE_PEIM_PRIVATE_DATA));
1321       ASSERT ((UINTN)DestPtr == (UINTN)NewCapsuleBase);
1322       CopyMem ((VOID *)DestPtr, (VOID *)(UINTN)CurrentBlockDesc->Union.DataBlock, (UINTN)CurrentBlockDesc->Length);
1323       DestPtr += sizeof (EFI_CAPSULE_PEIM_PRIVATE_DATA) + (CapsuleNumber - 1) * sizeof (UINT64);
1324     }
1325
1326     //
1327     // Walk through the block descriptor list.
1328     //
1329     CurrentBlockDesc++;
1330   }
1331
1332   //
1333   // We return the base of memory we want reserved, and the size.
1334   // The memory peim should handle it appropriately from there.
1335   //
1336   *MemorySize = (UINTN)CapsuleSize;
1337   *MemoryBase = (VOID *)NewCapsuleBase;
1338
1339   ASSERT (PrivateDataPtr->Signature == EFI_CAPSULE_PEIM_PRIVATE_DATA_SIGNATURE);
1340   ASSERT (PrivateDataPtr->CapsuleAllImageSize == CapsuleImageSize);
1341   ASSERT (PrivateDataPtr->CapsuleNumber == CapsuleIndex);
1342
1343   return EFI_SUCCESS;
1344 }