/** @file
  This library is used to share code between UEFI network stack modules.
  It provides the helper routines to access TCP service.

Copyright (c) 2010 - 2018, Intel Corporation. All rights reserved.<BR>
SPDX-License-Identifier: BSD-2-Clause-Patent

**/

#include <Uefi.h>

#include <Library/TcpIoLib.h>
#include <Library/BaseLib.h>
#include <Library/DebugLib.h>
#include <Library/UefiBootServicesTableLib.h>
#include <Library/MemoryAllocationLib.h>
#include <Library/BaseMemoryLib.h>

/**
  The common notify function associated with various TcpIo events.

  @param[in]  Event   The event signaled.
  @param[in]  Context The context.

**/
VOID
EFIAPI
TcpIoCommonNotify (
  IN EFI_EVENT  Event,
  IN VOID       *Context
  )
{
  if ((Event == NULL) || (Context == NULL)) {
    return ;
  }

  *((BOOLEAN *) Context) = TRUE;
}

/**
  The internal function for delay configuring TCP6 when IP6 driver is still in DAD.

  @param[in]  Tcp6               The EFI_TCP6_PROTOCOL protocol instance.
  @param[in]  Tcp6ConfigData     The Tcp6 configuration data.

  @retval EFI_SUCCESS            The operational settings successfully
                                 completed.
  @retval EFI_INVALID_PARAMETER  One or more parameters are invalid.
  @retval Others                 Failed to finish the operation.

**/
EFI_STATUS
TcpIoGetMapping (
  IN EFI_TCP6_PROTOCOL    *Tcp6,
  IN EFI_TCP6_CONFIG_DATA *Tcp6ConfigData
  )
{
  EFI_STATUS              Status;
  EFI_EVENT               Event;

  if ((Tcp6 == NULL) || (Tcp6ConfigData == NULL)) {
    return EFI_INVALID_PARAMETER;
  }

  Event  = NULL;
  Status = gBS->CreateEvent (
                  EVT_TIMER,
                  TPL_CALLBACK,
                  NULL,
                  NULL,
                  &Event
                  );
  if (EFI_ERROR (Status)) {
    goto ON_EXIT;
  }

  Status = gBS->SetTimer (
                  Event,
                  TimerRelative,
                  TCP_GET_MAPPING_TIMEOUT
                  );

  if (EFI_ERROR (Status)) {
    goto ON_EXIT;
  }

  while (EFI_ERROR (gBS->CheckEvent (Event))) {

    Tcp6->Poll (Tcp6);

    Status = Tcp6->Configure (Tcp6, Tcp6ConfigData);

    if (!EFI_ERROR (Status)) {
      break;
    }
  }

ON_EXIT:

  if (Event != NULL) {
    gBS->CloseEvent (Event);
  }

  return Status;
}

