MdeModulePkg/Application/CapsuleApp/CapsuleApp.c

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