/** @file

  Copyright (c) 2013-2014, ARM Ltd. All rights reserved.<BR>

  This program and the accompanying materials
  are licensed and made available under the terms and conditions of the BSD License
  which accompanies this distribution.  The full text of the license may be found at
  http://opensource.org/licenses/bsd-license.php

  THE PROGRAM IS DISTRIBUTED UNDER THE BSD LICENSE ON AN "AS IS" BASIS,
  WITHOUT WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.

**/

#include <Protocol/AndroidFastbootTransport.h>
#include <Protocol/AndroidFastbootPlatform.h>
#include <Protocol/SimpleTextOut.h>
#include <Protocol/SimpleTextIn.h>

#include <Library/AbootimgLib.h>
#include <Library/BaseMemoryLib.h>
#include <Library/PcdLib.h>
#include <Library/PrintLib.h>
#include <Library/UefiApplicationEntryPoint.h>
#include <Library/UefiBootServicesTableLib.h>
#include <Library/UefiRuntimeServicesTableLib.h>

#define ANDROID_FASTBOOT_VERSION    "0.7"

#define SPARSE_HEADER_MAGIC         0xED26FF3A
#define CHUNK_TYPE_RAW              0xCAC1
#define CHUNK_TYPE_FILL             0xCAC2
#define CHUNK_TYPE_DONT_CARE        0xCAC3
#define CHUNK_TYPE_CRC32            0xCAC4

#define FILL_BUF_SIZE               (16 * 1024 * 1024)
#define SPARSE_BLOCK_SIZE           4096

#define IS_DEVICE_PATH_NODE(node,type,subtype) (((node)->Type == (type)) && ((node)->SubType == (subtype)))

#define ALIGN(x, a)        (((x) + ((a) - 1)) & ~((a) - 1))

typedef struct _SPARSE_HEADER {
  UINT32       Magic;
  UINT16       MajorVersion;
  UINT16       MinorVersion;
  UINT16       FileHeaderSize;
  UINT16       ChunkHeaderSize;
  UINT32       BlockSize;
  UINT32       TotalBlocks;
  UINT32       TotalChunks;
  UINT32       ImageChecksum;
} SPARSE_HEADER;

typedef struct _CHUNK_HEADER {
  UINT16       ChunkType;
  UINT16       Reserved1;
  UINT32       ChunkSize;
  UINT32       TotalSize;
} CHUNK_HEADER;

/*
 * UEFI Application using the FASTBOOT_TRANSPORT_PROTOCOL and
 * FASTBOOT_PLATFORM_PROTOCOL to implement the Android Fastboot protocol.
 */

STATIC FASTBOOT_TRANSPORT_PROTOCOL *mTransport;
STATIC FASTBOOT_PLATFORM_PROTOCOL  *mPlatform;

STATIC EFI_SIMPLE_TEXT_OUTPUT_PROTOCOL *mTextOut;

typedef enum {
  ExpectCmdState,
  ExpectDataState,
  FastbootStateMax
} ANDROID_FASTBOOT_STATE;

STATIC ANDROID_FASTBOOT_STATE mState = ExpectCmdState;

// When in ExpectDataState, the number of bytes of data to expect:
STATIC UINT64 mNumDataBytes;
// .. and the number of bytes so far received this data phase
STATIC UINT64 mBytesReceivedSoFar;
// .. and the buffer to save data into
STATIC UINT8 *mDataBuffer = NULL;

// Event notify functions, from which gBS->Exit shouldn't be called, can signal
// this event when the application should exit
STATIC EFI_EVENT mFinishedEvent;

STATIC EFI_EVENT mFatalSendErrorEvent;

// This macro uses sizeof - only use it on arrays (i.e. string literals)
#define SEND_LITERAL(Str) mTransport->Send (                  \
                                        sizeof (Str) - 1,     \
                                        Str,                  \
                                        &mFatalSendErrorEvent \
                                        )
#define MATCH_CMD_LITERAL(Cmd, Buf) !AsciiStrnCmp (Cmd, Buf, sizeof (Cmd) - 1)

#define IS_LOWERCASE_ASCII(Char) (Char >= 'a' && Char <= 'z')

#define FASTBOOT_STRING_MAX_LENGTH  256
#define FASTBOOT_COMMAND_MAX_LENGTH 64

