/** @file
#
#  Copyright (c) 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/Dhcp4.h>
#include <Protocol/Tcp4.h>
#include <Protocol/ServiceBinding.h>
#include <Protocol/SimpleTextOut.h>

#include <Library/BaseLib.h>
#include <Library/BaseMemoryLib.h>
#include <Library/DebugLib.h>
#include <Library/MemoryAllocationLib.h>
#include <Library/PrintLib.h>
#include <Library/UefiBootServicesTableLib.h>
#include <Library/UefiDriverEntryPoint.h>
#include <Library/UefiRuntimeServicesTableLib.h>

#define IP4_ADDR_TO_STRING(IpAddr, IpAddrString) UnicodeSPrint (       \
                                                   IpAddrString,       \
                                                   16 * 2,             \
                                                   L"%d.%d.%d.%d",     \
                                                   IpAddr.Addr[0],     \
                                                   IpAddr.Addr[1],     \
                                                   IpAddr.Addr[2],     \
                                                   IpAddr.Addr[3]      \
                                                   );

// Fastboot says max packet size is 512, but FASTBOOT_TRANSPORT_PROTOCOL
// doesn't place a limit on the size of buffers returned by Receive.
// (This isn't actually a packet size - it's just the size of the buffers we
//  pass to the TCP driver to fill with received data.)
// We can achieve much better performance by doing this in larger chunks.
#define RX_FRAGMENT_SIZE 2048

STATIC EFI_SIMPLE_TEXT_OUTPUT_PROTOCOL *mTextOut;

STATIC EFI_TCP4_PROTOCOL *mTcpConnection;
STATIC EFI_TCP4_PROTOCOL *mTcpListener;

STATIC EFI_EVENT mReceiveEvent;

STATIC EFI_SERVICE_BINDING_PROTOCOL *mTcpServiceBinding;
STATIC EFI_HANDLE                    mTcpHandle = NULL;

// We only ever use one IO token for receive and one for transmit. To save
// repeatedly allocating and freeing, just allocate statically and re-use.
#define NUM_RX_TOKENS 16
#define TOKEN_NEXT(Index) (((Index) + 1) % NUM_RX_TOKENS)

STATIC UINTN                     mNextSubmitIndex;
STATIC UINTN                     mNextReceiveIndex;
STATIC EFI_TCP4_IO_TOKEN         mReceiveToken[NUM_RX_TOKENS];
STATIC EFI_TCP4_RECEIVE_DATA     mRxData[NUM_RX_TOKENS];
STATIC EFI_TCP4_IO_TOKEN         mTransmitToken;
STATIC EFI_TCP4_TRANSMIT_DATA    mTxData;
// We also reuse the accept token
STATIC EFI_TCP4_LISTEN_TOKEN     mAcceptToken;
// .. and the close token
STATIC EFI_TCP4_CLOSE_TOKEN      mCloseToken;

// List type for queued received packets
typedef struct _FASTBOOT_TCP_PACKET_LIST {
  LIST_ENTRY  Link;
  VOID       *Buffer;
  UINTN       BufferSize;
} FASTBOOT_TCP_PACKET_LIST;

STATIC LIST_ENTRY mPacketListHead;

STATIC
VOID
EFIAPI
DataReceived (
  IN EFI_EVENT Event,
  IN VOID     *Context
  );

/*
  Helper function to set up a receive IO token and call Tcp->Receive
*/
STATIC
EFI_STATUS
SubmitRecieveToken (
  VOID
  )
{
  EFI_STATUS             Status;
  VOID                  *FragmentBuffer;

  Status = EFI_SUCCESS;

  FragmentBuffer = AllocatePool (RX_FRAGMENT_SIZE);
  ASSERT (FragmentBuffer != NULL);
  if (FragmentBuffer == NULL) {
    DEBUG ((EFI_D_ERROR, "TCP Fastboot out of resources"));
    return EFI_OUT_OF_RESOURCES;
  }

  mRxData[mNextSubmitIndex].DataLength = RX_FRAGMENT_SIZE;
  mRxData[mNextSubmitIndex].FragmentTable[0].FragmentLength = RX_FRAGMENT_SIZE;
  mRxData[mNextSubmitIndex].FragmentTable[0].FragmentBuffer = FragmentBuffer;

  Status = mTcpConnection->Receive (mTcpConnection, &mReceiveToken[mNextSubmitIndex]);
   if (EFI_ERROR (Status)) {
    DEBUG ((EFI_D_ERROR, "TCP Receive: %r\n", Status));
    FreePool (FragmentBuffer);
  }

  mNextSubmitIndex = TOKEN_NEXT (mNextSubmitIndex);
  return Status;
}