/**
  Create a TCP socket with the specified configuration data.

  @param[in]  Image      The handle of the driver image.
  @param[in]  Controller The handle of the controller.
  @param[in]  TcpVersion The version of Tcp, TCP_VERSION_4 or TCP_VERSION_6.
  @param[in]  ConfigData The Tcp configuration data.
  @param[out] TcpIo      The TcpIo.

  @retval EFI_SUCCESS            The TCP socket is created and configured.
  @retval EFI_INVALID_PARAMETER  One or more parameters are invalid.
  @retval EFI_UNSUPPORTED        One or more of the control options are not
                                 supported in the implementation.
  @retval EFI_OUT_OF_RESOURCES   Failed to allocate memory.
  @retval Others                 Failed to create the TCP socket or configure it.

**/
EFI_STATUS
EFIAPI
TcpIoCreateSocket (
  IN EFI_HANDLE             Image,
  IN EFI_HANDLE             Controller,
  IN UINT8                  TcpVersion,
  IN TCP_IO_CONFIG_DATA     *ConfigData,
  OUT TCP_IO                *TcpIo
  )
{
  EFI_STATUS                Status;
  EFI_EVENT                 Event;
  EFI_GUID                  *ServiceBindingGuid;
  EFI_GUID                  *ProtocolGuid;
  VOID                      **Interface;
  EFI_TCP4_OPTION           ControlOption;
  EFI_TCP4_CONFIG_DATA      Tcp4ConfigData;
  EFI_TCP4_ACCESS_POINT     *AccessPoint4;
  EFI_TCP4_PROTOCOL         *Tcp4;
  EFI_TCP6_CONFIG_DATA      Tcp6ConfigData;
  EFI_TCP6_ACCESS_POINT     *AccessPoint6;
  EFI_TCP6_PROTOCOL         *Tcp6;
  EFI_TCP4_RECEIVE_DATA     *RxData;

  if ((Image == NULL) || (Controller == NULL) || (ConfigData == NULL) || (TcpIo == NULL)) {
    return EFI_INVALID_PARAMETER;
  }

  Tcp4 = NULL;
  Tcp6 = NULL;

  ZeroMem (TcpIo, sizeof (TCP_IO));

  if (TcpVersion == TCP_VERSION_4) {
    ServiceBindingGuid = &gEfiTcp4ServiceBindingProtocolGuid;
    ProtocolGuid       = &gEfiTcp4ProtocolGuid;
    Interface          = (VOID **) (&TcpIo->Tcp.Tcp4);
  } else if (TcpVersion == TCP_VERSION_6) {
    ServiceBindingGuid = &gEfiTcp6ServiceBindingProtocolGuid;
    ProtocolGuid       = &gEfiTcp6ProtocolGuid;
    Interface          = (VOID **) (&TcpIo->Tcp.Tcp6);
  } else {
    return EFI_UNSUPPORTED;
  }

  TcpIo->TcpVersion = TcpVersion;

  //
  // Create the TCP child instance and get the TCP protocol.
  //
  Status = NetLibCreateServiceChild (
             Controller,
             Image,
             ServiceBindingGuid,
             &TcpIo->Handle
             );
  if (EFI_ERROR (Status)) {
    return Status;
  }

  Status = gBS->OpenProtocol (
                  TcpIo->Handle,
                  ProtocolGuid,
                  Interface,
                  Image,
                  Controller,
                  EFI_OPEN_PROTOCOL_BY_DRIVER
                  );
  if (EFI_ERROR (Status) || (*Interface == NULL)) {
    goto ON_ERROR;
  }

  if (TcpVersion == TCP_VERSION_4) {
    Tcp4             = TcpIo->Tcp.Tcp4;
  } else {
    Tcp6             = TcpIo->Tcp.Tcp6;
  }

  TcpIo->Image       = Image;
  TcpIo->Controller  = Controller;

  //
  // Set the configuration parameters.
  //
  ControlOption.ReceiveBufferSize       = 0x200000;
  ControlOption.SendBufferSize          = 0x200000;
  ControlOption.MaxSynBackLog           = 0;
  ControlOption.ConnectionTimeout       = 0;
  ControlOption.DataRetries             = 6;
  ControlOption.FinTimeout              = 0;
  ControlOption.TimeWaitTimeout         = 0;
  ControlOption.KeepAliveProbes         = 4;
  ControlOption.KeepAliveTime           = 0;
  ControlOption.KeepAliveInterval       = 0;
  ControlOption.EnableNagle             = FALSE;
  ControlOption.EnableTimeStamp         = FALSE;
  ControlOption.EnableWindowScaling     = TRUE;
  ControlOption.EnableSelectiveAck      = FALSE;
  ControlOption.EnablePathMtuDiscovery  = FALSE;

  if (TcpVersion == TCP_VERSION_4) {
    Tcp4ConfigData.TypeOfService        = 8;
    Tcp4ConfigData.TimeToLive           = 255;
    Tcp4ConfigData.ControlOption        = &ControlOption;

    AccessPoint4                        = &Tcp4ConfigData.AccessPoint;

    ZeroMem (AccessPoint4, sizeof (EFI_TCP4_ACCESS_POINT));
    AccessPoint4->StationPort           = ConfigData->Tcp4IoConfigData.StationPort;
    AccessPoint4->RemotePort            = ConfigData->Tcp4IoConfigData.RemotePort;
    AccessPoint4->ActiveFlag            = ConfigData->Tcp4IoConfigData.ActiveFlag;

    CopyMem (
      &AccessPoint4->StationAddress,
      &ConfigData->Tcp4IoConfigData.LocalIp,
      sizeof (EFI_IPv4_ADDRESS)
      );
    CopyMem (
      &AccessPoint4->SubnetMask,
      &ConfigData->Tcp4IoConfigData.SubnetMask,
      sizeof (EFI_IPv4_ADDRESS)
      );
    CopyMem (
      &AccessPoint4->RemoteAddress,
      &ConfigData->Tcp4IoConfigData.RemoteIp,
      sizeof (EFI_IPv4_ADDRESS)
      );

    ASSERT (Tcp4 != NULL);

    //
    // Configure the TCP4 protocol.
    //
    Status = Tcp4->Configure (Tcp4, &Tcp4ConfigData);
    if (EFI_ERROR (Status)) {
      goto ON_ERROR;
    }

    if (!EFI_IP4_EQUAL (&ConfigData->Tcp4IoConfigData.Gateway, &mZeroIp4Addr)) {
      //
      // The gateway is not zero. Add the default route manually.
      //
      Status = Tcp4->Routes (
                       Tcp4,
                       FALSE,
                       &mZeroIp4Addr,
                       &mZeroIp4Addr,
                       &ConfigData->Tcp4IoConfigData.Gateway
                       );
      if (EFI_ERROR (Status)) {
        goto ON_ERROR;
      }
    }
  } else {
    Tcp6ConfigData.TrafficClass         = 0;
    Tcp6ConfigData.HopLimit             = 255;
    Tcp6ConfigData.ControlOption        = (EFI_TCP6_OPTION *) &ControlOption;

    AccessPoint6                        = &Tcp6ConfigData.AccessPoint;

    ZeroMem (AccessPoint6, sizeof (EFI_TCP6_ACCESS_POINT));
    AccessPoint6->StationPort           = ConfigData->Tcp6IoConfigData.StationPort;
    AccessPoint6->RemotePort            = ConfigData->Tcp6IoConfigData.RemotePort;
    AccessPoint6->ActiveFlag            = ConfigData->Tcp6IoConfigData.ActiveFlag;

    IP6_COPY_ADDRESS (&AccessPoint6->RemoteAddress, &ConfigData->Tcp6IoConfigData.RemoteIp);


    ASSERT (Tcp6 != NULL);
    //
    // Configure the TCP6 protocol.
    //
    Status = Tcp6->Configure (Tcp6, &Tcp6ConfigData);
    if (Status == EFI_NO_MAPPING) {
      Status = TcpIoGetMapping (Tcp6, &Tcp6ConfigData);
    }

    if (EFI_ERROR (Status)) {
      goto ON_ERROR;
    }
  }

  //
  // Create events for variuos asynchronous operations.
  //
  Status = gBS->CreateEvent (
                  EVT_NOTIFY_SIGNAL,
                  TPL_NOTIFY,
                  TcpIoCommonNotify,
                  &TcpIo->IsConnDone,
                  &Event
                  );
  if (EFI_ERROR (Status)) {
    goto ON_ERROR;
  }

  TcpIo->ConnToken.Tcp4Token.CompletionToken.Event = Event;

  Status = gBS->CreateEvent (
                  EVT_NOTIFY_SIGNAL,
                  TPL_NOTIFY,
                  TcpIoCommonNotify,
                  &TcpIo->IsListenDone,
                  &Event
                  );
  if (EFI_ERROR (Status)) {
    goto ON_ERROR;
  }

  TcpIo->ListenToken.Tcp4Token.CompletionToken.Event = Event;

  Status = gBS->CreateEvent (
                  EVT_NOTIFY_SIGNAL,
                  TPL_NOTIFY,
                  TcpIoCommonNotify,
                  &TcpIo->IsTxDone,
                  &Event
                  );
  if (EFI_ERROR (Status)) {
    goto ON_ERROR;
  }

  TcpIo->TxToken.Tcp4Token.CompletionToken.Event = Event;


  Status = gBS->CreateEvent (
                  EVT_NOTIFY_SIGNAL,
                  TPL_NOTIFY,
                  TcpIoCommonNotify,
                  &TcpIo->IsRxDone,
                  &Event
                  );
  if (EFI_ERROR (Status)) {
    goto ON_ERROR;
  }

  TcpIo->RxToken.Tcp4Token.CompletionToken.Event = Event;

  RxData = (EFI_TCP4_RECEIVE_DATA *) AllocateZeroPool (sizeof (EFI_TCP4_RECEIVE_DATA));
  if (RxData == NULL) {
    Status = EFI_OUT_OF_RESOURCES;
    goto ON_ERROR;
  }

  TcpIo->RxToken.Tcp4Token.Packet.RxData = RxData;

  Status = gBS->CreateEvent (
                  EVT_NOTIFY_SIGNAL,
                  TPL_NOTIFY,
                  TcpIoCommonNotify,
                  &TcpIo->IsCloseDone,
                  &Event
                  );
  if (EFI_ERROR (Status)) {
    goto ON_ERROR;
  }

  TcpIo->CloseToken.Tcp4Token.CompletionToken.Event = Event;


  return EFI_SUCCESS;

ON_ERROR:

  TcpIoDestroySocket (TcpIo);

  return Status;
}

