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