STATIC
VOID
HandleGetVar (
  IN CHAR8 *CmdArg
  )
{
  CHAR8      Response[FASTBOOT_COMMAND_MAX_LENGTH + 1] = "OKAY";
  EFI_STATUS Status;

  // Respond to getvar:version with 0.4 (version of Fastboot protocol)
  if (!AsciiStrnCmp ("version", CmdArg, sizeof ("version") - 1 )) {
    SEND_LITERAL ("OKAY" ANDROID_FASTBOOT_VERSION);
  } else {
    // All other variables are assumed to be platform specific
    Status = mPlatform->GetVar (CmdArg, Response + 4);
    if (EFI_ERROR (Status)) {
      SEND_LITERAL ("FAILSomething went wrong when looking up the variable");
    } else {
      mTransport->Send (AsciiStrLen (Response), Response, &mFatalSendErrorEvent);
    }
  }
}

STATIC
VOID
HandleDownload (
  IN CHAR8 *NumBytesString
  )
{
  CHAR8       Response[13];
  CHAR16      OutputString[FASTBOOT_STRING_MAX_LENGTH];

  // Argument is 8-character ASCII string hex representation of number of bytes
  // that will be sent in the data phase.
  // Response is "DATA" + that same 8-character string.

  // Replace any previously downloaded data
  if (mDataBuffer != NULL) {
    FreePool (mDataBuffer);
    mDataBuffer = NULL;
  }

  // Parse out number of data bytes to expect
  mNumDataBytes = AsciiStrHexToUint64 (NumBytesString);
  if (mNumDataBytes == 0) {
    mTextOut->OutputString (mTextOut, L"ERROR: Fail to get the number of bytes to download.\r\n");
    SEND_LITERAL ("FAILFailed to get the number of bytes to download");
    return;
  }

  UnicodeSPrint (OutputString, sizeof (OutputString), L"Downloading %d bytes\r\n", mNumDataBytes);
  mTextOut->OutputString (mTextOut, OutputString);

  mDataBuffer = AllocatePool (mNumDataBytes);
  if (mDataBuffer == NULL) {
    SEND_LITERAL ("FAILNot enough memory");
  } else {
    ZeroMem (Response, sizeof Response);
    if (mTransport->RequestReceive) {
      mTransport->RequestReceive (mNumDataBytes);
    }
    AsciiSPrint (Response, sizeof Response, "DATA%x",
      (UINT32)mNumDataBytes);
    mTransport->Send (sizeof Response - 1, Response, &mFatalSendErrorEvent);

    mState = ExpectDataState;
    mBytesReceivedSoFar = 0;
  }
}

STATIC
EFI_STATUS
FlashSparseImage (
  IN CHAR8         *PartitionName,
  IN SPARSE_HEADER *SparseHeader
  )
{
  EFI_STATUS        Status = EFI_SUCCESS;
  UINTN             Chunk, Offset = 0, Left, Count, FillBufSize;
  VOID             *Image;
  CHUNK_HEADER     *ChunkHeader;
  VOID             *FillBuf;
  CHAR16            OutputString[FASTBOOT_STRING_MAX_LENGTH];

  Image = (VOID *)SparseHeader;
  Image += SparseHeader->FileHeaderSize;

  // allocate the fill buf with dynamic size
  FillBufSize = FILL_BUF_SIZE;
  while (FillBufSize >= SPARSE_BLOCK_SIZE) {
    FillBuf = AllocatePool (FillBufSize);
    if (FillBuf == NULL) {
      FillBufSize = FillBufSize >> 1;
    } else {
      break;
    }
  };
  if (FillBufSize < SPARSE_BLOCK_SIZE) {
    UnicodeSPrint (
      OutputString,
      sizeof (OutputString),
      L"Fail to allocate the fill buffer\n"
      );
    mTextOut->OutputString (mTextOut, OutputString);
    return EFI_BUFFER_TOO_SMALL;
  }

  for (Chunk = 0; Chunk < SparseHeader->TotalChunks; Chunk++) {
    ChunkHeader = (CHUNK_HEADER *)Image;
    DEBUG ((DEBUG_INFO, "Chunk #%d - Type: 0x%x Size: %d TotalSize: %d Offset %d\n",
            (Chunk+1), ChunkHeader->ChunkType, ChunkHeader->ChunkSize,
            ChunkHeader->TotalSize, Offset));
    Image += sizeof (CHUNK_HEADER);
    switch (ChunkHeader->ChunkType) {
    case CHUNK_TYPE_RAW:
      Status = mPlatform->FlashPartitionEx (
                            PartitionName,
                            Offset,
                            ChunkHeader->ChunkSize * SparseHeader->BlockSize,
                            Image
                            );
      if (EFI_ERROR (Status)) {
        return Status;
      }
      Image += (UINTN)ChunkHeader->ChunkSize * SparseHeader->BlockSize;
      Offset += (UINTN)ChunkHeader->ChunkSize * SparseHeader->BlockSize;
      break;
    case CHUNK_TYPE_FILL:
      Left = (UINTN)ChunkHeader->ChunkSize * SparseHeader->BlockSize;
      while (Left > 0) {
        if (Left > FILL_BUF_SIZE) {
          Count = FILL_BUF_SIZE;
        } else {
          Count = Left;
        }
        SetMem32 (FillBuf, Count, *(UINT32 *)Image);
        Status = mPlatform->FlashPartitionEx (
                              PartitionName,
                              Offset,
                              Count,
                              FillBuf
                              );
        if (EFI_ERROR (Status)) {
          return Status;
        }
        Offset += Count;
        Left = Left - Count;
      }
      Image += sizeof (UINT32);
      break;
    case CHUNK_TYPE_DONT_CARE:
      Offset += (UINTN)ChunkHeader->ChunkSize * SparseHeader->BlockSize;
      break;
    default:
      UnicodeSPrint (
        OutputString,
        sizeof (OutputString),
        L"Unsupported Chunk Type:0x%x\n",
        ChunkHeader->ChunkType
        );
      mTextOut->OutputString (mTextOut, OutputString);
      break;
    }
  }
  FreePool ((VOID *)FillBuf);
  return Status;
}