/**
  Destroy the socket.

  @param[in]  TcpIo The TcpIo which wraps the socket to be destroyed.

**/
VOID
EFIAPI
TcpIoDestroySocket (
  IN TCP_IO                 *TcpIo
  )
{
  EFI_EVENT                 Event;
  EFI_TCP4_PROTOCOL         *Tcp4;
  EFI_TCP6_PROTOCOL         *Tcp6;
  UINT8                     TcpVersion;
  EFI_GUID                  *ServiceBindingGuid;
  EFI_GUID                  *ProtocolGuid;
  EFI_HANDLE                ChildHandle;

  if (TcpIo == NULL) {
    return ;
  }

  TcpVersion = TcpIo->TcpVersion;

  if ((TcpVersion != TCP_VERSION_4) && (TcpVersion != TCP_VERSION_6)) {
    return ;
  }

  Event = TcpIo->ConnToken.Tcp4Token.CompletionToken.Event;

  if (Event != NULL) {
    gBS->CloseEvent (Event);
  }

  Event = TcpIo->ListenToken.Tcp4Token.CompletionToken.Event;

  if (Event != NULL) {
    gBS->CloseEvent (Event);
  }

  Event = TcpIo->TxToken.Tcp4Token.CompletionToken.Event;

  if (Event != NULL) {
    gBS->CloseEvent (Event);
  }

  Event = TcpIo->RxToken.Tcp4Token.CompletionToken.Event;

  if (Event != NULL) {
    gBS->CloseEvent (Event);
  }

  Event = TcpIo->CloseToken.Tcp4Token.CompletionToken.Event;

  if (Event != NULL) {
    gBS->CloseEvent (Event);
  }

  if (TcpIo->RxToken.Tcp4Token.Packet.RxData != NULL) {
    FreePool (TcpIo->RxToken.Tcp4Token.Packet.RxData);
  }

  Tcp4 = NULL;
  Tcp6 = NULL;


  if (TcpVersion == TCP_VERSION_4) {
    ServiceBindingGuid = &gEfiTcp4ServiceBindingProtocolGuid;
    ProtocolGuid       = &gEfiTcp4ProtocolGuid;
    Tcp4 = TcpIo->Tcp.Tcp4;
    if (Tcp4 != NULL) {
      Tcp4->Configure (Tcp4, NULL);
    }
  } else {
    ServiceBindingGuid = &gEfiTcp6ServiceBindingProtocolGuid;
    ProtocolGuid       = &gEfiTcp6ProtocolGuid;
    Tcp6 = TcpIo->Tcp.Tcp6;
    if (Tcp6 != NULL) {
      Tcp6->Configure (Tcp6, NULL);
    }
  }

  if ((Tcp4 != NULL) || (Tcp6 != NULL)) {

    gBS->CloseProtocol (
           TcpIo->Handle,
           ProtocolGuid,
           TcpIo->Image,
           TcpIo->Controller
           );
  }

  ChildHandle = NULL;

  if (TcpIo->IsListenDone) {
    if (TcpVersion == TCP_VERSION_4) {
      Tcp4 = TcpIo->NewTcp.Tcp4;
      if (Tcp4 != NULL) {
        Tcp4->Configure (Tcp4, NULL);
        ChildHandle = TcpIo->ListenToken.Tcp4Token.NewChildHandle;
      }
    } else {
      Tcp6 = TcpIo->NewTcp.Tcp6;
      if (Tcp6 != NULL) {
        Tcp6->Configure (Tcp6, NULL);
        ChildHandle = TcpIo->ListenToken.Tcp6Token.NewChildHandle;
      }
    }

    if (ChildHandle != NULL) {

      gBS->CloseProtocol (
             ChildHandle,
             ProtocolGuid,
             TcpIo->Image,
             TcpIo->Controller
             );
    }
  }

  NetLibDestroyServiceChild (
    TcpIo->Controller,
    TcpIo->Image,
    ServiceBindingGuid,
    TcpIo->Handle
    );
}

