/** @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 various 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 transmitted.
  @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;

  //
  // Transmit 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 asynchronous 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 memory.
  @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;
}