/*
  Event notify function for when we have closed our TCP connection.
  We can now start listening for another connection.
*/
STATIC
VOID
ConnectionClosed (
  IN EFI_EVENT  Event,
  IN VOID      *Context
  )
{
  EFI_STATUS Status;

  // Possible bug in EDK2 TCP4 driver: closing a connection doesn't remove its
  // PCB from the list of live connections. Subsequent attempts to Configure()
  // a TCP instance with the same local port will fail with INVALID_PARAMETER.
  // Calling Configure with NULL is a workaround for this issue.
  Status = mTcpConnection->Configure (mTcpConnection, NULL);

  mTcpConnection = NULL;

  Status = mTcpListener->Accept (mTcpListener, &mAcceptToken);
  if (EFI_ERROR (Status)) {
    DEBUG ((EFI_D_ERROR, "TCP Accept: %r\n", Status));
  }
}

STATIC
VOID
CloseReceiveEvents (
  VOID
  )
{
  UINTN Index;

  for (Index = 0; Index < NUM_RX_TOKENS; Index++) {
    gBS->CloseEvent (mReceiveToken[Index].CompletionToken.Event);
  }
}

/*
  Event notify function to be called when we receive TCP data.
*/
STATIC
VOID
EFIAPI
DataReceived (
  IN EFI_EVENT Event,
  IN VOID     *Context
  )
{
  EFI_STATUS                 Status;
  FASTBOOT_TCP_PACKET_LIST  *NewEntry;
  EFI_TCP4_IO_TOKEN         *ReceiveToken;

  ReceiveToken = &mReceiveToken[mNextReceiveIndex];

  Status = ReceiveToken->CompletionToken.Status;

  if (Status == EFI_CONNECTION_FIN) {
    //
    // Remote host closed connection. Close our end.
    //

    CloseReceiveEvents ();

    Status = mTcpConnection->Close (mTcpConnection, &mCloseToken);
    ASSERT_EFI_ERROR (Status);

    return;
  }

  //
  // Add an element to the receive queue
  //

  NewEntry = AllocatePool (sizeof (FASTBOOT_TCP_PACKET_LIST));
  if (NewEntry == NULL) {
    DEBUG ((EFI_D_ERROR, "TCP Fastboot: Out of resources\n"));
    return;
  }

  mNextReceiveIndex = TOKEN_NEXT (mNextReceiveIndex);

  if (!EFI_ERROR (Status)) {
    NewEntry->Buffer
      = ReceiveToken->Packet.RxData->FragmentTable[0].FragmentBuffer;
    NewEntry->BufferSize
      = ReceiveToken->Packet.RxData->FragmentTable[0].FragmentLength;

    // Prepare to receive more data
    SubmitRecieveToken();
  } else {
    // Fatal receive error. Put an entry with NULL in the queue, signifying
    // to return EFI_DEVICE_ERROR from TcpFastbootTransportReceive.
    NewEntry->Buffer = NULL;
    NewEntry->BufferSize = 0;

    DEBUG ((EFI_D_ERROR, "\nTCP Fastboot Receive error: %r\n", Status));
  }

  InsertTailList (&mPacketListHead, &NewEntry->Link);

  Status = gBS->SignalEvent (mReceiveEvent);
  ASSERT_EFI_ERROR (Status);
}


/*
  Event notify function to be called when we accept an incoming TCP connection.
*/
STATIC
VOID
EFIAPI
ConnectionAccepted (
  IN EFI_EVENT Event,
  IN VOID     *Context
  )
{
  EFI_TCP4_LISTEN_TOKEN *AcceptToken;
  EFI_STATUS             Status;
  UINTN                  Index;

  AcceptToken = (EFI_TCP4_LISTEN_TOKEN *) Context;
  Status = AcceptToken->CompletionToken.Status;

  if (EFI_ERROR (Status)) {
    DEBUG ((EFI_D_ERROR, "TCP Fastboot: Connection Error: %r\n", Status));
    return;
  }
  DEBUG ((EFI_D_ERROR, "TCP Fastboot: Connection Received.\n"));

  //
  // Accepting a new TCP connection creates a new instance of the TCP protocol.
  // Open it and prepare to receive on it.
  //

  Status = gBS->OpenProtocol (
                  AcceptToken->NewChildHandle,
                  &gEfiTcp4ProtocolGuid,
                  (VOID **) &mTcpConnection,
                  gImageHandle,
                  NULL,
                  EFI_OPEN_PROTOCOL_GET_PROTOCOL
                  );
  if (EFI_ERROR (Status)) {
    DEBUG ((EFI_D_ERROR, "Open TCP Connection: %r\n", Status));
    return;
  }

  mNextSubmitIndex = 0;
  mNextReceiveIndex = 0;

  for (Index = 0; Index < NUM_RX_TOKENS; Index++) {
    Status = gBS->CreateEvent (
                    EVT_NOTIFY_SIGNAL,
                    TPL_CALLBACK,
                    DataReceived,
                    NULL,
                    &(mReceiveToken[Index].CompletionToken.Event)
                    );
    ASSERT_EFI_ERROR (Status);
  }

  for (Index = 0; Index < NUM_RX_TOKENS; Index++) {
    SubmitRecieveToken();
  }
}