/**
  Connect to the other endpoint of the TCP socket.

  @param[in, out]  TcpIo     The TcpIo wrapping the TCP socket.
  @param[in]       Timeout   The time to wait for connection done. Set to NULL for infinite wait.

  @retval EFI_SUCCESS            Connect to the other endpoint of the TCP socket
                                 successfully.
  @retval EFI_TIMEOUT            Failed to connect to the other endpoint of the
                                 TCP socket in the specified time period.
  @retval EFI_INVALID_PARAMETER  One or more parameters are invalid.
  @retval EFI_UNSUPPORTED        One or more of the control options are not
                                 supported in the implementation.
  @retval Others                 Other errors as indicated.

**/
EFI_STATUS
EFIAPI
TcpIoConnect (
  IN OUT TCP_IO             *TcpIo,
  IN     EFI_EVENT          Timeout        OPTIONAL
  )
{
  EFI_TCP4_PROTOCOL         *Tcp4;
  EFI_TCP6_PROTOCOL         *Tcp6;
  EFI_STATUS                Status;

  if ((TcpIo == NULL) || (TcpIo->Tcp.Tcp4 == NULL)) {
    return EFI_INVALID_PARAMETER;
  }

  TcpIo->IsConnDone = FALSE;

  Tcp4 = NULL;
  Tcp6 = NULL;

  if (TcpIo->TcpVersion == TCP_VERSION_4) {
    Tcp4   = TcpIo->Tcp.Tcp4;
    Status = Tcp4->Connect (Tcp4, &TcpIo->ConnToken.Tcp4Token);
  } else if (TcpIo->TcpVersion == TCP_VERSION_6) {
    Tcp6   = TcpIo->Tcp.Tcp6;
    Status = Tcp6->Connect (Tcp6, &TcpIo->ConnToken.Tcp6Token);
  } else {
    return EFI_UNSUPPORTED;
  }

  if (EFI_ERROR (Status)) {
    return Status;
  }

  while (!TcpIo->IsConnDone && ((Timeout == NULL) || EFI_ERROR (gBS->CheckEvent (Timeout)))) {
    if (TcpIo->TcpVersion == TCP_VERSION_4) {
      Tcp4->Poll (Tcp4);
    } else {
      Tcp6->Poll (Tcp6);
    }
  }

  if (!TcpIo->IsConnDone) {
    if (TcpIo->TcpVersion == TCP_VERSION_4) {
      Tcp4->Cancel (Tcp4, &TcpIo->ConnToken.Tcp4Token.CompletionToken);
    } else {
      Tcp6->Cancel (Tcp6, &TcpIo->ConnToken.Tcp6Token.CompletionToken);
    }
    Status = EFI_TIMEOUT;
  } else {
    Status = TcpIo->ConnToken.Tcp4Token.CompletionToken.Status;
  }

  return Status;
}

