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