STATIC
VOID
HandleFlash (
  IN CHAR8 *PartitionName
  )
{
  EFI_STATUS        Status;
  CHAR16            OutputString[FASTBOOT_STRING_MAX_LENGTH];
  SPARSE_HEADER    *SparseHeader;

  // Build output string
  UnicodeSPrint (OutputString, sizeof (OutputString), L"Flashing partition %a\r\n", PartitionName);
  mTextOut->OutputString (mTextOut, OutputString);

  if (mDataBuffer == NULL) {
    // Doesn't look like we were sent any data
    SEND_LITERAL ("FAILNo data to flash");
    return;
  }

  SparseHeader = (SPARSE_HEADER *)mDataBuffer;
  if (SparseHeader->Magic == SPARSE_HEADER_MAGIC) {
    DEBUG ((DEBUG_INFO, "Sparse Magic: 0x%x Major: %d Minor: %d fhs: %d chs: %d bs: %d tbs: %d tcs: %d checksum: %d \n",
                SparseHeader->Magic, SparseHeader->MajorVersion, SparseHeader->MinorVersion,  SparseHeader->FileHeaderSize,
                SparseHeader->ChunkHeaderSize, SparseHeader->BlockSize, SparseHeader->TotalBlocks,
                SparseHeader->TotalChunks, SparseHeader->ImageChecksum));
    if (SparseHeader->MajorVersion != 1) {
        DEBUG ((DEBUG_ERROR, "Sparse image version %d.%d not supported.\n",
                    SparseHeader->MajorVersion, SparseHeader->MinorVersion));
        return;
    }
    Status = FlashSparseImage (PartitionName, SparseHeader);
  } else {
    Status = mPlatform->FlashPartition (
                          PartitionName,
                          mNumDataBytes,
                          mDataBuffer
                          );
  }
  switch (Status) {
  case EFI_SUCCESS:
    mTextOut->OutputString (mTextOut, L"Done.\r\n");
    SEND_LITERAL ("OKAY");
    break;
  case EFI_NOT_FOUND:
    SEND_LITERAL ("FAILNo such partition.");
    mTextOut->OutputString (mTextOut, L"No such partition.\r\n");
    break;
  default:
    SEND_LITERAL ("FAILError flashing partition.");
    mTextOut->OutputString (mTextOut, L"Error flashing partition.\r\n");
    DEBUG ((EFI_D_ERROR, "Couldn't flash image:\n"));
    break;
  }
}

STATIC
VOID
HandleErase (
  IN CHAR8 *PartitionName
  )
{
  EFI_STATUS  Status;
  CHAR16      OutputString[FASTBOOT_STRING_MAX_LENGTH];

  // Build output string
  UnicodeSPrint (OutputString, sizeof (OutputString), L"Erasing partition %a\r\n", PartitionName);
  mTextOut->OutputString (mTextOut, OutputString);

  Status = mPlatform->ErasePartition (PartitionName);
  if (EFI_ERROR (Status)) {
    SEND_LITERAL ("FAILCheck device console.");
    DEBUG ((EFI_D_ERROR, "Couldn't erase image:  %r\n", Status));
  } else {
    SEND_LITERAL ("OKAY");
  }
}