/**
  Accept the incomding request from the other endpoint of the TCP socket.

  @param[in, out]  TcpIo     The TcpIo wrapping the TCP socket.
  @param[in]       Timeout   The time to wait for connection done. Set to NULL for infinite wait.


  @retval EFI_SUCCESS            Connect to the other endpoint of the TCP socket
                                 successfully.
  @retval EFI_INVALID_PARAMETER  One or more parameters are invalid.
  @retval EFI_UNSUPPORTED        One or more of the control options are not
                                 supported in the implementation.

  @retval EFI_TIMEOUT            Failed to connect to the other endpoint of the
                                 TCP socket in the specified time period.
  @retval Others                 Other errors as indicated.

**/
EFI_STATUS
EFIAPI
TcpIoAccept (
  IN OUT TCP_IO             *TcpIo,
  IN     EFI_EVENT          Timeout        OPTIONAL
  )
{
  EFI_STATUS                Status;
  EFI_GUID                  *ProtocolGuid;
  EFI_TCP4_PROTOCOL         *Tcp4;
  EFI_TCP6_PROTOCOL         *Tcp6;

  if ((TcpIo == NULL) || (TcpIo->Tcp.Tcp4 == NULL)) {
    return EFI_INVALID_PARAMETER;
  }

  TcpIo->IsListenDone = FALSE;

  Tcp4 = NULL;
  Tcp6 = NULL;

  if (TcpIo->TcpVersion == TCP_VERSION_4) {
    Tcp4   = TcpIo->Tcp.Tcp4;
    Status = Tcp4->Accept (Tcp4, &TcpIo->ListenToken.Tcp4Token);
  } else if (TcpIo->TcpVersion == TCP_VERSION_6) {
    Tcp6   = TcpIo->Tcp.Tcp6;
    Status = Tcp6->Accept (Tcp6, &TcpIo->ListenToken.Tcp6Token);
  } else {
    return EFI_UNSUPPORTED;
  }

  if (EFI_ERROR (Status)) {
    return Status;
  }

  while (!TcpIo->IsListenDone && ((Timeout == NULL) || EFI_ERROR (gBS->CheckEvent (Timeout)))) {
    if (TcpIo->TcpVersion == TCP_VERSION_4) {
      Tcp4->Poll (Tcp4);
    } else {
      Tcp6->Poll (Tcp6);
    }
  }

  if (!TcpIo->IsListenDone) {
    if (TcpIo->TcpVersion == TCP_VERSION_4) {
      Tcp4->Cancel (Tcp4, &TcpIo->ListenToken.Tcp4Token.CompletionToken);
    } else {
      Tcp6->Cancel (Tcp6, &TcpIo->ListenToken.Tcp6Token.CompletionToken);
    }
    Status = EFI_TIMEOUT;
  } else {
    Status = TcpIo->ListenToken.Tcp4Token.CompletionToken.Status;
  }

  //
  // The new TCP instance handle created for the established connection is
  // in ListenToken.
  //
  if (!EFI_ERROR (Status)) {
    if (TcpIo->TcpVersion == TCP_VERSION_4) {
      ProtocolGuid = &gEfiTcp4ProtocolGuid;
    } else {
      ProtocolGuid = &gEfiTcp6ProtocolGuid;
    }

    Status = gBS->OpenProtocol (
                    TcpIo->ListenToken.Tcp4Token.NewChildHandle,
                    ProtocolGuid,
                    (VOID **) (&TcpIo->NewTcp.Tcp4),
                    TcpIo->Image,
                    TcpIo->Controller,
                    EFI_OPEN_PROTOCOL_BY_DRIVER
                    );

  }

  return Status;
}