/*
  Set up TCP Fastboot transport: Configure the network device via DHCP then
  start waiting for a TCP connection on the Fastboot port.
*/
EFI_STATUS
TcpFastbootTransportStart (
  EFI_EVENT ReceiveEvent
  )
{
  EFI_STATUS                    Status;
  EFI_HANDLE                    NetDeviceHandle;
  EFI_HANDLE                   *HandleBuffer;
  EFI_IP4_MODE_DATA             Ip4ModeData;
  UINTN                         NumHandles;
  CHAR16                        IpAddrString[16];
  UINTN                         Index;

  EFI_TCP4_CONFIG_DATA TcpConfigData = {
    0x00,                                           // IPv4 Type of Service
    255,                                            // IPv4 Time to Live
    {                                               // AccessPoint:
      TRUE,                                         // Use default address
      { {0, 0, 0, 0} },                             // IP Address  (ignored - use default)
      { {0, 0, 0, 0} },                             // Subnet mask (ignored - use default)
      FixedPcdGet32 (PcdAndroidFastbootTcpPort),    // Station port
      { {0, 0, 0, 0} },                             // Remote address: accept any
      0,                                            // Remote Port: accept any
      FALSE                                         // ActiveFlag: be a "server"
    },
    NULL                                            // Default advanced TCP options
  };

  mReceiveEvent = ReceiveEvent;
  InitializeListHead (&mPacketListHead);

  mTextOut->OutputString (mTextOut, L"Initialising TCP Fastboot transport...\r\n");

  //
  // Open a passive TCP instance
  //

  Status = gBS->LocateHandleBuffer (
                  ByProtocol,
                  &gEfiTcp4ServiceBindingProtocolGuid,
                  NULL,
                  &NumHandles,
                  &HandleBuffer
                  );
  if (EFI_ERROR (Status)) {
    DEBUG ((EFI_D_ERROR, "Find TCP Service Binding: %r\n", Status));
    return Status;
  }

  // We just use the first network device
  NetDeviceHandle = HandleBuffer[0];

  Status =  gBS->OpenProtocol (
                    NetDeviceHandle,
                    &gEfiTcp4ServiceBindingProtocolGuid,
                    (VOID **) &mTcpServiceBinding,
                    gImageHandle,
                    NULL,
                    EFI_OPEN_PROTOCOL_GET_PROTOCOL
                    );
  if (EFI_ERROR (Status)) {
    DEBUG ((EFI_D_ERROR, "Open TCP Service Binding: %r\n", Status));
    return Status;
  }

  Status = mTcpServiceBinding->CreateChild (mTcpServiceBinding, &mTcpHandle);
  if (EFI_ERROR (Status)) {
    DEBUG ((EFI_D_ERROR, "TCP ServiceBinding Create: %r\n", Status));
    return Status;
  }

  Status =  gBS->OpenProtocol (
                    mTcpHandle,
                    &gEfiTcp4ProtocolGuid,
                    (VOID **) &mTcpListener,
                    gImageHandle,
                    NULL,
                    EFI_OPEN_PROTOCOL_GET_PROTOCOL
                    );
  if (EFI_ERROR (Status)) {
    DEBUG ((EFI_D_ERROR, "Open TCP Protocol: %r\n", Status));
  }

  //
  // Set up re-usable tokens
  //

  for (Index = 0; Index < NUM_RX_TOKENS; Index++) {
    mRxData[Index].UrgentFlag = FALSE;
    mRxData[Index].FragmentCount = 1;
    mReceiveToken[Index].Packet.RxData = &mRxData[Index];
  }

  mTxData.Push = TRUE;
  mTxData.Urgent = FALSE;
  mTxData.FragmentCount = 1;
  mTransmitToken.Packet.TxData = &mTxData;

  Status = gBS->CreateEvent (
                  EVT_NOTIFY_SIGNAL,
                  TPL_CALLBACK,
                  ConnectionAccepted,
                  &mAcceptToken,
                  &mAcceptToken.CompletionToken.Event
                  );
  ASSERT_EFI_ERROR (Status);

  Status = gBS->CreateEvent (
                  EVT_NOTIFY_SIGNAL,
                  TPL_CALLBACK,
                  ConnectionClosed,
                  &mCloseToken,
                  &mCloseToken.CompletionToken.Event
                  );
  ASSERT_EFI_ERROR (Status);

  //
  // Configure the TCP instance
  //

  Status = mTcpListener->Configure (mTcpListener, &TcpConfigData);
  if (Status == EFI_NO_MAPPING) {
    // Wait until the IP configuration process (probably DHCP) has finished
    do {
      Status = mTcpListener->GetModeData (mTcpListener,
                               NULL, NULL,
                               &Ip4ModeData,
                               NULL, NULL
                               );
      ASSERT_EFI_ERROR (Status);
    } while (!Ip4ModeData.IsConfigured);
    Status = mTcpListener->Configure (mTcpListener, &TcpConfigData);
  } else if (EFI_ERROR (Status)) {
    DEBUG ((EFI_D_ERROR, "TCP Configure: %r\n", Status));
    return Status;
  }

  //
  // Tell the user our address and hostname
  //
  IP4_ADDR_TO_STRING (Ip4ModeData.ConfigData.StationAddress, IpAddrString);

  mTextOut->OutputString (mTextOut, L"TCP Fastboot transport configured.");
  mTextOut->OutputString (mTextOut, L"\r\nIP address: ");
  mTextOut->OutputString (mTextOut ,IpAddrString);
  mTextOut->OutputString (mTextOut, L"\r\n");

  //
  // Start listening for a connection
  //

  Status = mTcpListener->Accept (mTcpListener, &mAcceptToken);
  if (EFI_ERROR (Status)) {
    DEBUG ((EFI_D_ERROR, "TCP Accept: %r\n", Status));
    return Status;
  }

  mTextOut->OutputString (mTextOut, L"TCP Fastboot transport initialised.\r\n");

  FreePool (HandleBuffer);

  return EFI_SUCCESS;
}