STATIC
VOID
HandleBoot (
  VOID
  )
{
  CHAR16     *BootPathStr;

  mTextOut->OutputString (mTextOut, L"Booting downloaded image\r\n");

  if (mDataBuffer == NULL) {
    // Doesn't look like we were sent any data
    SEND_LITERAL ("FAILNo image in memory");
    return;
  }

  // We don't really have any choice but to report success, because once we
  // boot we lose control of the system.
  SEND_LITERAL ("OKAY");

  BootPathStr = (CHAR16 *)PcdGetPtr (PcdAndroidBootDevicePath);
  AbootimgBootRam (mDataBuffer, mNumDataBytes, BootPathStr, NULL);
  // We shouldn't get here
}

STATIC
VOID
HandleOemCommand (
  IN CHAR8 *Command
  )
{
  EFI_STATUS  Status;

  Status = mPlatform->DoOemCommand (Command);
  if (Status == EFI_NOT_FOUND) {
    SEND_LITERAL ("FAILOEM Command not recognised.");
  } else if (Status == EFI_DEVICE_ERROR) {
    SEND_LITERAL ("FAILError while executing command");
  } else if (EFI_ERROR (Status)) {
    SEND_LITERAL ("FAIL");
  } else {
    SEND_LITERAL ("OKAY");
  }
}

STATIC
VOID
AcceptCmd (
  IN        UINTN  Size,
  IN  CONST CHAR8 *Data
  )
{
  EFI_STATUS  Status;
  CHAR8       Command[FASTBOOT_COMMAND_MAX_LENGTH + 1];

  // Max command size is 64 bytes
  if (Size > FASTBOOT_COMMAND_MAX_LENGTH) {
    SEND_LITERAL ("FAILCommand too large");
    return;
  }

  // Commands aren't null-terminated. Let's get a null-terminated version.
  AsciiStrnCpyS (Command, sizeof Command, Data, Size);

  // Parse command
  if (MATCH_CMD_LITERAL ("getvar", Command)) {
    HandleGetVar (Command + sizeof ("getvar"));
  } else if (MATCH_CMD_LITERAL ("download", Command)) {
    HandleDownload (Command + sizeof ("download"));
  } else if (MATCH_CMD_LITERAL ("verify", Command)) {
    SEND_LITERAL ("FAILNot supported");
  } else if (MATCH_CMD_LITERAL ("flash", Command)) {
    HandleFlash (Command + sizeof ("flash"));
  } else if (MATCH_CMD_LITERAL ("erase", Command)) {
    HandleErase (Command + sizeof ("erase"));
  } else if (MATCH_CMD_LITERAL ("boot", Command)) {
    HandleBoot ();
  } else if (MATCH_CMD_LITERAL ("continue", Command)) {
    SEND_LITERAL ("OKAY");
    mTextOut->OutputString (mTextOut, L"Received 'continue' command. Exiting Fastboot mode\r\n");

    gBS->SignalEvent (mFinishedEvent);
  } else if (MATCH_CMD_LITERAL ("reboot", Command)) {
    if (MATCH_CMD_LITERAL ("reboot-bootloader", Command)) {
      Status = mPlatform->DoOemCommand ("reboot-bootloader");
      if (EFI_ERROR (Status)) {
        SEND_LITERAL ("INFOreboot-bootloader not supported, rebooting normally.");
      }
    }
    SEND_LITERAL ("OKAY");
    gRT->ResetSystem (EfiResetCold, EFI_SUCCESS, 0, NULL);

    // Shouldn't get here
    DEBUG ((EFI_D_ERROR, "Fastboot: gRT->ResetSystem didn't work\n"));
  } else if (MATCH_CMD_LITERAL ("powerdown", Command)) {
    SEND_LITERAL ("OKAY");
    gRT->ResetSystem (EfiResetShutdown, EFI_SUCCESS, 0, NULL);

    // Shouldn't get here
    DEBUG ((EFI_D_ERROR, "Fastboot: gRT->ResetSystem didn't work\n"));
  } else if (MATCH_CMD_LITERAL ("oem", Command)) {
    // The "oem" command isn't in the specification, but it was observed in the
    // wild, followed by a space, followed by the actual command.
    HandleOemCommand (Command + sizeof ("oem"));
  } else if (IS_LOWERCASE_ASCII (Command[0])) {
    // Commands starting with lowercase ASCII characters are reserved for the
    // Fastboot protocol. If we don't recognise it, it's probably the future
    // and there are new commmands in the protocol.
    // (By the way, the "oem" command mentioned above makes this reservation
    //  redundant, but we handle it here to be spec-compliant)
    SEND_LITERAL ("FAILCommand not recognised. Check Fastboot version.");
  } else {
    HandleOemCommand (Command);
  }
}