/**
  Reset the socket.

  @param[in, out]  TcpIo The TcpIo wrapping the TCP socket.

**/
VOID
EFIAPI
TcpIoReset (
  IN OUT TCP_IO             *TcpIo
  )
{
  EFI_TCP4_PROTOCOL         *Tcp4;
  EFI_TCP6_PROTOCOL         *Tcp6;
  EFI_STATUS                Status;

  if ((TcpIo == NULL) || (TcpIo->Tcp.Tcp4 == NULL)) {
    return ;
  }

  TcpIo->IsCloseDone = FALSE;
  Tcp4               = NULL;
  Tcp6               = NULL;

  if (TcpIo->TcpVersion == TCP_VERSION_4) {
    TcpIo->CloseToken.Tcp4Token.AbortOnClose = TRUE;
    Tcp4 = TcpIo->Tcp.Tcp4;
    Status = Tcp4->Close (Tcp4, &TcpIo->CloseToken.Tcp4Token);
  } else if (TcpIo->TcpVersion == TCP_VERSION_6) {
    TcpIo->CloseToken.Tcp6Token.AbortOnClose = TRUE;
    Tcp6 = TcpIo->Tcp.Tcp6;
    Status = Tcp6->Close (Tcp6, &TcpIo->CloseToken.Tcp6Token);
  } else {
    return ;
  }

  if (EFI_ERROR (Status)) {
    return ;
  }

  while (!TcpIo->IsCloseDone) {
    if (TcpIo->TcpVersion == TCP_VERSION_4) {
      Tcp4->Poll (Tcp4);
    } else {
      Tcp6->Poll (Tcp6);
    }
  }
}