EFI_STATUS
TcpFastbootTransportStop (
  VOID
  )
{
  EFI_TCP4_CLOSE_TOKEN      CloseToken;
  EFI_STATUS                Status;
  UINTN                     EventIndex;
  FASTBOOT_TCP_PACKET_LIST *Entry;
  FASTBOOT_TCP_PACKET_LIST *NextEntry;

  // Close any existing TCP connection, blocking until it's done.
  if (mTcpConnection != NULL) {
    CloseReceiveEvents ();

    CloseToken.AbortOnClose = FALSE;

    Status = gBS->CreateEvent (0, 0, NULL, NULL, &CloseToken.CompletionToken.Event);
    ASSERT_EFI_ERROR (Status);

    Status = mTcpConnection->Close (mTcpConnection, &CloseToken);
    ASSERT_EFI_ERROR (Status);

    Status = gBS->WaitForEvent (
                    1,
                    &CloseToken.CompletionToken.Event,
                    &EventIndex
                    );
    ASSERT_EFI_ERROR (Status);

    ASSERT_EFI_ERROR (CloseToken.CompletionToken.Status);

    // Possible bug in EDK2 TCP4 driver: closing a connection doesn't remove its
    // PCB from the list of live connections. Subsequent attempts to Configure()
    // a TCP instance with the same local port will fail with INVALID_PARAMETER.
    // Calling Configure with NULL is a workaround for this issue.
    Status = mTcpConnection->Configure (mTcpConnection, NULL);
    ASSERT_EFI_ERROR (Status);
  }


  gBS->CloseEvent (mAcceptToken.CompletionToken.Event);

  // Stop listening for connections.
  // Ideally we would do this with Cancel, but it isn't implemented by EDK2.
  // So we just "reset this TCPv4 instance brutally".
  Status = mTcpListener->Configure (mTcpListener, NULL);
  ASSERT_EFI_ERROR (Status);

  Status = mTcpServiceBinding->DestroyChild (mTcpServiceBinding, &mTcpHandle);

  // Free any data the user didn't pick up
  Entry = (FASTBOOT_TCP_PACKET_LIST *) GetFirstNode (&mPacketListHead);
  while (!IsNull (&mPacketListHead, &Entry->Link)) {
    NextEntry = (FASTBOOT_TCP_PACKET_LIST *) GetNextNode (&mPacketListHead, &Entry->Link);

    RemoveEntryList (&Entry->Link);
    if (Entry->Buffer) {
      FreePool (Entry->Buffer);
    }
    FreePool (Entry);

    Entry = NextEntry;
  }

  return EFI_SUCCESS;
}