STATIC
VOID
AcceptData (
  IN  UINTN  Size,
  IN  VOID  *Data
  )
{
  UINT32 RemainingBytes = mNumDataBytes - mBytesReceivedSoFar;
  CHAR16 OutputString[FASTBOOT_STRING_MAX_LENGTH];
  STATIC UINTN Count = 0;

  // Protocol doesn't say anything about sending extra data so just ignore it.
  if (Size > RemainingBytes) {
    Size = RemainingBytes;
  }

  CopyMem (&mDataBuffer[mBytesReceivedSoFar], Data, Size);

  mBytesReceivedSoFar += Size;

  // Show download progress. Don't do it for every packet  as outputting text
  // might be time consuming - do it on the last packet and on every 32nd packet
  if ((Count++ % 32) == 0 || Size == RemainingBytes) {
    // (Note no newline in format string - it will overwrite the line each time)
    UnicodeSPrint (
      OutputString,
      sizeof (OutputString),
      L"\r%8d / %8d bytes downloaded (%d%%)",
      mBytesReceivedSoFar,
      mNumDataBytes,
      (mBytesReceivedSoFar * 100) / mNumDataBytes // percentage
      );
    mTextOut->OutputString (mTextOut, OutputString);
  }

  if (mBytesReceivedSoFar == mNumDataBytes) {
    // Download finished.

    mTextOut->OutputString (mTextOut, L"\r\n");
    SEND_LITERAL ("OKAY");
    mState = ExpectCmdState;
  }
}

/*
  This is the NotifyFunction passed to CreateEvent in the FastbootAppEntryPoint
  It will be called by the UEFI event framework when the transport protocol
  implementation signals that data has been received from the Fastboot host.
  The parameters are ignored.
*/
STATIC
VOID
DataReady (
  IN EFI_EVENT  Event,
  IN VOID      *Context
  )
{
  UINTN       Size;
  VOID       *Data;
  EFI_STATUS  Status;

  do {
    // Indicate lower layer driver that how much bytes are expected.
    if (mState == ExpectDataState) {
      Size = mNumDataBytes;
    } else {
      Size = 0;
    }
    Status = mTransport->Receive (&Size, &Data);
    if (!EFI_ERROR (Status)) {
      if (mState == ExpectCmdState) {
        AcceptCmd (Size, (CHAR8 *) Data);
      } else if (mState == ExpectDataState) {
        AcceptData (Size, Data);
      } else {
        ASSERT (FALSE);
      }
      FreePool (Data);
    }
  } while (!EFI_ERROR (Status));

  // Quit if there was a fatal error
  if (Status != EFI_NOT_READY) {
    ASSERT (Status == EFI_DEVICE_ERROR);
    // (Put a newline at the beginning as we are probably in the data phase,
    //  so the download progress line, with no '\n' is probably on the console)
    mTextOut->OutputString (mTextOut, L"\r\nFatal error receiving data. Exiting.\r\n");
    gBS->SignalEvent (mFinishedEvent);
  }
}

/*
  Event notify for a fatal error in transmission.
*/
STATIC
VOID
FatalErrorNotify (
  IN EFI_EVENT  Event,
  IN VOID      *Context
  )
{
  mTextOut->OutputString (mTextOut, L"Fatal error sending command response. Exiting.\r\n");
  gBS->SignalEvent (mFinishedEvent);
}