/**
  Transmit the Packet to the other endpoint of the socket.

  @param[in]   TcpIo           The TcpIo wrapping the TCP socket.
  @param[in]   Packet          The packet to transmit.

  @retval EFI_SUCCESS            The packet is trasmitted.
  @retval EFI_INVALID_PARAMETER  One or more parameters are invalid.
  @retval EFI_UNSUPPORTED        One or more of the control options are not
                                 supported in the implementation.
  @retval EFI_OUT_OF_RESOURCES   Failed to allocate memory.
  @retval EFI_DEVICE_ERROR       An unexpected network or system error occurred.
  @retval Others                 Other errors as indicated.

**/
EFI_STATUS
EFIAPI
TcpIoTransmit (
  IN TCP_IO                 *TcpIo,
  IN NET_BUF                *Packet
  )
{
  EFI_STATUS                Status;
  VOID                      *Data;
  EFI_TCP4_PROTOCOL         *Tcp4;
  EFI_TCP6_PROTOCOL         *Tcp6;
  UINTN                     Size;

  if ((TcpIo == NULL) || (TcpIo->Tcp.Tcp4 == NULL)|| (Packet == NULL)) {
    return EFI_INVALID_PARAMETER;
  }

  if (TcpIo->TcpVersion == TCP_VERSION_4) {

    Size = sizeof (EFI_TCP4_TRANSMIT_DATA) +
           (Packet->BlockOpNum - 1) * sizeof (EFI_TCP4_FRAGMENT_DATA);
  } else if (TcpIo->TcpVersion == TCP_VERSION_6) {
    Size = sizeof (EFI_TCP6_TRANSMIT_DATA) +
           (Packet->BlockOpNum - 1) * sizeof (EFI_TCP6_FRAGMENT_DATA);
  } else {
    return EFI_UNSUPPORTED;
  }

  Data = AllocatePool (Size);
  if (Data == NULL) {
    return EFI_OUT_OF_RESOURCES;
  }

  ((EFI_TCP4_TRANSMIT_DATA *) Data)->Push        = TRUE;
  ((EFI_TCP4_TRANSMIT_DATA *) Data)->Urgent      = FALSE;
  ((EFI_TCP4_TRANSMIT_DATA *) Data)->DataLength  = Packet->TotalSize;

  //
  // Build the fragment table.
  //
  ((EFI_TCP4_TRANSMIT_DATA *) Data)->FragmentCount = Packet->BlockOpNum;

  NetbufBuildExt (
    Packet,
    (NET_FRAGMENT *) &((EFI_TCP4_TRANSMIT_DATA *) Data)->FragmentTable[0],
    &((EFI_TCP4_TRANSMIT_DATA *) Data)->FragmentCount
    );

  Tcp4   = NULL;
  Tcp6   = NULL;
  Status = EFI_DEVICE_ERROR;

  //
  // Trasnmit the packet.
  //
  if (TcpIo->TcpVersion == TCP_VERSION_4) {
    TcpIo->TxToken.Tcp4Token.Packet.TxData = (EFI_TCP4_TRANSMIT_DATA *) Data;
    Tcp4    = TcpIo->Tcp.Tcp4;
    if (TcpIo->IsListenDone) {
      Tcp4 = TcpIo->NewTcp.Tcp4;
    }

    if (Tcp4 == NULL) {
      goto ON_EXIT;
    }

    Status  = Tcp4->Transmit (Tcp4, &TcpIo->TxToken.Tcp4Token);
  } else {
    TcpIo->TxToken.Tcp6Token.Packet.TxData = (EFI_TCP6_TRANSMIT_DATA *) Data;
    Tcp6    = TcpIo->Tcp.Tcp6;
    if (TcpIo->IsListenDone) {
      Tcp6 = TcpIo->NewTcp.Tcp6;
    }

    if (Tcp6 == NULL) {
      goto ON_EXIT;
    }

    Status  = Tcp6->Transmit (Tcp6, &TcpIo->TxToken.Tcp6Token);
  }

  if (EFI_ERROR (Status)) {
    goto ON_EXIT;
  }

  while (!TcpIo->IsTxDone) {
    if (TcpIo->TcpVersion == TCP_VERSION_4) {
      Tcp4->Poll (Tcp4);
    } else {
      Tcp6->Poll (Tcp6);
    }
  }

  TcpIo->IsTxDone  = FALSE;
  Status           = TcpIo->TxToken.Tcp4Token.CompletionToken.Status;

ON_EXIT:

  FreePool (Data);

  return Status;
}