/*
  Event notify function for when data has been sent. Free resources and report
  errors.
  Context should point to the transmit IO token passed to
  TcpConnection->Transmit.
*/
STATIC
VOID
DataSent (
  EFI_EVENT Event,
  VOID     *Context
  )
{
  EFI_STATUS           Status;

  Status = mTransmitToken.CompletionToken.Status;
  if (EFI_ERROR (Status)) {
    DEBUG ((EFI_D_ERROR, "TCP Fastboot transmit result: %r\n", Status));
    gBS->SignalEvent (*(EFI_EVENT *) Context);
  }

  FreePool (mTransmitToken.Packet.TxData->FragmentTable[0].FragmentBuffer);
}

EFI_STATUS
TcpFastbootTransportSend (
  IN        UINTN      BufferSize,
  IN  CONST VOID      *Buffer,
  IN        EFI_EVENT *FatalErrorEvent
  )
{
  EFI_STATUS                Status;

  if (BufferSize > 512) {
    return EFI_INVALID_PARAMETER;
  }

  //
  // Build transmit IO token
  //

  // Create an event so we are notified when a transmission is complete.
  // We use this to free resources and report errors.
  Status = gBS->CreateEvent (
                  EVT_NOTIFY_SIGNAL,
                  TPL_CALLBACK,
                  DataSent,
                  FatalErrorEvent,
                  &mTransmitToken.CompletionToken.Event
                  );
  ASSERT_EFI_ERROR (Status);

  mTxData.DataLength = BufferSize;

  mTxData.FragmentTable[0].FragmentLength = BufferSize;
  mTxData.FragmentTable[0].FragmentBuffer = AllocateCopyPool (
                                             BufferSize,
                                             Buffer
                                             );

  Status = mTcpConnection->Transmit (mTcpConnection, &mTransmitToken);
  if (EFI_ERROR (Status)) {
    DEBUG ((EFI_D_ERROR, "TCP Transmit: %r\n", Status));
    return Status;
  }

  return EFI_SUCCESS;
}


EFI_STATUS
TcpFastbootTransportReceive (
  OUT UINTN  *BufferSize,
  OUT VOID  **Buffer
  )
{
  FASTBOOT_TCP_PACKET_LIST *Entry;

  if (IsListEmpty (&mPacketListHead)) {
    return EFI_NOT_READY;
  }

  Entry = (FASTBOOT_TCP_PACKET_LIST *) GetFirstNode (&mPacketListHead);

  if (Entry->Buffer == NULL) {
    // There was an error receiving this packet.
    return EFI_DEVICE_ERROR;
  }

  *Buffer = Entry->Buffer;
  *BufferSize = Entry->BufferSize;

  RemoveEntryList (&Entry->Link);
  FreePool (Entry);

  return EFI_SUCCESS;
}

FASTBOOT_TRANSPORT_PROTOCOL mTransportProtocol = {
  TcpFastbootTransportStart,
  TcpFastbootTransportStop,
  TcpFastbootTransportSend,
  TcpFastbootTransportReceive
};

EFI_STATUS
TcpFastbootTransportEntryPoint (
  IN EFI_HANDLE        ImageHandle,
  IN EFI_SYSTEM_TABLE *SystemTable
  )
{
  EFI_STATUS Status;


  Status = gBS->LocateProtocol(
    &gEfiSimpleTextOutProtocolGuid,
    NULL,
    (VOID **) &mTextOut
    );
  if (EFI_ERROR (Status)) {
    DEBUG ((EFI_D_ERROR, "Fastboot: Open Text Output Protocol: %r\n", Status));
    return Status;
  }

  Status = gBS->InstallProtocolInterface (
                  &ImageHandle,
                  &gAndroidFastbootTransportProtocolGuid,
                  EFI_NATIVE_INTERFACE,
                  &mTransportProtocol
                  );
  if (EFI_ERROR (Status)) {
    DEBUG ((EFI_D_ERROR, "Fastboot: Install transport Protocol: %r\n", Status));
  }

  return Status;
}