EFI_STATUS
EFIAPI
FastbootAppEntryPoint (
  IN EFI_HANDLE                            ImageHandle,
  IN EFI_SYSTEM_TABLE                      *SystemTable
  )
{
  EFI_STATUS                      Status;
  EFI_EVENT                       ReceiveEvent;
  EFI_EVENT                       WaitEventArray[2];
  UINTN                           EventIndex;
  EFI_SIMPLE_TEXT_INPUT_PROTOCOL *TextIn;
  EFI_INPUT_KEY                   Key;

  mDataBuffer = NULL;

  Status = gBS->LocateProtocol (
    &gAndroidFastbootTransportProtocolGuid,
    NULL,
    (VOID **) &mTransport
    );
  if (EFI_ERROR (Status)) {
    DEBUG ((EFI_D_ERROR, "Fastboot: Couldn't open Fastboot Transport Protocol: %r\n", Status));
    return Status;
  }

  Status = gBS->LocateProtocol (&gAndroidFastbootPlatformProtocolGuid, NULL, (VOID **) &mPlatform);
  if (EFI_ERROR (Status)) {
    DEBUG ((EFI_D_ERROR, "Fastboot: Couldn't open Fastboot Platform Protocol: %r\n", Status));
    return Status;
  }

  Status = mPlatform->Init ();
  if (EFI_ERROR (Status)) {
    DEBUG ((EFI_D_ERROR, "Fastboot: Couldn't initialise Fastboot Platform Protocol: %r\n", Status));
    return Status;
  }

  Status = gBS->LocateProtocol (&gEfiSimpleTextOutProtocolGuid, NULL, (VOID **) &mTextOut);
  if (EFI_ERROR (Status)) {
    DEBUG ((EFI_D_ERROR,
      "Fastboot: Couldn't open Text Output Protocol: %r\n", Status
      ));
    return Status;
  }

  Status = gBS->LocateProtocol (&gEfiSimpleTextInProtocolGuid, NULL, (VOID **) &TextIn);
  if (EFI_ERROR (Status)) {
    DEBUG ((EFI_D_ERROR, "Fastboot: Couldn't open Text Input Protocol: %r\n", Status));
    return Status;
  }

  // Disable watchdog
  Status = gBS->SetWatchdogTimer (0, 0x10000, 0, NULL);
  if (EFI_ERROR (Status)) {
    DEBUG ((EFI_D_ERROR, "Fastboot: Couldn't disable watchdog timer: %r\n", Status));
  }

  // Create event for receipt of data from the host
  Status = gBS->CreateEvent (
                  EVT_NOTIFY_SIGNAL,
                  TPL_CALLBACK,
                  DataReady,
                  NULL,
                  &ReceiveEvent
                  );
  ASSERT_EFI_ERROR (Status);

  // Create event for exiting application when "continue" command is received
  Status = gBS->CreateEvent (0, TPL_CALLBACK, NULL, NULL, &mFinishedEvent);
  ASSERT_EFI_ERROR (Status);

  // Create event to pass to FASTBOOT_TRANSPORT_PROTOCOL.Send, signalling a
  // fatal error
  Status = gBS->CreateEvent (
                 EVT_NOTIFY_SIGNAL,
                 TPL_CALLBACK,
                 FatalErrorNotify,
                 NULL,
                 &mFatalSendErrorEvent
                 );
  ASSERT_EFI_ERROR (Status);


  // Start listening for data
  Status = mTransport->Start (
    ReceiveEvent
    );
  if (EFI_ERROR (Status)) {
    DEBUG ((EFI_D_ERROR, "Fastboot: Couldn't start transport: %r\n", Status));
    return Status;
  }

  // Talk to the user
  mTextOut->OutputString (mTextOut,
      L"Android Fastboot mode - version " ANDROID_FASTBOOT_VERSION ".\r\n");
  mTextOut->OutputString (mTextOut, L"Press RETURN or SPACE key to quit.\r\n");

  // Quit when the user presses any key, or mFinishedEvent is signalled
  WaitEventArray[0] = mFinishedEvent;
  WaitEventArray[1] = TextIn->WaitForKey;
  while (1) {
    gBS->WaitForEvent (2, WaitEventArray, &EventIndex);
    if (EventIndex == 0) {
      break;
    }
    Status = TextIn->ReadKeyStroke (gST->ConIn, &Key);
    if (Key.ScanCode == SCAN_NULL) {
      if ((Key.UnicodeChar == CHAR_CARRIAGE_RETURN) || (Key.UnicodeChar == ' ')) {
        break;
      }
    }
  }

  mTransport->Stop ();
  if (EFI_ERROR (Status)) {
    DEBUG ((EFI_D_ERROR, "Warning: Fastboot Transport Stop: %r\n", Status));
  }
  mPlatform->UnInit ();

  return EFI_SUCCESS;
}