MdeModulePkg/Application/CapsuleApp/CapsuleApp.c

   1 /** @file
   2   A shell application that triggers capsule update process.
   3
   4   Copyright (c) 2016, Intel Corporation. All rights reserved.<BR>
   5   This program and the accompanying materials
   6   are licensed and made available under the terms and conditions of the BSD License
   7   which accompanies this distribution.  The full text of the license may be found at
   8   http://opensource.org/licenses/bsd-license.php
   9
  10   THE PROGRAM IS DISTRIBUTED UNDER THE BSD LICENSE ON AN "AS IS" BASIS,
  11   WITHOUT WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.
  12
  13 **/
  14
  15 #include <Uefi.h>
  16 #include <Library/BaseLib.h>
  17 #include <Library/DebugLib.h>
  18 #include <Library/BaseMemoryLib.h>
  19 #include <Library/MemoryAllocationLib.h>
  20 #include <Library/UefiBootServicesTableLib.h>
  21 #include <Library/UefiRuntimeServicesTableLib.h>
  22 #include <Library/UefiLib.h>
  23 #include <Library/PrintLib.h>
  24 #include <Protocol/LoadedImage.h>
  25 #include <Protocol/SimpleFileSystem.h>
  26 #include <Protocol/GraphicsOutput.h>
  27 #include <Guid/FileInfo.h>
  28 #include <Guid/Gpt.h>
  29 #include <Guid/GlobalVariable.h>
  30 #include <Guid/CapsuleReport.h>
  31 #include <Guid/SystemResourceTable.h>
  32 #include <Guid/FmpCapsule.h>
  33 #include <IndustryStandard/WindowsUxCapsule.h>
  34
  35 #define CAPSULE_HEADER_SIZE  0x20
  36
  37 #define NESTED_CAPSULE_HEADER_SIZE  SIZE_4KB
  38 #define SYSTEM_FIRMWARE_FLAG 0x50000
  39 #define DEVICE_FIRMWARE_FLAG 0x78010
  40
  41 #define EFI_CAPSULE_FROM_FILE_DIR                           L"\\EFI\\UpdateCapsule\\"
  42 #define EFI_OS_INDICATIONS_FILE_CAPSULE_DELIVERY_SUPPORTED  0x0000000000000004
  43
  44 #define MAJOR_VERSION   1
  45 #define MINOR_VERSION   0
  46
  47 #define MAX_CAPSULE_NUM 10
  48
  49 extern UINTN  Argc;
  50 extern CHAR16 **Argv;
  51
  52 //
  53 // Define how many block descriptors we want to test with.
  54 //
  55 UINTN  NumberOfDescriptors = 1;
  56 UINTN  CapsuleFirstIndex;
  57 UINTN  CapsuleLastIndex;
  58
  59 /**
  60   Dump capsule information
  61
  62   @param[in] CapsuleName  The name of the capsule image.
  63
  64   @retval EFI_SUCCESS            The capsule information is dumped.
  65   @retval EFI_UNSUPPORTED        Input parameter is not valid.
  66 **/
  67 EFI_STATUS
  68 DumpCapsule (
  69   IN CHAR16                                        *CapsuleName
  70   );
  71
  72 /**
  73   Dump capsule status variable.
  74
  75   @retval EFI_SUCCESS            The capsule status variable is dumped.
  76   @retval EFI_UNSUPPORTED        Input parameter is not valid.
  77 **/
  78 EFI_STATUS
  79 DmpCapsuleStatusVariable (
  80   VOID
  81   );
  82
  83 /**
  84   Dump FMP protocol info.
  85 **/
  86 VOID
  87 DumpFmpData (
  88   VOID
  89   );
  90
  91 /**
  92   Dump ESRT info.
  93 **/
  94 VOID
  95 DumpEsrtData (
  96   VOID
  97   );
  98
  99 /**
 100   Read a file.
 101
 102   @param[in]  FileName        The file to be read.
 103   @param[out] BufferSize      The file buffer size
 104   @param[out] Buffer          The file buffer
 105
 106   @retval EFI_SUCCESS    Read file successfully
 107   @retval EFI_NOT_FOUND  File not found
 108 **/
 109 EFI_STATUS
 110 ReadFileToBuffer (
 111   IN  CHAR16                               *FileName,
 112   OUT UINTN                                *BufferSize,
 113   OUT VOID                                 **Buffer
 114   );
 115
 116 /**
 117   Write a file.
 118
 119   @param[in] FileName        The file to be written.
 120   @param[in] BufferSize      The file buffer size
 121   @param[in] Buffer          The file buffer
 122
 123   @retval EFI_SUCCESS    Write file successfully
 124 **/
 125 EFI_STATUS
 126 WriteFileFromBuffer (
 127   IN  CHAR16                               *FileName,
 128   IN  UINTN                                BufferSize,
 129   IN  VOID                                 *Buffer
 130   );
 131
 132 /**
 133
 134   This function parse application ARG.
 135
 136   @return Status
 137 **/
 138 EFI_STATUS
 139 GetArg (
 140   VOID
 141   );
 142
 143 /**
 144   Create UX capsule.
 145
 146   @retval EFI_SUCCESS            The capsule header is appended.
 147   @retval EFI_UNSUPPORTED        Input parameter is not valid.
 148   @retval EFI_OUT_OF_RESOURCES   No enough resource to create UX capsule.
 149 **/
 150 EFI_STATUS
 151 CreateBmpFmp (
 152   VOID
 153   )
 154 {
 155   CHAR16                                        *OutputCapsuleName;
 156   VOID                                          *BmpBuffer;
 157   UINTN                                         FileSize;
 158   CHAR16                                        *BmpName;
 159   UINT8                                         *FullCapsuleBuffer;
 160   UINTN                                         FullCapsuleBufferSize;
 161   EFI_DISPLAY_CAPSULE                           *DisplayCapsule;
 162   EFI_STATUS                                    Status;
 163   EFI_GRAPHICS_OUTPUT_PROTOCOL                  *Gop;
 164
 165   Status = gBS->LocateProtocol(&gEfiGraphicsOutputProtocolGuid, NULL, (VOID **)&Gop);
 166   if (EFI_ERROR(Status)) {
 167     Print(L"CapsuleApp: NO GOP is found.\n");
 168     return EFI_UNSUPPORTED;
 169   }
 170   Print(L"Current GOP: Mode - %d, ", Gop->Mode->Mode);
 171   Print(L"HorizontalResolution - %d, ", Gop->Mode->Info->HorizontalResolution);
 172   Print(L"VerticalResolution - %d\n", Gop->Mode->Info->VerticalResolution);
 173   // HorizontalResolution >= BMP_IMAGE_HEADER.PixelWidth
 174   // VerticalResolution   >= BMP_IMAGE_HEADER.PixelHeight
 175
 176   if (Argc != 5) {
 177     Print(L"CapsuleApp: Invalid Parameter.\n");
 178     return EFI_UNSUPPORTED;
 179   }
 180
 181   if (StrCmp(Argv[3], L"-O") != 0) {
 182     Print(L"CapsuleApp: NO output capsule name.\n");
 183     return EFI_UNSUPPORTED;
 184   }
 185   OutputCapsuleName = Argv[4];
 186
 187   BmpBuffer = NULL;
 188   FileSize = 0;
 189   FullCapsuleBuffer = NULL;
 190
 191   BmpName = Argv[2];
 192   Status = ReadFileToBuffer(BmpName, &FileSize, &BmpBuffer);
 193   if (EFI_ERROR(Status)) {
 194     Print(L"CapsuleApp: BMP image (%s) is not found.\n", BmpName);
 195     goto Done;
 196   }
 197
 198   FullCapsuleBufferSize = sizeof(EFI_DISPLAY_CAPSULE) + FileSize;
 199   FullCapsuleBuffer = AllocatePool(FullCapsuleBufferSize);
 200   if (FullCapsuleBuffer == NULL) {
 201     Print(L"CapsuleApp: Capsule Buffer size (0x%x) too big.\n", FullCapsuleBufferSize);
 202     Status = EFI_OUT_OF_RESOURCES;
 203     goto Done;
 204   }
 205
 206   DisplayCapsule = (EFI_DISPLAY_CAPSULE *)FullCapsuleBuffer;
 207   CopyGuid(&DisplayCapsule->CapsuleHeader.CapsuleGuid, &gWindowsUxCapsuleGuid);
 208   DisplayCapsule->CapsuleHeader.HeaderSize = sizeof(DisplayCapsule->CapsuleHeader);
 209   DisplayCapsule->CapsuleHeader.Flags = CAPSULE_FLAGS_PERSIST_ACROSS_RESET;
 210   DisplayCapsule->CapsuleHeader.CapsuleImageSize = (UINT32)FullCapsuleBufferSize;
 211
 212   DisplayCapsule->ImagePayload.Version = 1;
 213   DisplayCapsule->ImagePayload.Checksum = 0;
 214   DisplayCapsule->ImagePayload.ImageType = 0; // BMP
 215   DisplayCapsule->ImagePayload.Reserved = 0;
 216   DisplayCapsule->ImagePayload.Mode = Gop->Mode->Mode;
 217   DisplayCapsule->ImagePayload.OffsetX = 0;
 218   DisplayCapsule->ImagePayload.OffsetY = 0;
 219
 220   CopyMem((DisplayCapsule + 1), BmpBuffer, FileSize);
 221
 222   DisplayCapsule->ImagePayload.Checksum = CalculateCheckSum8(FullCapsuleBuffer, FullCapsuleBufferSize);
 223
 224   Status = WriteFileFromBuffer(OutputCapsuleName, FullCapsuleBufferSize, FullCapsuleBuffer);
 225   Print(L"CapsuleApp: Write %s %r\n", OutputCapsuleName, Status);
 226
 227 Done:
 228   if (BmpBuffer != NULL) {
 229     FreePool(BmpBuffer);
 230   }
 231
 232   if (FullCapsuleBuffer != NULL) {
 233     FreePool(FullCapsuleBuffer);
 234   }
 235
 236   return Status;
 237 }
 238
 239 /**
 240   Get ImageTypeId in the FMP capsule header.
 241
 242   @param[in] CapsuleHeader  The FMP capsule image header.
 243
 244   @return ImageTypeId
 245 **/
 246 EFI_GUID *
 247 GetCapsuleImageTypeId (
 248   IN EFI_CAPSULE_HEADER                            *CapsuleHeader
 249   )
 250 {
 251   EFI_FIRMWARE_MANAGEMENT_CAPSULE_HEADER       *FmpCapsuleHeader;
 252   UINT64                                       *ItemOffsetList;
 253   EFI_FIRMWARE_MANAGEMENT_CAPSULE_IMAGE_HEADER *ImageHeader;
 254
 255   FmpCapsuleHeader = (EFI_FIRMWARE_MANAGEMENT_CAPSULE_HEADER *)((UINT8 *)CapsuleHeader + CapsuleHeader->HeaderSize);
 256   ItemOffsetList = (UINT64 *)(FmpCapsuleHeader + 1);
 257   if (FmpCapsuleHeader->PayloadItemCount == 0) {
 258     return NULL;
 259   }
 260   ImageHeader = (EFI_FIRMWARE_MANAGEMENT_CAPSULE_IMAGE_HEADER *)((UINT8 *)FmpCapsuleHeader + ItemOffsetList[FmpCapsuleHeader->EmbeddedDriverCount]);
 261   return &ImageHeader->UpdateImageTypeId;
 262 }
 263
 264 /**
 265   Get ESRT FwType according to ImageTypeId
 266
 267   @param[in]  ImageTypeId   ImageTypeId of an FMP capsule.
 268
 269   @return ESRT FwType
 270 **/
 271 UINT32
 272 GetEsrtFwType (
 273   IN  EFI_GUID                                      *ImageTypeId
 274   )
 275 {
 276   EFI_STATUS                 Status;
 277   EFI_SYSTEM_RESOURCE_TABLE  *Esrt;
 278   EFI_SYSTEM_RESOURCE_ENTRY  *EsrtEntry;
 279   UINTN                      Index;
 280
 281   //
 282   // Check ESRT
 283   //
 284   Status = EfiGetSystemConfigurationTable(&gEfiSystemResourceTableGuid, (VOID **)&Esrt);
 285   if (!EFI_ERROR(Status)) {
 286     ASSERT(Esrt != NULL);
 287     EsrtEntry = (VOID *)(Esrt + 1);
 288     for (Index = 0; Index < Esrt->FwResourceCount; Index++, EsrtEntry++) {
 289       if (CompareGuid(&EsrtEntry->FwClass, ImageTypeId)) {
 290         return EsrtEntry->FwType;
 291       }
 292     }
 293   }
 294
 295   return ESRT_FW_TYPE_UNKNOWN;
 296 }
 297
 298 /**
 299   Append a capsule header on top of current image.
 300   This function follows Windows UEFI Firmware Update Platform document.
 301
 302   @retval EFI_SUCCESS            The capsule header is appended.
 303   @retval EFI_UNSUPPORTED        Input parameter is not valid.
 304   @retval EFI_OUT_OF_RESOURCES   No enough resource to append capsule header.
 305 **/
 306 EFI_STATUS
 307 CreateNestedFmp (
 308   VOID
 309   )
 310 {
 311   CHAR16                                        *OutputCapsuleName;
 312   VOID                                          *CapsuleBuffer;
 313   UINTN                                         FileSize;
 314   CHAR16                                        *CapsuleName;
 315   UINT8                                         *FullCapsuleBuffer;
 316   UINTN                                         FullCapsuleBufferSize;
 317   EFI_CAPSULE_HEADER                            *NestedCapsuleHeader;
 318   EFI_GUID                                      *ImageTypeId;
 319   UINT32                                        FwType;
 320   EFI_STATUS                                    Status;
 321
 322   if (Argc != 5) {
 323     Print(L"CapsuleApp: Invalid Parameter.\n");
 324     return EFI_UNSUPPORTED;
 325   }
 326
 327   if (StrCmp(Argv[3], L"-O") != 0) {
 328     Print(L"CapsuleApp: NO output capsule name.\n");
 329     return EFI_UNSUPPORTED;
 330   }
 331   OutputCapsuleName = Argv[4];
 332
 333   CapsuleBuffer = NULL;
 334   FileSize = 0;
 335   FullCapsuleBuffer = NULL;
 336
 337   CapsuleName = Argv[2];
 338   Status = ReadFileToBuffer(CapsuleName, &FileSize, &CapsuleBuffer);
 339   if (EFI_ERROR(Status)) {
 340     Print(L"CapsuleApp: Capsule image (%s) is not found.\n", CapsuleName);
 341     goto Done;
 342   }
 343
 344   ImageTypeId = GetCapsuleImageTypeId(CapsuleBuffer);
 345   if (ImageTypeId == NULL) {
 346     Print(L"CapsuleApp: Capsule ImageTypeId is not found.\n");
 347     goto Done;
 348   }
 349   FwType = GetEsrtFwType(ImageTypeId);
 350   if ((FwType != ESRT_FW_TYPE_SYSTEMFIRMWARE) && (FwType != ESRT_FW_TYPE_DEVICEFIRMWARE)) {
 351     Print(L"CapsuleApp: Capsule FwType is invalid.\n");
 352     goto Done;
 353   }
 354
 355   FullCapsuleBufferSize = NESTED_CAPSULE_HEADER_SIZE + FileSize;
 356   FullCapsuleBuffer = AllocatePool(FullCapsuleBufferSize);
 357   if (FullCapsuleBuffer == NULL) {
 358     Print(L"CapsuleApp: Capsule Buffer size (0x%x) too big.\n", FullCapsuleBufferSize);
 359     Status = EFI_OUT_OF_RESOURCES;
 360     goto Done;
 361   }
 362
 363   NestedCapsuleHeader = (EFI_CAPSULE_HEADER *)FullCapsuleBuffer;
 364   ZeroMem(NestedCapsuleHeader, NESTED_CAPSULE_HEADER_SIZE);
 365   CopyGuid(&NestedCapsuleHeader->CapsuleGuid, ImageTypeId);
 366   NestedCapsuleHeader->HeaderSize = NESTED_CAPSULE_HEADER_SIZE;
 367   NestedCapsuleHeader->Flags = (FwType == ESRT_FW_TYPE_DEVICEFIRMWARE) ? SYSTEM_FIRMWARE_FLAG : DEVICE_FIRMWARE_FLAG;
 368   NestedCapsuleHeader->CapsuleImageSize = (UINT32)FullCapsuleBufferSize;
 369
 370   CopyMem((UINT8 *)NestedCapsuleHeader + NestedCapsuleHeader->HeaderSize, CapsuleBuffer, FileSize);
 371
 372   Status = WriteFileFromBuffer(OutputCapsuleName, FullCapsuleBufferSize, FullCapsuleBuffer);
 373   Print(L"CapsuleApp: Write %s %r\n", OutputCapsuleName, Status);
 374
 375 Done:
 376   if (CapsuleBuffer != NULL) {
 377     FreePool(CapsuleBuffer);
 378   }
 379
 380   if (FullCapsuleBuffer != NULL) {
 381     FreePool(FullCapsuleBuffer);
 382   }
 383
 384   return Status;
 385 }
 386
 387
 388 /**
 389   Clear capsule status variable.
 390
 391   @retval EFI_SUCCESS            The capsule status variable is cleared.
 392 **/
 393 EFI_STATUS
 394 ClearCapsuleStatusVariable (
 395   VOID
 396   )
 397 {
 398   EFI_STATUS                          Status;
 399   UINT32                              Index;
 400   CHAR16                              CapsuleVarName[20];
 401   CHAR16                              *TempVarName;
 402
 403   StrCpyS (CapsuleVarName, sizeof(CapsuleVarName)/sizeof(CapsuleVarName[0]), L"Capsule");
 404   TempVarName = CapsuleVarName + StrLen (CapsuleVarName);
 405   Index = 0;
 406
 407   while (TRUE) {
 408     UnicodeSPrint (TempVarName, 5 * sizeof(CHAR16), L"%04x", Index);
 409
 410     Status = gRT->SetVariable (
 411                     CapsuleVarName,
 412                     &gEfiCapsuleReportGuid,
 413                     EFI_VARIABLE_NON_VOLATILE | EFI_VARIABLE_RUNTIME_ACCESS | EFI_VARIABLE_BOOTSERVICE_ACCESS,
 414                     0,
 415                     (VOID *)NULL
 416                     );
 417     if (EFI_ERROR(Status)) {
 418       //
 419       // There is no capsule variables, quit
 420       //
 421       break;
 422     }
 423
 424     Index++;
 425     if (Index > 0xFFFF) {
 426       break;
 427     }
 428   }
 429
 430   return EFI_SUCCESS;
 431 }
 432
 433 /**
 434   Build Gather list for a list of capsule images.
 435
 436   @param[in]  CapsuleBuffer    An array of pointer to capsule images
 437   @param[in]  FileSize         An array of UINTN to capsule images size
 438   @param[in]  CapsuleNum       The count of capsule images
 439   @param[out] BlockDescriptors The block descriptors for the capsule images
 440
 441   @retval EFI_SUCCESS The block descriptors for the capsule images are constructed.
 442 **/
 443 EFI_STATUS
 444 BuildGatherList (
 445   IN VOID                          **CapsuleBuffer,
 446   IN UINTN                         *FileSize,
 447   IN UINTN                         CapsuleNum,
 448   OUT EFI_CAPSULE_BLOCK_DESCRIPTOR **BlockDescriptors
 449   )
 450 {
 451   EFI_STATUS                    Status;
 452   EFI_CAPSULE_BLOCK_DESCRIPTOR  *BlockDescriptors1;
 453   EFI_CAPSULE_BLOCK_DESCRIPTOR  *BlockDescriptors2;
 454   EFI_CAPSULE_BLOCK_DESCRIPTOR  *BlockDescriptorPre;
 455   EFI_CAPSULE_BLOCK_DESCRIPTOR  *BlockDescriptorsHeader;
 456   EFI_CAPSULE_BLOCK_DESCRIPTOR  *TempBlockPtr;
 457   UINT8                         *TempDataPtr;
 458   UINTN                         SizeLeft;
 459   UINTN                         Size;
 460   INT32                         Count;
 461   INT32                         Number;
 462   UINTN                         Index;
 463
 464   TempBlockPtr           = NULL;
 465   BlockDescriptors1      = NULL;
 466   BlockDescriptors2      = NULL;
 467   BlockDescriptorPre     = NULL;
 468   BlockDescriptorsHeader = NULL;
 469
 470   for (Index = 0; Index < CapsuleNum; Index++) {
 471     //
 472     // Allocate memory for the descriptors.
 473     //
 474     if (NumberOfDescriptors == 1) {
 475       Count = 2;
 476     } else {
 477       Count = (INT32)(NumberOfDescriptors + 2) / 2;
 478     }
 479
 480     Size               = Count * sizeof (EFI_CAPSULE_BLOCK_DESCRIPTOR);
 481     BlockDescriptors1  = AllocateRuntimeZeroPool (Size);
 482     if (BlockDescriptors1 == NULL) {
 483       Print (L"CapsuleApp: failed to allocate memory for descriptors\n");
 484       Status = EFI_OUT_OF_RESOURCES;
 485       goto ERREXIT;
 486     } else {
 487       Print (L"CapsuleApp: creating capsule descriptors at 0x%X\n", (UINTN) BlockDescriptors1);
 488       Print (L"CapsuleApp: capsule data starts          at 0x%X with size 0x%X\n", (UINTN) CapsuleBuffer, FileSize);
 489     }
 490
 491     //
 492     // Record descirptor header
 493     //
 494     if (Index == 0) {
 495       BlockDescriptorsHeader = BlockDescriptors1;
 496     }
 497
 498     if (BlockDescriptorPre != NULL) {
 499       BlockDescriptorPre->Union.ContinuationPointer = (UINTN) BlockDescriptors1;
 500       BlockDescriptorPre->Length = 0;
 501     }
 502
 503     //
 504     // Fill them in
 505     //
 506     TempBlockPtr  = BlockDescriptors1;
 507     TempDataPtr   = CapsuleBuffer[Index];
 508     SizeLeft      = FileSize[Index];
 509     for (Number = 0; (Number < Count - 1) && (SizeLeft != 0); Number++) {
 510       //
 511       // Divide remaining data in half
 512       //
 513       if (NumberOfDescriptors != 1) {
 514         if (SizeLeft == 1) {
 515           Size = 1;
 516         } else {
 517           Size = SizeLeft / 2;
 518         }
 519       } else {
 520         Size = SizeLeft;
 521       }
 522       TempBlockPtr->Union.DataBlock    = (UINTN)TempDataPtr;
 523       TempBlockPtr->Length  = Size;
 524       Print (L"CapsuleApp: capsule block/size              0x%X/0x%X\n", (UINTN) TempDataPtr, Size);
 525       SizeLeft -= Size;
 526       TempDataPtr += Size;
 527       TempBlockPtr++;
 528     }
 529
 530     //
 531     // Allocate the second list, point the first block's last entry to point
 532     // to this one, and fill this one in. Worst case is that the previous
 533     // list only had one element that pointed here, so we need at least two
 534     // elements -- one to point to all the data, another to terminate the list.
 535     //
 536     if ((NumberOfDescriptors != 1) && (SizeLeft != 0)) {
 537       Count = (INT32)(NumberOfDescriptors + 2) - Count;
 538       if (Count == 1) {
 539         Count++;
 540       }
 541
 542       Size              = Count * sizeof (EFI_CAPSULE_BLOCK_DESCRIPTOR);
 543       BlockDescriptors2 = AllocateRuntimeZeroPool (Size);
 544       if (BlockDescriptors2 == NULL) {
 545         Print (L"CapsuleApp: failed to allocate memory for descriptors\n");
 546         Status = EFI_OUT_OF_RESOURCES;
 547         goto ERREXIT;
 548       }
 549
 550       //
 551       // Point the first list's last element to point to this second list.
 552       //
 553       TempBlockPtr->Union.ContinuationPointer   = (UINTN) BlockDescriptors2;
 554
 555       TempBlockPtr->Length  = 0;
 556       TempBlockPtr = BlockDescriptors2;
 557       for (Number = 0; Number < Count - 1; Number++) {
 558         //
 559         // If second-to-last one, then dump rest to this element
 560         //
 561         if (Number == (Count - 2)) {
 562           Size = SizeLeft;
 563         } else {
 564           //
 565           // Divide remaining data in half
 566           //
 567           if (SizeLeft == 1) {
 568             Size = 1;
 569           } else {
 570             Size = SizeLeft / 2;
 571           }
 572         }
 573
 574         TempBlockPtr->Union.DataBlock    = (UINTN)TempDataPtr;
 575         TempBlockPtr->Length  = Size;
 576         Print (L"CapsuleApp: capsule block/size              0x%X/0x%X\n", (UINTN) TempDataPtr, Size);
 577         SizeLeft -= Size;
 578         TempDataPtr += Size;
 579         TempBlockPtr++;
 580         if (SizeLeft == 0) {
 581           break;
 582         }
 583       }
 584     }
 585
 586     BlockDescriptorPre = TempBlockPtr;
 587     BlockDescriptors1  = NULL;
 588   }
 589
 590   //
 591   // Null-terminate.
 592   //
 593   if (TempBlockPtr != NULL) {
 594     TempBlockPtr->Union.ContinuationPointer    = (UINTN)NULL;
 595     TempBlockPtr->Length  = 0;
 596     *BlockDescriptors = BlockDescriptorsHeader;
 597   }
 598
 599   return EFI_SUCCESS;
 600
 601 ERREXIT:
 602   if (BlockDescriptors1 != NULL) {
 603     FreePool(BlockDescriptors1);
 604   }
 605
 606   if (BlockDescriptors2 != NULL) {
 607     FreePool(BlockDescriptors2);
 608   }
 609
 610   return Status;
 611 }
 612
 613 /**
 614   Clear the Gather list for a list of capsule images.
 615
 616   @param[in]  BlockDescriptors The block descriptors for the capsule images
 617   @param[in]  CapsuleNum       The count of capsule images
 618 **/
 619 VOID
 620 CleanGatherList (
 621   IN EFI_CAPSULE_BLOCK_DESCRIPTOR   *BlockDescriptors,
 622   IN UINTN                          CapsuleNum
 623   )
 624 {
 625   EFI_CAPSULE_BLOCK_DESCRIPTOR   *TempBlockPtr;
 626   EFI_CAPSULE_BLOCK_DESCRIPTOR   *TempBlockPtr1;
 627   EFI_CAPSULE_BLOCK_DESCRIPTOR   *TempBlockPtr2;
 628   UINTN                          Index;
 629
 630   if (BlockDescriptors != NULL) {
 631     TempBlockPtr1 = BlockDescriptors;
 632     while (1){
 633       TempBlockPtr = TempBlockPtr1;
 634       for (Index = 0; Index < CapsuleNum; Index++) {
 635         if (TempBlockPtr[Index].Length == 0) {
 636           break;
 637         }
 638       }
 639
 640       if (TempBlockPtr[Index].Union.ContinuationPointer == (UINTN)NULL) {
 641         break;
 642       }
 643
 644       TempBlockPtr2 = (VOID *) ((UINTN) TempBlockPtr->Union.ContinuationPointer);
 645       FreePool(TempBlockPtr1);
 646       TempBlockPtr1 = TempBlockPtr2;
 647     }
 648   }
 649 }
 650
 651 /**
 652   Print APP usage.
 653 **/
 654 VOID
 655 PrintUsage (
 656   VOID
 657   )
 658 {
 659   Print(L"CapsuleApp:  usage\n");
 660   Print(L"  CapsuleApp <Capsule...>\n");
 661   Print(L"  CapsuleApp -S\n");
 662   Print(L"  CapsuleApp -C\n");
 663   Print(L"  CapsuleApp -P\n");
 664   Print(L"  CapsuleApp -E\n");
 665   Print(L"  CapsuleApp -G <BMP> -O <Capsule>\n");
 666   Print(L"  CapsuleApp -N <Capsule> -O <NestedCapsule>\n");
 667   Print(L"  CapsuleApp -D <Capsule>\n");
 668   Print(L"Parameter:\n");
 669   Print(L"  -S:  Dump capsule report variable (EFI_CAPSULE_REPORT_GUID),\n");
 670   Print(L"       which is defined in UEFI specification.\n");
 671   Print(L"  -C:  Clear capsule report variable (EFI_CAPSULE_RPORT_GUID),\n");
 672   Print(L"       which is defined in UEFI specification.\n");
 673   Print(L"  -P:  Dump UEFI FMP protocol info.\n");
 674   Print(L"  -E:  Dump UEFI ESRT table info.\n");
 675   Print(L"  -G:  Convert a BMP file to be a UX capsule,\n");
 676   Print(L"       according to Windows Firmware Update document\n");
 677   Print(L"  -N:  Append a Capsule Header to an existing capsule image,\n");
 678   Print(L"       according to Windows Firmware Update document\n");
 679   Print(L"  -O:  Output new Capsule file name\n");
 680   Print(L"  -D:  Dump Capsule image header information and FMP header information,\n");
 681   Print(L"       if it is an FMP capsule.\n");
 682 }
 683
 684 /**
 685   Update Capsule image.
 686
 687   @param[in]  ImageHandle     The image handle.
 688   @param[in]  SystemTable     The system table.
 689
 690   @retval EFI_SUCCESS            Command completed successfully.
 691   @retval EFI_INVALID_PARAMETER  Command usage error.
 692   @retval EFI_NOT_FOUND          The input file can't be found.
 693 **/
 694 EFI_STATUS
 695 EFIAPI
 696 UefiMain (
 697   IN EFI_HANDLE        ImageHandle,
 698   IN EFI_SYSTEM_TABLE  *SystemTable
 699   )
 700 {
 701   EFI_STATUS                    Status;
 702   UINTN                         FileSize[MAX_CAPSULE_NUM];
 703   VOID                          *CapsuleBuffer[MAX_CAPSULE_NUM];
 704   EFI_CAPSULE_BLOCK_DESCRIPTOR  *BlockDescriptors;
 705   EFI_CAPSULE_HEADER             *CapsuleHeaderArray[MAX_CAPSULE_NUM + 1];
 706   UINT64                         MaxCapsuleSize;
 707   EFI_RESET_TYPE                 ResetType;
 708   BOOLEAN                        NeedReset;
 709   CHAR16                         *CapsuleName;
 710   UINTN                          CapsuleNum;
 711   UINTN                          Index;
 712
 713   Status = GetArg();
 714   if (EFI_ERROR(Status)) {
 715     Print(L"Please use UEFI SHELL to run this application!\n", Status);
 716     return Status;
 717   }
 718   if (Argc < 2) {
 719     PrintUsage();
 720     return EFI_INVALID_PARAMETER;
 721   }
 722   if (StrCmp(Argv[1], L"-D") == 0) {
 723     Status = DumpCapsule(Argv[2]);
 724     return Status;
 725   }
 726   if (StrCmp(Argv[1], L"-G") == 0) {
 727     Status = CreateBmpFmp();
 728     return Status;
 729   }
 730   if (StrCmp(Argv[1], L"-N") == 0) {
 731     Status = CreateNestedFmp();
 732     return Status;
 733   }
 734   if (StrCmp(Argv[1], L"-S") == 0) {
 735     Status = DmpCapsuleStatusVariable();
 736     return EFI_SUCCESS;
 737   }
 738   if (StrCmp(Argv[1], L"-C") == 0) {
 739     Status = ClearCapsuleStatusVariable();
 740     return Status;
 741   }
 742   if (StrCmp(Argv[1], L"-P") == 0) {
 743     DumpFmpData();
 744     return EFI_SUCCESS;
 745   }
 746   if (StrCmp(Argv[1], L"-E") == 0) {
 747     DumpEsrtData();
 748     return EFI_SUCCESS;
 749   }
 750   CapsuleFirstIndex = 1;
 751   CapsuleLastIndex = Argc - 1;
 752   CapsuleNum = CapsuleLastIndex - CapsuleFirstIndex + 1;
 753
 754   if (CapsuleFirstIndex > CapsuleLastIndex) {
 755     Print(L"CapsuleApp: NO capsule image.\n");
 756     return EFI_UNSUPPORTED;
 757   }
 758   if (CapsuleNum > MAX_CAPSULE_NUM) {
 759     Print(L"CapsuleApp: Too many capsule images.\n");
 760     return EFI_UNSUPPORTED;
 761   }
 762
 763   ZeroMem(&CapsuleBuffer, sizeof(CapsuleBuffer));
 764   ZeroMem(&FileSize, sizeof(FileSize));
 765   BlockDescriptors = NULL;
 766
 767   for (Index = 0; Index < CapsuleNum; Index++) {
 768     CapsuleName = Argv[CapsuleFirstIndex + Index];
 769     Status = ReadFileToBuffer(CapsuleName, &FileSize[Index], &CapsuleBuffer[Index]);
 770     if (EFI_ERROR(Status)) {
 771       Print(L"CapsuleApp: capsule image (%s) is not found.\n", CapsuleName);
 772       goto Done;
 773     }
 774   }
 775
 776   //
 777   // Every capsule use 2 descriptor 1 for data 1 for end
 778   //
 779   Status = BuildGatherList(CapsuleBuffer, FileSize, CapsuleNum, &BlockDescriptors);
 780   if (EFI_ERROR(Status)) {
 781     goto Done;
 782   }
 783
 784   //
 785   // Call the runtime service capsule.
 786   //
 787   NeedReset = FALSE;
 788   for (Index = 0; Index < CapsuleNum; Index++) {
 789     CapsuleHeaderArray[Index] = (EFI_CAPSULE_HEADER *) CapsuleBuffer[Index];
 790     if ((CapsuleHeaderArray[Index]->Flags & CAPSULE_FLAGS_PERSIST_ACROSS_RESET) != 0) {
 791       NeedReset = TRUE;
 792     }
 793   }
 794   CapsuleHeaderArray[CapsuleNum] = NULL;
 795
 796   //
 797   // Inquire platform capability of UpdateCapsule.
 798   //
 799   Status = gRT->QueryCapsuleCapabilities (CapsuleHeaderArray, CapsuleNum, &MaxCapsuleSize, &ResetType);
 800   if (EFI_ERROR(Status)) {
 801     Print (L"CapsuleApp: failed to query capsule capability - %r\n", Status);
 802     goto Done;
 803   }
 804
 805   for (Index = 0; Index < CapsuleNum; Index++) {
 806     if (FileSize[Index] > MaxCapsuleSize) {
 807       Print (L"CapsuleApp: capsule is too large to update, %ld is allowed\n", MaxCapsuleSize);
 808       Status = EFI_UNSUPPORTED;
 809       goto Done;
 810     }
 811   }
 812
 813   //
 814   // Check whether the input capsule image has the flag of persist across system reset.
 815   //
 816   if (NeedReset) {
 817     Status = gRT->UpdateCapsule(CapsuleHeaderArray,CapsuleNum,(UINTN) BlockDescriptors);
 818     if (Status != EFI_SUCCESS) {
 819       Print (L"CapsuleApp: failed to update capsule - %r\n", Status);
 820       goto Done;
 821     }
 822     //
 823     // For capsule who has reset flag, after calling UpdateCapsule service,triger a
 824     // system reset to process capsule persist across a system reset.
 825     //
 826     gRT->ResetSystem (ResetType, EFI_SUCCESS, 0, NULL);
 827   } else {
 828     //
 829     // For capsule who has no reset flag, only call UpdateCapsule Service without a
 830     // system reset. The service will process the capsule immediately.
 831     //
 832     Status = gRT->UpdateCapsule (CapsuleHeaderArray,CapsuleNum,(UINTN) BlockDescriptors);
 833     if (Status != EFI_SUCCESS) {
 834       Print (L"CapsuleApp: failed to update capsule - %r\n", Status);
 835     }
 836   }
 837
 838   Status = EFI_SUCCESS;
 839
 840 Done:
 841   for (Index = 0; Index < CapsuleNum; Index++) {
 842     if (CapsuleBuffer[Index] != NULL) {
 843       FreePool (CapsuleBuffer[Index]);
 844     }
 845   }
 846
 847   CleanGatherList(BlockDescriptors, CapsuleNum);
 848
 849   return Status;
 850 }