MdeModulePkg/Application/CapsuleApp/CapsuleApp.c

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