/**
  Receive data from the socket.

  @param[in, out]  TcpIo       The TcpIo which wraps the socket to be destroyed.
  @param[in]       Packet      The buffer to hold the data copy from the socket rx buffer.
  @param[in]       AsyncMode   Is this receive asyncronous or not.
  @param[in]       Timeout     The time to wait for receiving the amount of data the Packet
                               can hold. Set to NULL for infinite wait.

  @retval EFI_SUCCESS            The required amount of data is received from the socket.
  @retval EFI_INVALID_PARAMETER  One or more parameters are invalid.
  @retval EFI_DEVICE_ERROR       An unexpected network or system error occurred.
  @retval EFI_OUT_OF_RESOURCES   Failed to allocate momery.
  @retval EFI_TIMEOUT            Failed to receive the required amount of data in the
                                 specified time period.
  @retval Others                 Other errors as indicated.

**/
EFI_STATUS
EFIAPI
TcpIoReceive (
  IN OUT TCP_IO             *TcpIo,
  IN     NET_BUF            *Packet,
  IN     BOOLEAN            AsyncMode,
  IN     EFI_EVENT          Timeout       OPTIONAL
  )
{
  EFI_TCP4_PROTOCOL         *Tcp4;
  EFI_TCP6_PROTOCOL         *Tcp6;
  EFI_TCP4_RECEIVE_DATA     *RxData;
  EFI_STATUS                Status;
  NET_FRAGMENT              *Fragment;
  UINT32                    FragmentCount;
  UINT32                    CurrentFragment;

  if ((TcpIo == NULL) || (TcpIo->Tcp.Tcp4 == NULL)|| (Packet == NULL)) {
    return EFI_INVALID_PARAMETER;
  }

  RxData = TcpIo->RxToken.Tcp4Token.Packet.RxData;
  if (RxData == NULL) {
    return EFI_INVALID_PARAMETER;
  }

  Tcp4 = NULL;
  Tcp6 = NULL;

  if (TcpIo->TcpVersion == TCP_VERSION_4) {
    Tcp4 = TcpIo->Tcp.Tcp4;

    if (TcpIo->IsListenDone) {
      Tcp4 = TcpIo->NewTcp.Tcp4;
    }

    if (Tcp4 == NULL) {
      return EFI_DEVICE_ERROR;
    }

  } else if (TcpIo->TcpVersion == TCP_VERSION_6) {
    Tcp6 = TcpIo->Tcp.Tcp6;

    if (TcpIo->IsListenDone) {
      Tcp6 = TcpIo->NewTcp.Tcp6;
    }

    if (Tcp6 == NULL) {
      return EFI_DEVICE_ERROR;
    }

  } else {
    return EFI_UNSUPPORTED;
  }

  FragmentCount = Packet->BlockOpNum;
  Fragment      = AllocatePool (FragmentCount * sizeof (NET_FRAGMENT));
  if (Fragment == NULL) {
    Status = EFI_OUT_OF_RESOURCES;
    goto ON_EXIT;
  }
  //
  // Build the fragment table.
  //
  NetbufBuildExt (Packet, Fragment, &FragmentCount);

  RxData->FragmentCount         = 1;
  CurrentFragment               = 0;
  Status                        = EFI_SUCCESS;

  while (CurrentFragment < FragmentCount) {
    RxData->DataLength                       = Fragment[CurrentFragment].Len;
    RxData->FragmentTable[0].FragmentLength  = Fragment[CurrentFragment].Len;
    RxData->FragmentTable[0].FragmentBuffer  = Fragment[CurrentFragment].Bulk;

    if (TcpIo->TcpVersion == TCP_VERSION_4) {
      Status = Tcp4->Receive (Tcp4, &TcpIo->RxToken.Tcp4Token);
    } else {
      Status = Tcp6->Receive (Tcp6, &TcpIo->RxToken.Tcp6Token);
    }

    if (EFI_ERROR (Status)) {
      goto ON_EXIT;
    }

    while (!TcpIo->IsRxDone && ((Timeout == NULL) || EFI_ERROR (gBS->CheckEvent (Timeout)))) {
      //
      // Poll until some data is received or an error occurs.
      //
      if (TcpIo->TcpVersion == TCP_VERSION_4) {
        Tcp4->Poll (Tcp4);
      } else {
        Tcp6->Poll (Tcp6);
      }
    }

    if (!TcpIo->IsRxDone) {
      //
      // Timeout occurs, cancel the receive request.
      //
      if (TcpIo->TcpVersion == TCP_VERSION_4) {
        Tcp4->Cancel (Tcp4, &TcpIo->RxToken.Tcp4Token.CompletionToken);
      } else {
        Tcp6->Cancel (Tcp6, &TcpIo->RxToken.Tcp6Token.CompletionToken);
      }

      Status = EFI_TIMEOUT;
      goto ON_EXIT;
    } else {
      TcpIo->IsRxDone = FALSE;
    }

    Status = TcpIo->RxToken.Tcp4Token.CompletionToken.Status;

    if (EFI_ERROR (Status)) {
      goto ON_EXIT;
    }

    Fragment[CurrentFragment].Len -= RxData->FragmentTable[0].FragmentLength;
    if (Fragment[CurrentFragment].Len == 0) {
      CurrentFragment++;
    } else {
      Fragment[CurrentFragment].Bulk += RxData->FragmentTable[0].FragmentLength;
    }
  }

ON_EXIT:

  if (Fragment != NULL) {
    FreePool (Fragment);
  }

  return Status;
}