Fixed GCC 4.4 build issues due to EFIAPI not being used when required.

[mirror_edk2.git] / MdeModulePkg / Library / DxeUdpIoLib / DxeUdpIoLib.c

diff --git a/MdeModulePkg/Library/DxeUdpIoLib/DxeUdpIoLib.c b/MdeModulePkg/Library/DxeUdpIoLib/DxeUdpIoLib.c
index c24f88827e384bf4615760c72ef589d6369780ad..a27f03aab91414c891918e777eb27bb0bf6f83aa 100644 (file)
--- a/MdeModulePkg/Library/DxeUdpIoLib/DxeUdpIoLib.c
+++ b/MdeModulePkg/Library/DxeUdpIoLib/DxeUdpIoLib.c
@@ -1,29 +1,20 @@
 /** @file\r
 /** @file\r

+  Help functions to access UDP service, it is used by both the DHCP and MTFTP.\r

 \r
 \r

-Copyright (c) 2006 - 2007, Intel Corporation\r
+Copyright (c) 2005 - 2009, Intel Corporation.<BR>\r

 All rights reserved. This program and the accompanying materials\r
 are licensed and made available under the terms and conditions of the BSD License\r
 All rights reserved. This program and the accompanying materials\r
 are licensed and made available under the terms and conditions of the BSD License\r

-which accompanies this distribution.  The full text of the license may be found at\r
+which accompanies this distribution.  The full text of the license may be found at<BR>\r

 http://opensource.org/licenses/bsd-license.php\r
 \r
 THE PROGRAM IS DISTRIBUTED UNDER THE BSD LICENSE ON AN "AS IS" BASIS,\r
 WITHOUT WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.\r
 http://opensource.org/licenses/bsd-license.php\r
 \r
 THE PROGRAM IS DISTRIBUTED UNDER THE BSD LICENSE ON AN "AS IS" BASIS,\r
 WITHOUT WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.\r

-\r
-\r
-Module Name:\r
-\r
-  Udp4Io.c\r
-\r
-Abstract:\r
-\r
-  Help functions to access UDP service, it is used by both the DHCP and MTFTP.\r
-\r
-\r

 **/\r
 \r
 **/\r
 \r

-#include <PiDxe.h>\r
+#include <Uefi.h>\r

 \r
 #include <Protocol/Udp4.h>\r
 \r
 #include <Protocol/Udp4.h>\r

+#include <Protocol/Udp6.h>\r

 \r
 #include <Library/UdpIoLib.h>\r
 #include <Library/BaseLib.h>\r
 \r
 #include <Library/UdpIoLib.h>\r
 #include <Library/BaseLib.h>\r


@@ -31,155 +22,311 @@ Abstract:
 #include <Library/UefiBootServicesTableLib.h>\r
 #include <Library/MemoryAllocationLib.h>\r
 #include <Library/BaseMemoryLib.h>\r
 #include <Library/UefiBootServicesTableLib.h>\r
 #include <Library/MemoryAllocationLib.h>\r
 #include <Library/BaseMemoryLib.h>\r

+#include <Library/DpcLib.h>\r
+\r
+\r
+/**\r
+  Free a UDP_TX_TOKEN. The TX event is closed.\r
+\r
+  @param[in]  TxToken                 The UDP_TX_TOKEN to release.\r

 \r
 \r

+**/\r
+VOID\r
+UdpIoFreeTxToken (\r
+  IN UDP_TX_TOKEN           *TxToken\r
+  )\r
+{\r
+\r
+  if (TxToken->UdpIo->UdpVersion == UDP_IO_UDP4_VERSION) {\r
+    gBS->CloseEvent (TxToken->Token.Udp4.Event);\r
+  } else if (TxToken->UdpIo->UdpVersion == UDP_IO_UDP6_VERSION) {\r
+    gBS->CloseEvent (TxToken->Token.Udp6.Event);\r
+  } else {\r
+    ASSERT (FALSE);\r
+  }\r
+\r
+  FreePool (TxToken);\r
+}\r
+\r
+/**\r
+  Free a UDP_RX_TOKEN. The RX event is closed.\r
+\r
+  @param[in]  RxToken                 The UDP_RX_TOKEN to release.\r
+\r
+**/\r
+VOID\r
+UdpIoFreeRxToken (\r
+  IN UDP_RX_TOKEN           *RxToken\r
+  )\r
+{\r
+  if (RxToken->UdpIo->UdpVersion == UDP_IO_UDP4_VERSION) {\r
+    gBS->CloseEvent (RxToken->Token.Udp4.Event);\r
+  } else if (RxToken->UdpIo->UdpVersion == UDP_IO_UDP6_VERSION) {\r
+    gBS->CloseEvent (RxToken->Token.Udp6.Event);\r
+  } else {\r
+    ASSERT (FALSE);\r
+  }\r
+\r
+  FreePool (RxToken);\r
+}\r
+\r
+/**\r
+  The callback function when the packet is sent by UDP.\r
+\r
+  It will remove the packet from the local list then call\r
+  the packet owner's callback function set by UdpIoSendDatagram.\r
+\r
+  @param[in]  Context               The UDP TX Token.\r
+\r
+**/\r

 VOID\r
 EFIAPI\r
 UdpIoOnDgramSentDpc (\r
   IN VOID                   *Context\r
 VOID\r
 EFIAPI\r
 UdpIoOnDgramSentDpc (\r
   IN VOID                   *Context\r

-  );\r
+  )\r
+{\r
+  UDP_TX_TOKEN              *TxToken;\r
+\r
+  TxToken = (UDP_TX_TOKEN *) Context;\r
+  ASSERT (TxToken->Signature == UDP_IO_TX_SIGNATURE);\r
+  ASSERT ((TxToken->UdpIo->UdpVersion == UDP_IO_UDP4_VERSION) ||\r
+          (TxToken->UdpIo->UdpVersion == UDP_IO_UDP6_VERSION));\r
+\r
+  RemoveEntryList (&TxToken->Link);\r
+\r
+  if (TxToken->UdpIo->UdpVersion == UDP_IO_UDP4_VERSION) {\r
+    TxToken->CallBack (TxToken->Packet, NULL, TxToken->Token.Udp4.Status, TxToken->Context);\r
+  } else {\r
+    TxToken->CallBack (TxToken->Packet, NULL, TxToken->Token.Udp6.Status, TxToken->Context);\r
+  }\r
+\r
+  UdpIoFreeTxToken (TxToken);\r
+}\r
+\r
+/**\r
+  Request UdpIoOnDgramSentDpc as a DPC at TPL_CALLBACK.\r
+\r
+  @param[in]  Event                 The event signaled.\r
+  @param[in]  Context               The UDP TX Token.\r

 \r
 \r

+**/\r

 VOID\r
 EFIAPI\r
 UdpIoOnDgramSent (\r
   IN EFI_EVENT              Event,\r
   IN VOID                   *Context\r
 VOID\r
 EFIAPI\r
 UdpIoOnDgramSent (\r
   IN EFI_EVENT              Event,\r
   IN VOID                   *Context\r

-  );\r
+  )\r
+{\r
+  //\r
+  // Request UdpIoOnDgramSentDpc as a DPC at TPL_CALLBACK\r
+  //\r
+  QueueDpc (TPL_CALLBACK, UdpIoOnDgramSentDpc, Context);\r
+}\r
+\r
+/**\r
+  Recycle the received UDP data.\r
+\r
+  @param[in]  Context               The UDP_RX_TOKEN.\r

 \r
 \r

+**/\r

 VOID\r
 EFIAPI\r
 VOID\r
 EFIAPI\r

-UdpIoOnDgramRcvd (\r
-  IN EFI_EVENT              Event,\r
+UdpIoRecycleDgram (\r

   IN VOID                   *Context\r
   IN VOID                   *Context\r

-  );\r
+  )\r
+{\r
+  UDP_RX_TOKEN              *RxToken;\r
+\r
+  RxToken = (UDP_RX_TOKEN *) Context;\r

 \r
 \r

+  if (RxToken->UdpIo->UdpVersion == UDP_IO_UDP4_VERSION) {\r
+    gBS->SignalEvent (RxToken->Token.Udp4.Packet.RxData->RecycleSignal);\r
+  } else if (RxToken->UdpIo->UdpVersion == UDP_IO_UDP6_VERSION) {\r
+    gBS->SignalEvent (RxToken->Token.Udp6.Packet.RxData->RecycleSignal);\r
+  } else {\r
+    ASSERT (FALSE);\r
+  }\r
+\r
+  UdpIoFreeRxToken (RxToken);\r
+}\r

 \r
 /**\r
 \r
 /**\r

-  Wrap a transmit request into a UDP_TX_TOKEN.\r
+  The event handle for UDP receive request.\r

 \r
 \r

-  @param  UdpIo                 The UdpIo port to send packet to\r
-  @param  Packet                The user's packet\r
-  @param  EndPoint              The local and remote access point\r
-  @param  Gateway               The overrided next hop\r
-  @param  CallBack              The function to call when transmission completed.\r
-  @param  Context               The opaque parameter to the call back\r
+  It will build a NET_BUF from the recieved UDP data, then deliver it\r
+  to the receiver.\r

 \r
 \r

-  @return The wrapped transmission request or NULL if failed to allocate resources.\r
+  @param[in]  Context               The UDP RX token.\r

 \r
 **/\r
 \r
 **/\r

-UDP_TX_TOKEN *\r
-UdpIoWrapTx (\r
-  IN UDP_IO_PORT            *UdpIo,\r
-  IN NET_BUF                *Packet,\r
-  IN UDP_POINTS             *EndPoint, OPTIONAL\r
-  IN IP4_ADDR               Gateway,\r
-  IN UDP_IO_CALLBACK        CallBack,\r
+VOID\r
+EFIAPI\r
+UdpIoOnDgramRcvdDpc (\r

   IN VOID                   *Context\r
   )\r
 {\r
   IN VOID                   *Context\r
   )\r
 {\r

-  UDP_TX_TOKEN              *Token;\r
-  EFI_UDP4_COMPLETION_TOKEN *UdpToken;\r
-  EFI_UDP4_TRANSMIT_DATA    *UdpTxData;\r

   EFI_STATUS                Status;\r
   EFI_STATUS                Status;\r

-  UINT32                    Count;\r
-  IP4_ADDR                  Ip;\r
+  VOID                      *Token;\r
+  VOID                      *RxData;\r
+  VOID                      *Session;\r
+  UDP_RX_TOKEN              *RxToken;\r
+  UDP_END_POINT             EndPoint;\r
+  NET_BUF                   *Netbuf;\r

 \r
 \r

-  Token = AllocatePool (sizeof (UDP_TX_TOKEN) +\r
-                           sizeof (EFI_UDP4_FRAGMENT_DATA) * (Packet->BlockOpNum - 1));\r
+  RxToken = (UDP_RX_TOKEN *) Context;\r

 \r
 \r

-  if (Token == NULL) {\r
-    return NULL;\r
-  }\r
+  ZeroMem (&EndPoint, sizeof(UDP_END_POINT));\r

 \r
 \r

-  Token->Signature  = UDP_IO_TX_SIGNATURE;\r
-  InitializeListHead (&Token->Link);\r
+  ASSERT ((RxToken->Signature == UDP_IO_RX_SIGNATURE) &&\r
+          (RxToken == RxToken->UdpIo->RecvRequest));\r

 \r
 \r

-  Token->UdpIo      = UdpIo;\r
-  Token->CallBack   = CallBack;\r
-  Token->Packet     = Packet;\r
-  Token->Context    = Context;\r
+  ASSERT ((RxToken->UdpIo->UdpVersion == UDP_IO_UDP4_VERSION) ||\r
+          (RxToken->UdpIo->UdpVersion == UDP_IO_UDP6_VERSION));\r

 \r
 \r

-  UdpToken          = &(Token->UdpToken);\r
-  UdpToken->Status  = EFI_NOT_READY;\r
+  //\r
+  // Clear the receive request first in case that the caller\r
+  // wants to restart the receive in the callback.\r
+  //\r
+  RxToken->UdpIo->RecvRequest = NULL;\r
+\r
+  if (RxToken->UdpIo->UdpVersion == UDP_IO_UDP4_VERSION) {\r
+    Token  = &RxToken->Token.Udp4;\r
+    RxData = ((EFI_UDP4_COMPLETION_TOKEN *) Token)->Packet.RxData;\r
+    Status = ((EFI_UDP4_COMPLETION_TOKEN *) Token)->Status;\r
+  } else {\r
+    Token  = &RxToken->Token.Udp6;\r
+    RxData = ((EFI_UDP6_COMPLETION_TOKEN *) Token)->Packet.RxData;\r
+    Status = ((EFI_UDP6_COMPLETION_TOKEN *) Token)->Status;\r
+  }\r

 \r
 \r

-  Status = gBS->CreateEvent (\r
-                  EVT_NOTIFY_SIGNAL,\r
-                  TPL_NOTIFY,\r
-                  UdpIoOnDgramSent,\r
-                  Token,\r
-                  &UdpToken->Event\r
-                  );\r
+  if (EFI_ERROR (Status) || RxData == NULL) {\r
+    if (Status != EFI_ABORTED) {\r
+      //\r
+      // Invoke the CallBack only if the reception is not actively aborted.\r
+      //\r
+      RxToken->CallBack (NULL, NULL, Status, RxToken->Context);\r
+    }\r

 \r
 \r

-  if (EFI_ERROR (Status)) {\r
-    gBS->FreePool (Token);\r
-    return NULL;\r
+    UdpIoFreeRxToken (RxToken);\r
+    return;\r

   }\r
 \r
   }\r
 \r

-  UdpTxData                 = &Token->UdpTxData;\r
-  UdpToken->Packet.TxData   = UdpTxData;\r
+  //\r
+  // Build a NET_BUF from the UDP receive data, then deliver it up.\r
+  //\r
+  if (RxToken->UdpIo->UdpVersion == UDP_IO_UDP4_VERSION) {\r
+\r
+    Netbuf = NetbufFromExt (\r
+               (NET_FRAGMENT *)((EFI_UDP4_RECEIVE_DATA *) RxData)->FragmentTable,\r
+               ((EFI_UDP4_RECEIVE_DATA *) RxData)->FragmentCount,\r
+               0,\r
+               (UINT32) RxToken->HeadLen,\r
+               UdpIoRecycleDgram,\r
+               RxToken\r
+               );\r
+\r
+    if (Netbuf == NULL) {\r
+      gBS->SignalEvent (((EFI_UDP4_RECEIVE_DATA *) RxData)->RecycleSignal);\r
+      RxToken->CallBack (NULL, NULL, EFI_OUT_OF_RESOURCES, RxToken->Context);\r
+\r
+      UdpIoFreeRxToken (RxToken);\r
+      return;\r
+    }\r

 \r
 \r

-  UdpTxData->UdpSessionData = NULL;\r
-  UdpTxData->GatewayAddress = NULL;\r
+    Session             = &((EFI_UDP4_RECEIVE_DATA *) RxData)->UdpSession;\r
+    EndPoint.LocalPort  = ((EFI_UDP4_SESSION_DATA *) Session)->DestinationPort;\r
+    EndPoint.RemotePort = ((EFI_UDP4_SESSION_DATA *) Session)->SourcePort;\r
+\r
+    CopyMem (\r
+      &EndPoint.LocalAddr,\r
+      &((EFI_UDP4_SESSION_DATA *) Session)->DestinationAddress,\r
+      sizeof (EFI_IPv4_ADDRESS)\r
+      );\r
+\r
+    CopyMem (\r
+      &EndPoint.RemoteAddr,\r
+      &((EFI_UDP4_SESSION_DATA *) Session)->SourceAddress,\r
+      sizeof (EFI_IPv4_ADDRESS)\r
+      );\r
+\r
+    EndPoint.LocalAddr.Addr[0]  = NTOHL (EndPoint.LocalAddr.Addr[0]);\r
+    EndPoint.RemoteAddr.Addr[0] = NTOHL (EndPoint.RemoteAddr.Addr[0]);\r
+  } else {\r
+\r
+    Netbuf = NetbufFromExt (\r
+               (NET_FRAGMENT *)((EFI_UDP6_RECEIVE_DATA *) RxData)->FragmentTable,\r
+               ((EFI_UDP6_RECEIVE_DATA *) RxData)->FragmentCount,\r
+               0,\r
+               (UINT32) RxToken->HeadLen,\r
+               UdpIoRecycleDgram,\r
+               RxToken\r
+               );\r
+\r
+    if (Netbuf == NULL) {\r
+      gBS->SignalEvent (((EFI_UDP6_RECEIVE_DATA *) RxData)->RecycleSignal);\r
+      RxToken->CallBack (NULL, NULL, EFI_OUT_OF_RESOURCES, RxToken->Context);\r
+\r
+      UdpIoFreeRxToken (RxToken);\r
+      return;\r
+    }\r

 \r
 \r

-  if (EndPoint != NULL) {\r
-    Ip = HTONL (EndPoint->LocalAddr);\r
-    CopyMem (&Token->UdpSession.SourceAddress, &Ip, sizeof (EFI_IPv4_ADDRESS));\r
+    Session             = &((EFI_UDP6_RECEIVE_DATA *) RxData)->UdpSession;\r
+    EndPoint.LocalPort  = ((EFI_UDP6_SESSION_DATA *) Session)->DestinationPort;\r
+    EndPoint.RemotePort = ((EFI_UDP6_SESSION_DATA *) Session)->SourcePort;\r

 \r
 \r

-    Ip = HTONL (EndPoint->RemoteAddr);\r
-    CopyMem (&Token->UdpSession.DestinationAddress, &Ip, sizeof (EFI_IPv4_ADDRESS));\r
+    CopyMem (\r
+      &EndPoint.LocalAddr,\r
+      &((EFI_UDP6_SESSION_DATA *) Session)->DestinationAddress,\r
+      sizeof (EFI_IPv6_ADDRESS)\r
+      );\r

 \r
 \r

-    Token->UdpSession.SourcePort      = EndPoint->LocalPort;\r
-    Token->UdpSession.DestinationPort = EndPoint->RemotePort;\r
-    UdpTxData->UdpSessionData         = &Token->UdpSession;\r
-  }\r
+    CopyMem (\r
+      &EndPoint.RemoteAddr,\r
+      &((EFI_UDP6_SESSION_DATA *) Session)->SourceAddress,\r
+      sizeof (EFI_IPv6_ADDRESS)\r
+      );\r

 \r
 \r

-  if (Gateway != 0) {\r
-    Ip = HTONL (Gateway);\r
-    CopyMem (&Token->Gateway, &Ip, sizeof (EFI_IPv4_ADDRESS));\r
-\r
-    UdpTxData->GatewayAddress = &Token->Gateway;\r
+    Ip6Swap128 (&EndPoint.LocalAddr.v6);\r
+    Ip6Swap128 (&EndPoint.RemoteAddr.v6);\r

   }\r
 \r
   }\r
 \r

-  UdpTxData->DataLength = Packet->TotalSize;\r
-  Count                 = Packet->BlockOpNum;\r
-  NetbufBuildExt (Packet, (NET_FRAGMENT *) UdpTxData->FragmentTable, &Count);\r
-  UdpTxData->FragmentCount = Count;\r
-\r
-  return Token;\r
+  RxToken->CallBack (Netbuf, &EndPoint, EFI_SUCCESS, RxToken->Context);\r

 }\r
 \r
 }\r
 \r

-\r

 /**\r
 /**\r

-  Free a UDP_TX_TOKEN. The event is closed and memory released.\r
-\r
-  @param  Token                 The UDP_TX_TOKEN to release.\r
+  Request UdpIoOnDgramRcvdDpc() as a DPC at TPL_CALLBACK.\r

 \r
 \r

-  @return None\r
+  @param[in]  Event                 The UDP receive request event.\r
+  @param[in]  Context               The UDP RX token.\r

 \r
 **/\r
 VOID\r
 \r
 **/\r
 VOID\r

-UdpIoFreeTxToken (\r
-  IN UDP_TX_TOKEN           *Token\r
+EFIAPI\r
+UdpIoOnDgramRcvd (\r
+  IN EFI_EVENT              Event,\r
+  IN VOID                   *Context\r

   )\r
 {\r
   )\r
 {\r

-  gBS->CloseEvent (Token->UdpToken.Event);\r
-  gBS->FreePool (Token);\r
+  //\r
+  // Request UdpIoOnDgramRcvdDpc as a DPC at TPL_CALLBACK\r
+  //\r
+  QueueDpc (TPL_CALLBACK, UdpIoOnDgramRcvdDpc, Context);\r

 }\r
 \r
 }\r
 \r

-\r

 /**\r
   Create a UDP_RX_TOKEN to wrap the request.\r
 \r
 /**\r
   Create a UDP_RX_TOKEN to wrap the request.\r
 \r

-  @param  UdpIo                 The UdpIo to receive packets from\r
-  @param  CallBack              The function to call when receive finished.\r
-  @param  Context               The opaque parameter to the CallBack\r
-  @param  HeadLen               The head length to reserver for the packet.\r
+  @param[in]  UdpIo                 The UdpIo to receive packets from.\r
+  @param[in]  CallBack              The function to call when receive finished.\r
+  @param[in]  Context               The opaque parameter to the CallBack.\r
+  @param[in]  HeadLen               The head length to reserver for the packet.\r

 \r
 \r

-  @return The Wrapped request or NULL if failed to allocate resources.\r
+  @return The Wrapped request or NULL if failed to allocate resources or some errors happened.\r

 \r
 **/\r
 UDP_RX_TOKEN *\r
 UdpIoCreateRxToken (\r
 \r
 **/\r
 UDP_RX_TOKEN *\r
 UdpIoCreateRxToken (\r

-  IN UDP_IO_PORT            *UdpIo,\r
+  IN UDP_IO                 *UdpIo,\r

   IN UDP_IO_CALLBACK        CallBack,\r
   IN VOID                   *Context,\r
   IN UINT32                 HeadLen\r
   IN UDP_IO_CALLBACK        CallBack,\r
   IN VOID                   *Context,\r
   IN UINT32                 HeadLen\r


@@ -188,6 +335,9 @@ UdpIoCreateRxToken (
   UDP_RX_TOKEN              *Token;\r
   EFI_STATUS                Status;\r
 \r
   UDP_RX_TOKEN              *Token;\r
   EFI_STATUS                Status;\r
 \r

+  ASSERT ((UdpIo->UdpVersion == UDP_IO_UDP4_VERSION) ||\r
+          (UdpIo->UdpVersion == UDP_IO_UDP6_VERSION));\r
+\r

   Token = AllocatePool (sizeof (UDP_RX_TOKEN));\r
 \r
   if (Token == NULL) {\r
   Token = AllocatePool (sizeof (UDP_RX_TOKEN));\r
 \r
   if (Token == NULL) {\r


@@ -200,160 +350,396 @@ UdpIoCreateRxToken (
   Token->Context                = Context;\r
   Token->HeadLen                = HeadLen;\r
 \r
   Token->Context                = Context;\r
   Token->HeadLen                = HeadLen;\r
 \r

-  Token->UdpToken.Status        = EFI_NOT_READY;\r
-  Token->UdpToken.Packet.RxData = NULL;\r
+  if (UdpIo->UdpVersion == UDP_IO_UDP4_VERSION) {\r
+\r
+    Token->Token.Udp4.Status        = EFI_NOT_READY;\r
+    Token->Token.Udp4.Packet.RxData = NULL;\r
+\r
+    Status = gBS->CreateEvent (\r
+                    EVT_NOTIFY_SIGNAL,\r
+                    TPL_NOTIFY,\r
+                    UdpIoOnDgramRcvd,\r
+                    Token,\r
+                    &Token->Token.Udp4.Event\r
+                    );\r
+    } else {\r
+\r
+    Token->Token.Udp6.Status        = EFI_NOT_READY;\r
+    Token->Token.Udp6.Packet.RxData = NULL;\r
+\r
+    Status = gBS->CreateEvent (\r
+                    EVT_NOTIFY_SIGNAL,\r
+                    TPL_NOTIFY,\r
+                    UdpIoOnDgramRcvd,\r
+                    Token,\r
+                    &Token->Token.Udp6.Event\r
+                    );\r
+  }\r

 \r
 \r

-  Status = gBS->CreateEvent (\r
-                  EVT_NOTIFY_SIGNAL,\r
-                  TPL_NOTIFY,\r
-                  UdpIoOnDgramRcvd,\r
-                  Token,\r
-                  &Token->UdpToken.Event\r
-                  );\r

 \r
   if (EFI_ERROR (Status)) {\r
 \r
   if (EFI_ERROR (Status)) {\r

-    gBS->FreePool (Token);\r
+    FreePool (Token);\r

     return NULL;\r
   }\r
 \r
   return Token;\r
 }\r
 \r
     return NULL;\r
   }\r
 \r
   return Token;\r
 }\r
 \r

-\r

 /**\r
 /**\r

-  Free a receive request wrap.\r
+  Wrap a transmit request into a new created UDP_TX_TOKEN.\r

 \r
 \r

-  @param  Token                 The receive request to release.\r
+  @param[in]  UdpIo                 The UdpIo to send packet to.\r
+  @param[in]  Packet                The user's packet.\r
+  @param[in]  EndPoint              The local and remote access point.\r
+  @param[in]  Gateway               The overrided next hop.\r
+  @param[in]  CallBack              The function to call when transmission completed.\r
+  @param[in]  Context               The opaque parameter to the call back.\r

 \r
 \r

-  @return None\r
+  @return The wrapped transmission request or NULL if failed to allocate resources\r
+          or for some errors.\r

 \r
 **/\r
 \r
 **/\r

-VOID\r
-UdpIoFreeRxToken (\r
-  IN UDP_RX_TOKEN           *Token\r
+UDP_TX_TOKEN *\r
+UdpIoCreateTxToken (\r
+  IN UDP_IO                 *UdpIo,\r
+  IN NET_BUF                *Packet,\r
+  IN UDP_END_POINT          *EndPoint OPTIONAL,\r
+  IN EFI_IP_ADDRESS         *Gateway  OPTIONAL,\r
+  IN UDP_IO_CALLBACK        CallBack,\r
+  IN VOID                   *Context\r

   )\r
 {\r
   )\r
 {\r

-  gBS->CloseEvent (Token->UdpToken.Event);\r
-  gBS->FreePool (Token);\r
-}\r
+  UDP_TX_TOKEN              *TxToken;\r
+  VOID                      *Token;\r
+  VOID                      *Data;\r
+  EFI_STATUS                Status;\r
+  UINT32                    Count;\r
+  UINTN                     Size;\r
+  IP4_ADDR                  Ip;\r
+\r
+  ASSERT (Packet != NULL);\r
+  ASSERT ((UdpIo->UdpVersion == UDP_IO_UDP4_VERSION) ||\r
+          (UdpIo->UdpVersion == UDP_IO_UDP6_VERSION));\r
+\r
+  if (UdpIo->UdpVersion == UDP_IO_UDP4_VERSION) {\r
+    Size = sizeof (UDP_TX_TOKEN) + sizeof (EFI_UDP4_FRAGMENT_DATA) * (Packet->BlockOpNum - 1);\r
+  } else {\r
+    Size = sizeof (UDP_TX_TOKEN) + sizeof (EFI_UDP6_FRAGMENT_DATA) * (Packet->BlockOpNum - 1);\r
+  }\r
+\r
+  TxToken = AllocatePool (Size);\r
+\r
+  if (TxToken == NULL) {\r
+    return NULL;\r
+  }\r
+\r
+  TxToken->Signature = UDP_IO_TX_SIGNATURE;\r
+  InitializeListHead (&TxToken->Link);\r
+\r
+  TxToken->UdpIo     = UdpIo;\r
+  TxToken->CallBack  = CallBack;\r
+  TxToken->Packet    = Packet;\r
+  TxToken->Context   = Context;\r
+\r
+  Token              = &(TxToken->Token);\r
+  Count              = Packet->BlockOpNum;\r
+\r
+  if (UdpIo->UdpVersion == UDP_IO_UDP4_VERSION) {\r
+\r
+    ((EFI_UDP4_COMPLETION_TOKEN *) Token)->Status = EFI_NOT_READY;\r
+\r
+    Status = gBS->CreateEvent (\r
+                    EVT_NOTIFY_SIGNAL,\r
+                    TPL_NOTIFY,\r
+                    UdpIoOnDgramSent,\r
+                    TxToken,\r
+                    &((EFI_UDP4_COMPLETION_TOKEN *) Token)->Event\r
+                    );\r
+\r
+    if (EFI_ERROR (Status)) {\r
+      FreePool (TxToken);\r
+      return NULL;\r
+    }\r

 \r
 \r

+    Data = &(TxToken->Data.Udp4);\r
+    ((EFI_UDP4_COMPLETION_TOKEN *) Token)->Packet.TxData = Data;\r
+\r
+    ((EFI_UDP4_TRANSMIT_DATA *) Data)->UdpSessionData    = NULL;\r
+    ((EFI_UDP4_TRANSMIT_DATA *) Data)->GatewayAddress    = NULL;\r
+    ((EFI_UDP4_TRANSMIT_DATA *) Data)->DataLength        = Packet->TotalSize;\r
+\r
+    NetbufBuildExt (\r
+      Packet,\r
+      (NET_FRAGMENT *)((EFI_UDP4_TRANSMIT_DATA *) Data)->FragmentTable,\r
+      &Count\r
+      );\r
+\r
+    ((EFI_UDP4_TRANSMIT_DATA *) Data)->FragmentCount     = Count;\r
+\r
+    if (EndPoint != NULL) {\r
+      Ip = HTONL (EndPoint->LocalAddr.Addr[0]);\r
+      CopyMem (\r
+        &TxToken->Session.Udp4.SourceAddress,\r
+        &Ip,\r
+        sizeof (EFI_IPv4_ADDRESS)\r
+        );\r
+\r
+      Ip = HTONL (EndPoint->RemoteAddr.Addr[0]);\r
+      CopyMem (\r
+        &TxToken->Session.Udp4.DestinationAddress,\r
+        &Ip,\r
+        sizeof (EFI_IPv4_ADDRESS)\r
+        );\r
+\r
+      TxToken->Session.Udp4.SourcePort                   = EndPoint->LocalPort;\r
+      TxToken->Session.Udp4.DestinationPort              = EndPoint->RemotePort;\r
+      ((EFI_UDP4_TRANSMIT_DATA *) Data)->UdpSessionData  = &(TxToken->Session.Udp4);\r
+    }\r
+\r
+    if (Gateway != NULL && (Gateway->Addr[0] != 0)) {\r
+      Ip = HTONL (Gateway->Addr[0]);\r
+      CopyMem (&TxToken->Gateway, &Ip, sizeof (EFI_IPv4_ADDRESS));\r
+      ((EFI_UDP4_TRANSMIT_DATA *) Data)->GatewayAddress = &TxToken->Gateway;\r
+    }\r
+\r
+  } else {\r
+\r
+    ((EFI_UDP6_COMPLETION_TOKEN *) Token)->Status = EFI_NOT_READY;\r
+\r
+    Status = gBS->CreateEvent (\r
+                    EVT_NOTIFY_SIGNAL,\r
+                    TPL_NOTIFY,\r
+                    UdpIoOnDgramSent,\r
+                    TxToken,\r
+                    &((EFI_UDP6_COMPLETION_TOKEN *) Token)->Event\r
+                    );\r
+\r
+    if (EFI_ERROR (Status)) {\r
+      FreePool (TxToken);\r
+      return NULL;\r
+    }\r
+\r
+    Data = &(TxToken->Data.Udp6);\r
+    ((EFI_UDP6_COMPLETION_TOKEN *) Token)->Packet.TxData  = Data;\r
+    ((EFI_UDP6_TRANSMIT_DATA *) Data)->UdpSessionData     = NULL;\r
+    ((EFI_UDP6_TRANSMIT_DATA *) Data)->DataLength         = Packet->TotalSize;\r
+\r
+    NetbufBuildExt (\r
+      Packet,\r
+      (NET_FRAGMENT *)((EFI_UDP6_TRANSMIT_DATA *) Data)->FragmentTable,\r
+      &Count\r
+      );\r
+\r
+    ((EFI_UDP6_TRANSMIT_DATA *) Data)->FragmentCount      = Count;\r
+\r
+    if (EndPoint != NULL) {\r
+      CopyMem (\r
+        &TxToken->Session.Udp6.SourceAddress,\r
+        &EndPoint->LocalAddr.v6,\r
+        sizeof(EFI_IPv6_ADDRESS)\r
+        );\r
+\r
+      CopyMem (\r
+        &TxToken->Session.Udp6.DestinationAddress,\r
+        &EndPoint->RemoteAddr.v6,\r
+        sizeof(EFI_IPv6_ADDRESS)\r
+        );\r
+\r
+      TxToken->Session.Udp6.SourcePort                   = EndPoint->LocalPort;\r
+      TxToken->Session.Udp6.DestinationPort              = EndPoint->RemotePort;\r
+      ((EFI_UDP6_TRANSMIT_DATA *) Data)->UdpSessionData  = &(TxToken->Session.Udp6);\r
+    }\r
+  }\r
+\r
+  return TxToken;\r
+}\r

 \r
 /**\r
 \r
 /**\r

-  Create a UDP IO port to access the UDP service. It will\r
-  create and configure a UDP child.\r
+  Creates a UDP_IO to access the UDP service. It creates and configures\r
+  a UDP child.\r
+\r
+  It locates the UDP service binding prototype on the Controller parameter\r
+  uses the UDP service binding prototype to create a UDP child (also known as\r
+  a UDP instance) configures the UDP child by calling Configure function prototype.\r
+  Any failures in creating or configuring the UDP child return NULL for failure.\r

 \r
 \r

-  @param  Controller            The controller that has the UDP service binding\r
-                                protocol installed.\r
-  @param  Image                 The image handle for the driver.\r
-  @param  Configure             The function to configure the created UDP child\r
-  @param  Context               The opaque parameter for the Configure funtion.\r
+  @param[in]  Controller            The controller that has the UDP service binding.\r
+                                    protocol installed.\r
+  @param[in]  ImageHandle           The image handle for the driver.\r
+  @param[in]  Configure             The function to configure the created UDP child.\r
+  @param[in]  UdpVersion            The UDP protocol version, UDP4 or UDP6.\r
+  @param[in]  Context               The opaque parameter for the Configure funtion.\r

 \r
 \r

-  @return A point to just created UDP IO port or NULL if failed.\r
+  @return Newly-created UDP_IO or NULL if failed.\r

 \r
 **/\r
 \r
 **/\r

-UDP_IO_PORT *\r
+UDP_IO *\r

 EFIAPI\r
 EFIAPI\r

-UdpIoCreatePort (\r
+UdpIoCreateIo (\r

   IN  EFI_HANDLE            Controller,\r
   IN  EFI_HANDLE            Controller,\r

-  IN  EFI_HANDLE            Image,\r
+  IN  EFI_HANDLE            ImageHandle,\r

   IN  UDP_IO_CONFIG         Configure,\r
   IN  UDP_IO_CONFIG         Configure,\r

+  IN  UINT8                 UdpVersion,\r

   IN  VOID                  *Context\r
   )\r
 {\r
   IN  VOID                  *Context\r
   )\r
 {\r

-  UDP_IO_PORT               *UdpIo;\r
+  UDP_IO                    *UdpIo;\r

   EFI_STATUS                Status;\r
 \r
   ASSERT (Configure != NULL);\r
   EFI_STATUS                Status;\r
 \r
   ASSERT (Configure != NULL);\r

+  ASSERT ((UdpVersion == UDP_IO_UDP4_VERSION) || (UdpVersion == UDP_IO_UDP6_VERSION));\r

 \r
 \r

-  UdpIo = AllocatePool (sizeof (UDP_IO_PORT));\r
+  UdpIo = AllocatePool (sizeof (UDP_IO));\r

 \r
   if (UdpIo == NULL) {\r
     return NULL;\r
   }\r
 \r
 \r
   if (UdpIo == NULL) {\r
     return NULL;\r
   }\r
 \r

+  UdpIo->UdpVersion   = UdpVersion;\r

   UdpIo->Signature    = UDP_IO_SIGNATURE;\r
   InitializeListHead (&UdpIo->Link);\r
   UdpIo->RefCnt       = 1;\r
 \r
   UdpIo->Controller   = Controller;\r
   UdpIo->Signature    = UDP_IO_SIGNATURE;\r
   InitializeListHead (&UdpIo->Link);\r
   UdpIo->RefCnt       = 1;\r
 \r
   UdpIo->Controller   = Controller;\r

-  UdpIo->Image        = Image;\r
+  UdpIo->Image        = ImageHandle;\r

 \r
   InitializeListHead (&UdpIo->SentDatagram);\r
   UdpIo->RecvRequest  = NULL;\r
   UdpIo->UdpHandle    = NULL;\r
 \r
 \r
   InitializeListHead (&UdpIo->SentDatagram);\r
   UdpIo->RecvRequest  = NULL;\r
   UdpIo->UdpHandle    = NULL;\r
 \r

-  //\r
-  // Create a UDP child then open and configure it\r
-  //\r
-  Status = NetLibCreateServiceChild (\r
-             Controller,\r
-             Image,\r
-             &gEfiUdp4ServiceBindingProtocolGuid,\r
-             &UdpIo->UdpHandle\r
-             );\r
+  if (UdpVersion == UDP_IO_UDP4_VERSION) {\r
+    //\r
+    // Create a UDP child then open and configure it\r
+    //\r
+    Status = NetLibCreateServiceChild (\r
+               Controller,\r
+               ImageHandle,\r
+               &gEfiUdp4ServiceBindingProtocolGuid,\r
+               &UdpIo->UdpHandle\r
+               );\r
+\r
+    if (EFI_ERROR (Status)) {\r
+      goto FREE_MEM;\r
+    }\r

 \r
 \r

-  if (EFI_ERROR (Status)) {\r
-    goto FREE_MEM;\r
-  }\r
+    Status = gBS->OpenProtocol (\r
+                    UdpIo->UdpHandle,\r
+                    &gEfiUdp4ProtocolGuid,\r
+                    (VOID **) &UdpIo->Protocol.Udp4,\r
+                    ImageHandle,\r
+                    Controller,\r
+                    EFI_OPEN_PROTOCOL_BY_DRIVER\r
+                    );\r
+\r
+    if (EFI_ERROR (Status)) {\r
+      goto FREE_CHILD;\r
+    }\r

 \r
 \r

-  Status = gBS->OpenProtocol (\r
-                  UdpIo->UdpHandle,\r
-                  &gEfiUdp4ProtocolGuid,\r
-                  (VOID **) &UdpIo->Udp,\r
-                  Image,\r
-                  Controller,\r
-                  EFI_OPEN_PROTOCOL_BY_DRIVER\r
-                  );\r
+    if (EFI_ERROR (Configure (UdpIo, Context))) {\r
+      goto CLOSE_PROTOCOL;\r
+    }\r

 \r
 \r

-  if (EFI_ERROR (Status)) {\r
-    goto FREE_CHILD;\r
-  }\r
+    Status = UdpIo->Protocol.Udp4->GetModeData (\r
+                                     UdpIo->Protocol.Udp4,\r
+                                     NULL,\r
+                                     NULL,\r
+                                     NULL,\r
+                                     &UdpIo->SnpMode\r
+                                     );\r

 \r
 \r

-  if (EFI_ERROR (Configure (UdpIo, Context))) {\r
-    goto CLOSE_PROTOCOL;\r
-  }\r
+    if (EFI_ERROR (Status)) {\r
+      goto CLOSE_PROTOCOL;\r
+    }\r

 \r
 \r

-  Status = UdpIo->Udp->GetModeData (UdpIo->Udp, NULL, NULL, NULL, &UdpIo->SnpMode);\r
+  } else {\r

 \r
 \r

-  if (EFI_ERROR (Status)) {\r
-    goto CLOSE_PROTOCOL;\r
+    Status = NetLibCreateServiceChild (\r
+               Controller,\r
+               ImageHandle,\r
+               &gEfiUdp6ServiceBindingProtocolGuid,\r
+               &UdpIo->UdpHandle\r
+               );\r
+\r
+    if (EFI_ERROR (Status)) {\r
+      goto FREE_MEM;\r
+    }\r
+\r
+    Status = gBS->OpenProtocol (\r
+                    UdpIo->UdpHandle,\r
+                    &gEfiUdp6ProtocolGuid,\r
+                    (VOID **) &UdpIo->Protocol.Udp6,\r
+                    ImageHandle,\r
+                    Controller,\r
+                    EFI_OPEN_PROTOCOL_BY_DRIVER\r
+                    );\r
+\r
+    if (EFI_ERROR (Status)) {\r
+      goto FREE_CHILD;\r
+    }\r
+\r
+    if (EFI_ERROR (Configure (UdpIo, Context))) {\r
+      goto CLOSE_PROTOCOL;\r
+    }\r
+\r
+    Status = UdpIo->Protocol.Udp6->GetModeData (\r
+                                     UdpIo->Protocol.Udp6,\r
+                                     NULL,\r
+                                     NULL,\r
+                                     NULL,\r
+                                     &UdpIo->SnpMode\r
+                                     );\r
+\r
+    if (EFI_ERROR (Status)) {\r
+      goto CLOSE_PROTOCOL;\r
+    }\r

   }\r
 \r
   return UdpIo;\r
 \r
 CLOSE_PROTOCOL:\r
   }\r
 \r
   return UdpIo;\r
 \r
 CLOSE_PROTOCOL:\r

-  gBS->CloseProtocol (UdpIo->UdpHandle, &gEfiUdp4ProtocolGuid, Image, Controller);\r
+  if (UdpVersion == UDP_IO_UDP4_VERSION) {\r
+    gBS->CloseProtocol (UdpIo->UdpHandle, &gEfiUdp4ProtocolGuid, ImageHandle, Controller);\r
+  } else {\r
+    gBS->CloseProtocol (UdpIo->UdpHandle, &gEfiUdp6ProtocolGuid, ImageHandle, Controller);\r
+  }\r

 \r
 FREE_CHILD:\r
 \r
 FREE_CHILD:\r

-  NetLibDestroyServiceChild (\r
-    Controller,\r
-    Image,\r
-    &gEfiUdp4ServiceBindingProtocolGuid,\r
-    UdpIo->UdpHandle\r
-    );\r
+  if (UdpVersion == UDP_IO_UDP4_VERSION) {\r
+    NetLibDestroyServiceChild (\r
+      Controller,\r
+      ImageHandle,\r
+      &gEfiUdp4ServiceBindingProtocolGuid,\r
+      UdpIo->UdpHandle\r
+      );\r
+  } else {\r
+    NetLibDestroyServiceChild (\r
+      Controller,\r
+      ImageHandle,\r
+      &gEfiUdp6ServiceBindingProtocolGuid,\r
+      UdpIo->UdpHandle\r
+      );\r
+  }\r

 \r
 FREE_MEM:\r
 \r
 FREE_MEM:\r

-  gBS->FreePool (UdpIo);\r
+  FreePool (UdpIo);\r

   return NULL;\r
 }\r
 \r
   return NULL;\r
 }\r
 \r

-\r

 /**\r
 /**\r

-  Cancel all the sent datagram that pass the selection of ToCancel.\r
+  Cancel all the sent datagram that pass the selection criteria of ToCancel.\r

   If ToCancel is NULL, all the datagrams are cancelled.\r
 \r
   If ToCancel is NULL, all the datagrams are cancelled.\r
 \r

-  @param  UdpIo                 The UDP IO port to cancel packet\r
-  @param  IoStatus              The IoStatus to return to the packet owners.\r
-  @param  ToCancel              The select funtion to test whether to cancel this\r
-                                packet or not.\r
-  @param  Context               The opaque parameter to the ToCancel.\r
-\r
-  @return None\r
+  @param[in]  UdpIo                 The UDP_IO to cancel packet.\r
+  @param[in]  IoStatus              The IoStatus to return to the packet owners.\r
+  @param[in]  ToCancel              The select funtion to test whether to cancel this\r
+                                    packet or not.\r
+  @param[in]  Context               The opaque parameter to the ToCancel.\r

 \r
 **/\r
 VOID\r
 \r
 **/\r
 VOID\r

+EFIAPI\r

 UdpIoCancelDgrams (\r
 UdpIoCancelDgrams (\r

-  IN UDP_IO_PORT            *UdpIo,\r
+  IN UDP_IO                 *UdpIo,\r

   IN EFI_STATUS             IoStatus,\r
   IN UDP_IO_TO_CANCEL       ToCancel,        OPTIONAL\r
   IN VOID                   *Context\r
   IN EFI_STATUS             IoStatus,\r
   IN UDP_IO_TO_CANCEL       ToCancel,        OPTIONAL\r
   IN VOID                   *Context\r


@@ -361,34 +747,45 @@ UdpIoCancelDgrams (
 {\r
   LIST_ENTRY                *Entry;\r
   LIST_ENTRY                *Next;\r
 {\r
   LIST_ENTRY                *Entry;\r
   LIST_ENTRY                *Next;\r

-  UDP_TX_TOKEN              *Token;\r
+  UDP_TX_TOKEN              *TxToken;\r
+\r
+  ASSERT ((UdpIo->UdpVersion == UDP_IO_UDP4_VERSION) ||\r
+          (UdpIo->UdpVersion == UDP_IO_UDP6_VERSION));\r

 \r
   NET_LIST_FOR_EACH_SAFE (Entry, Next, &UdpIo->SentDatagram) {\r
 \r
   NET_LIST_FOR_EACH_SAFE (Entry, Next, &UdpIo->SentDatagram) {\r

-    Token = NET_LIST_USER_STRUCT (Entry, UDP_TX_TOKEN, Link);\r
+    TxToken = NET_LIST_USER_STRUCT (Entry, UDP_TX_TOKEN, Link);\r

 \r
 \r

-    if ((ToCancel == NULL) || (ToCancel (Token, Context))) {\r
-      UdpIo->Udp->Cancel (UdpIo->Udp, &Token->UdpToken);\r
+    if ((ToCancel == NULL) || (ToCancel (TxToken, Context))) {\r
+\r
+      if (UdpIo->UdpVersion == UDP_IO_UDP4_VERSION) {\r
+        UdpIo->Protocol.Udp4->Cancel (UdpIo->Protocol.Udp4, &TxToken->Token.Udp4);\r
+      } else {\r
+        UdpIo->Protocol.Udp6->Cancel (UdpIo->Protocol.Udp6, &TxToken->Token.Udp6);\r
+      }\r

     }\r
   }\r
 }\r
 \r
     }\r
   }\r
 }\r
 \r

-\r

 /**\r
 /**\r

-  Free the UDP IO port and all its related resources including\r
-  all the transmitted packet.\r
+  Free the UDP_IO and all its related resources.\r

 \r
 \r

-  @param  UdpIo                 The UDP IO port to free.\r
+  The function will cancel all sent datagram and receive request.\r

 \r
 \r

-  @retval EFI_SUCCESS           The UDP IO port is freed.\r
+  @param[in]  UdpIo             The UDP_IO to free.\r
+\r
+  @retval EFI_SUCCESS           The UDP_IO is freed.\r

 \r
 **/\r
 EFI_STATUS\r
 EFIAPI\r
 \r
 **/\r
 EFI_STATUS\r
 EFIAPI\r

-UdpIoFreePort (\r
-  IN  UDP_IO_PORT           *UdpIo\r
+UdpIoFreeIo (\r
+  IN  UDP_IO           *UdpIo\r

   )\r
 {\r
   )\r
 {\r

-  UDP_RX_TOKEN  *RxToken;\r
+  UDP_RX_TOKEN         *RxToken;\r
+\r
+  ASSERT ((UdpIo->UdpVersion == UDP_IO_UDP4_VERSION) ||\r
+          (UdpIo->UdpVersion == UDP_IO_UDP6_VERSION));\r

 \r
   //\r
   // Cancel all the sent datagram and receive requests. The\r
 \r
   //\r
   // Cancel all the sent datagram and receive requests. The\r


@@ -400,129 +797,121 @@ UdpIoFreePort (
   //\r
   UdpIoCancelDgrams (UdpIo, EFI_ABORTED, NULL, NULL);\r
 \r
   //\r
   UdpIoCancelDgrams (UdpIo, EFI_ABORTED, NULL, NULL);\r
 \r

-  if ((RxToken = UdpIo->RecvRequest) != NULL) {\r
-    UdpIo->Udp->Cancel (UdpIo->Udp, &RxToken->UdpToken);\r
-  }\r
+  if (UdpIo->UdpVersion == UDP_IO_UDP4_VERSION) {\r

 \r
 \r

-  //\r
-  // Close then destory the UDP child\r
-  //\r
-  gBS->CloseProtocol (\r
-         UdpIo->UdpHandle,\r
-         &gEfiUdp4ProtocolGuid,\r
-         UdpIo->Image,\r
-         UdpIo->Controller\r
-         );\r
-\r
-  NetLibDestroyServiceChild (\r
-    UdpIo->Controller,\r
-    UdpIo->Image,\r
-    &gEfiUdp4ServiceBindingProtocolGuid,\r
-    UdpIo->UdpHandle\r
-    );\r
+    if ((RxToken = UdpIo->RecvRequest) != NULL) {\r
+      UdpIo->Protocol.Udp4->Cancel (UdpIo->Protocol.Udp4, &RxToken->Token.Udp4);\r
+    }\r
+\r
+    //\r
+    // Close then destory the Udp4 child\r
+    //\r
+    gBS->CloseProtocol (\r
+           UdpIo->UdpHandle,\r
+           &gEfiUdp4ProtocolGuid,\r
+           UdpIo->Image,\r
+           UdpIo->Controller\r
+           );\r
+\r
+    NetLibDestroyServiceChild (\r
+      UdpIo->Controller,\r
+      UdpIo->Image,\r
+      &gEfiUdp4ServiceBindingProtocolGuid,\r
+      UdpIo->UdpHandle\r
+      );\r
+\r
+  } else {\r
+\r
+    if ((RxToken = UdpIo->RecvRequest) != NULL) {\r
+      UdpIo->Protocol.Udp6->Cancel (UdpIo->Protocol.Udp6, &RxToken->Token.Udp6);\r
+    }\r
+\r
+    //\r
+    // Close then destory the Udp6 child\r
+    //\r
+    gBS->CloseProtocol (\r
+           UdpIo->UdpHandle,\r
+           &gEfiUdp6ProtocolGuid,\r
+           UdpIo->Image,\r
+           UdpIo->Controller\r
+           );\r
+\r
+    NetLibDestroyServiceChild (\r
+      UdpIo->Controller,\r
+      UdpIo->Image,\r
+      &gEfiUdp6ServiceBindingProtocolGuid,\r
+      UdpIo->UdpHandle\r
+      );\r
+    }\r

 \r
   if (!IsListEmpty(&UdpIo->Link)) {\r
     RemoveEntryList (&UdpIo->Link);\r
   }\r
 \r
 \r
   if (!IsListEmpty(&UdpIo->Link)) {\r
     RemoveEntryList (&UdpIo->Link);\r
   }\r
 \r

-  gBS->FreePool (UdpIo);\r
+  FreePool (UdpIo);\r

   return EFI_SUCCESS;\r
 }\r
 \r
 \r
 /**\r
   return EFI_SUCCESS;\r
 }\r
 \r
 \r
 /**\r

-  Clean up the UDP IO port. It will release all the transmitted\r
-  datagrams and receive request. It will also configure NULL the\r
-  UDP child.\r
+  Clean up the UDP_IO without freeing it. The function is called when\r
+  user wants to re-use the UDP_IO later.\r

 \r
 \r

-  @param  UdpIo                 UDP IO port to clean up.\r
+  It will release all the transmitted datagrams and receive request. It will\r
+  also configure NULL for the UDP instance.\r

 \r
 \r

-  @return None\r
+  @param[in]  UdpIo                 The UDP_IO to clean up.\r

 \r
 **/\r
 VOID\r
 EFIAPI\r
 \r
 **/\r
 VOID\r
 EFIAPI\r

-UdpIoCleanPort (\r
-  IN  UDP_IO_PORT           *UdpIo\r
+UdpIoCleanIo (\r
+  IN  UDP_IO                *UdpIo\r

   )\r
 {\r
   UDP_RX_TOKEN              *RxToken;\r
 \r
   )\r
 {\r
   UDP_RX_TOKEN              *RxToken;\r
 \r

+  ASSERT ((UdpIo->UdpVersion == UDP_IO_UDP4_VERSION) ||\r
+          (UdpIo->UdpVersion == UDP_IO_UDP6_VERSION));\r
+\r

   //\r
   // Cancel all the sent datagram and receive requests.\r
   //\r
   UdpIoCancelDgrams (UdpIo, EFI_ABORTED, NULL, NULL);\r
 \r
   //\r
   // Cancel all the sent datagram and receive requests.\r
   //\r
   UdpIoCancelDgrams (UdpIo, EFI_ABORTED, NULL, NULL);\r
 \r

-  if ((RxToken = UdpIo->RecvRequest) != NULL) {\r
-    UdpIo->Udp->Cancel (UdpIo->Udp, &RxToken->UdpToken);\r
-  }\r
-\r
-  UdpIo->Udp->Configure (UdpIo->Udp, NULL);\r
-}\r
-\r
-\r
-/**\r
-  The callback function when the packet is sent by UDP.\r
-  It will remove the packet from the local list then call\r
-  the packet owner's callback function.\r
-\r
-  @param  Context               The UDP TX Token.\r
-\r
-  @return None\r
-\r
-**/\r
-VOID\r
-EFIAPI\r
-UdpIoOnDgramSentDpc (\r
-  IN VOID                   *Context\r
-  )\r
-{\r
-  UDP_TX_TOKEN              *Token;\r
-\r
-  Token   = (UDP_TX_TOKEN *) Context;\r
-  ASSERT (Token->Signature == UDP_IO_TX_SIGNATURE);\r
-\r
-  RemoveEntryList (&Token->Link);\r
-  Token->CallBack (Token->Packet, NULL, Token->UdpToken.Status, Token->Context);\r
-\r
-  UdpIoFreeTxToken (Token);\r
-}\r
-\r
-/**\r
-  Request UdpIoOnDgramSentDpc as a DPC at TPL_CALLBACK.\r
+  if (UdpIo->UdpVersion == UDP_IO_UDP4_VERSION) {\r
+    if ((RxToken = UdpIo->RecvRequest) != NULL) {\r
+      UdpIo->Protocol.Udp4->Cancel (UdpIo->Protocol.Udp4, &RxToken->Token.Udp4);\r
+    }\r

 \r
 \r

-  @param  Event                 The event signalled.\r
-  @param  Context               The UDP TX Token.\r
+    UdpIo->Protocol.Udp4->Configure (UdpIo->Protocol.Udp4, NULL);\r

 \r
 \r

-  @return None\r
+  } else {\r
+    if ((RxToken = UdpIo->RecvRequest) != NULL) {\r
+      UdpIo->Protocol.Udp6->Cancel (UdpIo->Protocol.Udp6, &RxToken->Token.Udp6);\r
+    }\r

 \r
 \r

-**/\r
-VOID\r
-EFIAPI\r
-UdpIoOnDgramSent (\r
-  IN EFI_EVENT              Event,\r
-  IN VOID                   *Context\r
-  )\r
-{\r
-  //\r
-  // Request UdpIoOnDgramSentDpc as a DPC at TPL_CALLBACK\r
-  //\r
-  NetLibQueueDpc (TPL_CALLBACK, UdpIoOnDgramSentDpc, Context);\r
+    UdpIo->Protocol.Udp6->Configure (UdpIo->Protocol.Udp6, NULL);\r
+  }\r

 }\r
 \r
 }\r
 \r

-\r

 /**\r
 /**\r

-  Send a packet through the UDP IO port.\r
-\r
-  @param  UdpIo                 The UDP IO Port to send the packet through\r
-  @param  Packet                The packet to send\r
-  @param  EndPoint              The local and remote access point\r
-  @param  Gateway               The gateway to use\r
-  @param  CallBack              The call back function to call when packet is\r
-                                transmitted or failed.\r
-  @param  Context               The opque parameter to the CallBack\r
-\r
-  @retval EFI_OUT_OF_RESOURCES  Failed to allocate resource for the packet\r
+  Send a packet through the UDP_IO.\r
+\r
+  The packet will be wrapped in UDP_TX_TOKEN. Function Callback will be called\r
+  when the packet is sent. The optional parameter EndPoint overrides the default\r
+  address pair if specified.\r
+\r
+  @param[in]  UdpIo                 The UDP_IO to send the packet through.\r
+  @param[in]  Packet                The packet to send.\r
+  @param[in]  EndPoint              The local and remote access point. Override the\r
+                                    default address pair set during configuration.\r
+  @param[in]  Gateway               The gateway to use.\r
+  @param[in]  CallBack              The function being called when packet is\r
+                                    transmitted or failed.\r
+  @param[in]  Context               The opaque parameter passed to CallBack.\r
+\r
+  @retval EFI_OUT_OF_RESOURCES  Failed to allocate resource for the packet.\r

   @retval EFI_SUCCESS           The packet is successfully delivered to UDP  for\r
                                 transmission.\r
 \r
   @retval EFI_SUCCESS           The packet is successfully delivered to UDP  for\r
                                 transmission.\r
 \r


@@ -530,20 +919,23 @@ UdpIoOnDgramSent (
 EFI_STATUS\r
 EFIAPI\r
 UdpIoSendDatagram (\r
 EFI_STATUS\r
 EFIAPI\r
 UdpIoSendDatagram (\r

-  IN  UDP_IO_PORT           *UdpIo,\r
+  IN  UDP_IO                *UdpIo,\r

   IN  NET_BUF               *Packet,\r
   IN  NET_BUF               *Packet,\r

-  IN  UDP_POINTS            *EndPoint, OPTIONAL\r
-  IN  IP4_ADDR              Gateway,\r
+  IN  UDP_END_POINT         *EndPoint OPTIONAL,\r
+  IN  EFI_IP_ADDRESS        *Gateway  OPTIONAL,\r

   IN  UDP_IO_CALLBACK       CallBack,\r
   IN  VOID                  *Context\r
   )\r
 {\r
   IN  UDP_IO_CALLBACK       CallBack,\r
   IN  VOID                  *Context\r
   )\r
 {\r

-  UDP_TX_TOKEN              *Token;\r
+  UDP_TX_TOKEN              *TxToken;\r

   EFI_STATUS                Status;\r
 \r
   EFI_STATUS                Status;\r
 \r

-  Token = UdpIoWrapTx (UdpIo, Packet, EndPoint, Gateway, CallBack, Context);\r
+  ASSERT ((UdpIo->UdpVersion == UDP_IO_UDP4_VERSION) ||\r
+          (UdpIo->UdpVersion == UDP_IO_UDP6_VERSION));\r

 \r
 \r

-  if (Token == NULL) {\r
+  TxToken = UdpIoCreateTxToken (UdpIo, Packet, EndPoint, Gateway, CallBack, Context);\r
+\r
+  if (TxToken == NULL) {\r

     return EFI_OUT_OF_RESOURCES;\r
   }\r
 \r
     return EFI_OUT_OF_RESOURCES;\r
   }\r
 \r


@@ -551,11 +943,17 @@ UdpIoSendDatagram (
   // Insert the tx token into SendDatagram list before transmitting it. Remove\r
   // it from the list if the returned status is not EFI_SUCCESS.\r
   //\r
   // Insert the tx token into SendDatagram list before transmitting it. Remove\r
   // it from the list if the returned status is not EFI_SUCCESS.\r
   //\r

-  InsertHeadList (&UdpIo->SentDatagram, &Token->Link);\r
-  Status = UdpIo->Udp->Transmit (UdpIo->Udp, &Token->UdpToken);\r
+  InsertHeadList (&UdpIo->SentDatagram, &TxToken->Link);\r
+\r
+  if (UdpIo->UdpVersion == UDP_IO_UDP4_VERSION) {\r
+    Status = UdpIo->Protocol.Udp4->Transmit (UdpIo->Protocol.Udp4, &TxToken->Token.Udp4);\r
+  } else {\r
+    Status = UdpIo->Protocol.Udp6->Transmit (UdpIo->Protocol.Udp6, &TxToken->Token.Udp6);\r
+  }\r
+\r

   if (EFI_ERROR (Status)) {\r
   if (EFI_ERROR (Status)) {\r

-    RemoveEntryList (&Token->Link);\r
-    UdpIoFreeTxToken (Token);\r
+    RemoveEntryList (&TxToken->Link);\r
+    UdpIoFreeTxToken (TxToken);\r

     return Status;\r
   }\r
 \r
     return Status;\r
   }\r
 \r


@@ -564,15 +962,16 @@ UdpIoSendDatagram (
 \r
 \r
 /**\r
 \r
 \r
 /**\r

-  The selection function to cancel a single sent datagram.\r
+  The select function to cancel a single sent datagram.\r

 \r
 \r

-  @param  Token                 The UDP TX token to test againist.\r
-  @param  Context               The context\r
-\r
-  @return TRUE if the packet is to be cancelled, otherwise FALSE.\r
+  @param[in]  Token                 The UDP_TX_TOKEN to test against\r
+  @param[in]  Context               The NET_BUF of the sent datagram\r

 \r
 \r

+  @retval TRUE              The packet is to be cancelled.\r
+  @retval FALSE             The packet is not to be cancelled.\r

 **/\r
 BOOLEAN\r
 **/\r
 BOOLEAN\r

+EFIAPI\r

 UdpIoCancelSingleDgram (\r
   IN UDP_TX_TOKEN           *Token,\r
   IN VOID                   *Context\r
 UdpIoCancelSingleDgram (\r
   IN UDP_TX_TOKEN           *Token,\r
   IN VOID                   *Context\r


@@ -589,210 +988,79 @@ UdpIoCancelSingleDgram (
   return FALSE;\r
 }\r
 \r
   return FALSE;\r
 }\r
 \r

-\r

 /**\r
   Cancel a single sent datagram.\r
 \r
 /**\r
   Cancel a single sent datagram.\r
 \r

-  @param  UdpIo                 The UDP IO port to cancel the packet from\r
-  @param  Packet                The packet to cancel\r
-\r
-  @return None\r
+  @param[in]  UdpIo                 The UDP_IO to cancel the packet from\r
+  @param[in]  Packet                The packet to cancel\r

 \r
 **/\r
 VOID\r
 EFIAPI\r
 UdpIoCancelSentDatagram (\r
 \r
 **/\r
 VOID\r
 EFIAPI\r
 UdpIoCancelSentDatagram (\r

-  IN  UDP_IO_PORT           *UdpIo,\r
+  IN  UDP_IO                *UdpIo,\r

   IN  NET_BUF               *Packet\r
   )\r
 {\r
   UdpIoCancelDgrams (UdpIo, EFI_ABORTED, UdpIoCancelSingleDgram, Packet);\r
 }\r
 \r
   IN  NET_BUF               *Packet\r
   )\r
 {\r
   UdpIoCancelDgrams (UdpIo, EFI_ABORTED, UdpIoCancelSingleDgram, Packet);\r
 }\r
 \r

-\r
-/**\r
-  Recycle the received UDP data.\r
-\r
-  @param  Context               The UDP_RX_TOKEN\r
-\r
-  @return None\r
-\r
-**/\r
-VOID\r
-UdpIoRecycleDgram (\r
-  IN VOID                   *Context\r
-  )\r
-{\r
-  UDP_RX_TOKEN              *Token;\r
-\r
-  Token = (UDP_RX_TOKEN *) Context;\r
-  gBS->SignalEvent (Token->UdpToken.Packet.RxData->RecycleSignal);\r
-  UdpIoFreeRxToken (Token);\r
-}\r
-\r

 /**\r
 /**\r

-  The event handle for UDP receive request. It will build\r
-  a NET_BUF from the recieved UDP data, then deliver it\r
-  to the receiver.\r
+  Issue a receive request to the UDP_IO.\r

 \r
 \r

-  @param  Context               The UDP RX token.\r
+  This function is called when upper-layer needs packet from UDP for processing.\r
+  Only one receive request is acceptable at a time so a common usage model is\r
+  to invoke this function inside its Callback function when the former packet\r
+  is processed.\r

 \r
 \r

-  @return None\r
-\r
-**/\r
-VOID\r
-EFIAPI\r
-UdpIoOnDgramRcvdDpc (\r
-  IN VOID                   *Context\r
-  )\r
-{\r
-  EFI_UDP4_COMPLETION_TOKEN *UdpToken;\r
-  EFI_UDP4_RECEIVE_DATA     *UdpRxData;\r
-  EFI_UDP4_SESSION_DATA     *UdpSession;\r
-  UDP_RX_TOKEN              *Token;\r
-  UDP_POINTS                Points;\r
-  NET_BUF                   *Netbuf;\r
-\r
-  Token   = (UDP_RX_TOKEN *) Context;\r
-\r
-  ASSERT ((Token->Signature == UDP_IO_RX_SIGNATURE) &&\r
-          (Token == Token->UdpIo->RecvRequest));\r
-\r
-  //\r
-  // Clear the receive request first in case that the caller\r
-  // wants to restart the receive in the callback.\r
-  //\r
-  Token->UdpIo->RecvRequest = NULL;\r
-\r
-  UdpToken  = &Token->UdpToken;\r
-  UdpRxData = UdpToken->Packet.RxData;\r
-\r
-  if (EFI_ERROR (UdpToken->Status) || (UdpRxData == NULL)) {\r
-    if (UdpToken->Status != EFI_ABORTED) {\r
-      //\r
-      // Invoke the CallBack only if the reception is not actively aborted.\r
-      //\r
-      Token->CallBack (NULL, NULL, UdpToken->Status, Token->Context);\r
-    }\r
-\r
-    UdpIoFreeRxToken (Token);\r
-    return;\r
-  }\r
-\r
-  //\r
-  // Build a NET_BUF from the UDP receive data, then deliver it up.\r
-  //\r
-  Netbuf = NetbufFromExt (\r
-             (NET_FRAGMENT *) UdpRxData->FragmentTable,\r
-             UdpRxData->FragmentCount,\r
-             0,\r
-             (UINT32) Token->HeadLen,\r
-             UdpIoRecycleDgram,\r
-             Token\r
-             );\r
-\r
-  if (Netbuf == NULL) {\r
-    gBS->SignalEvent (UdpRxData->RecycleSignal);\r
-    Token->CallBack (NULL, NULL, EFI_OUT_OF_RESOURCES, Token->Context);\r
-\r
-    UdpIoFreeRxToken (Token);\r
-    return;\r
-  }\r
-\r
-  UdpSession        = &UdpRxData->UdpSession;\r
-  Points.LocalPort  = UdpSession->DestinationPort;\r
-  Points.RemotePort = UdpSession->SourcePort;\r
-\r
-  CopyMem (&Points.LocalAddr, &UdpSession->DestinationAddress, sizeof (IP4_ADDR));\r
-  CopyMem (&Points.RemoteAddr, &UdpSession->SourceAddress, sizeof (IP4_ADDR));\r
-  Points.LocalAddr  = NTOHL (Points.LocalAddr);\r
-  Points.RemoteAddr = NTOHL (Points.RemoteAddr);\r
-\r
-  Token->CallBack (Netbuf, &Points, EFI_SUCCESS, Token->Context);\r
-}\r
-\r
-/**\r
-  Request UdpIoOnDgramRcvdDpc as a DPC at TPL_CALLBACK.\r
-\r
-  @param  Event                 The UDP receive request event.\r
-  @param  Context               The UDP RX token.\r
-\r
-  @return None\r
-\r
-**/\r
-VOID\r
-EFIAPI\r
-UdpIoOnDgramRcvd (\r
-  IN EFI_EVENT              Event,\r
-  IN VOID                   *Context\r
-  )\r
-/*++\r
-\r
-Routine Description:\r
-\r
-  Request UdpIoOnDgramRcvdDpc as a DPC at TPL_CALLBACK\r
-\r
-Arguments:\r
-\r
-  Event   - The UDP receive request event\r
-  Context - The UDP RX token.\r
-\r
-Returns:\r
-\r
-  None\r
-\r
---*/\r
-{\r
-  //\r
-  // Request UdpIoOnDgramRcvdDpc as a DPC at TPL_CALLBACK\r
-  //\r
-  NetLibQueueDpc (TPL_CALLBACK, UdpIoOnDgramRcvdDpc, Context);\r
-}\r
-\r
-\r
-/**\r
-  Issue a receive request to the UDP IO port.\r
-\r
-  @param  UdpIo                 The UDP IO port to recieve the packet from.\r
-  @param  CallBack              The call back function to execute when receive\r
-                                finished.\r
-  @param  Context               The opque context to the call back\r
-  @param  HeadLen               The lenght of the application's header\r
+  @param[in]  UdpIo                 The UDP_IO to receive the packet from.\r
+  @param[in]  CallBack              The call back function to execute when the packet\r
+                                    is received.\r
+  @param[in]  Context               The opaque context passed to Callback.\r
+  @param[in]  HeadLen               The length of the upper-layer's protocol header.\r

 \r
   @retval EFI_ALREADY_STARTED   There is already a pending receive request. Only\r
 \r
   @retval EFI_ALREADY_STARTED   There is already a pending receive request. Only\r

-                                one receive request is supported.\r
-  @retval EFI_OUT_OF_RESOURCES  Failed to allocate some resource.\r
+                                one receive request is supported at a time.\r
+  @retval EFI_OUT_OF_RESOURCES  Failed to allocate needed resources.\r

   @retval EFI_SUCCESS           The receive request is issued successfully.\r
   @retval EFI_SUCCESS           The receive request is issued successfully.\r

+  @retval EFI_UNSUPPORTED       The UDP version in UDP_IO is not supported.\r

 \r
 **/\r
 EFI_STATUS\r
 EFIAPI\r
 UdpIoRecvDatagram (\r
 \r
 **/\r
 EFI_STATUS\r
 EFIAPI\r
 UdpIoRecvDatagram (\r

-  IN  UDP_IO_PORT           *UdpIo,\r
+  IN  UDP_IO                *UdpIo,\r

   IN  UDP_IO_CALLBACK       CallBack,\r
   IN  VOID                  *Context,\r
   IN  UINT32                HeadLen\r
   )\r
 {\r
   IN  UDP_IO_CALLBACK       CallBack,\r
   IN  VOID                  *Context,\r
   IN  UINT32                HeadLen\r
   )\r
 {\r

-  UDP_RX_TOKEN              *Token;\r
+  UDP_RX_TOKEN              *RxToken;\r

   EFI_STATUS                Status;\r
 \r
   EFI_STATUS                Status;\r
 \r

+  ASSERT ((UdpIo->UdpVersion == UDP_IO_UDP4_VERSION) ||\r
+          (UdpIo->UdpVersion == UDP_IO_UDP6_VERSION));\r
+\r

   if (UdpIo->RecvRequest != NULL) {\r
     return EFI_ALREADY_STARTED;\r
   }\r
 \r
   if (UdpIo->RecvRequest != NULL) {\r
     return EFI_ALREADY_STARTED;\r
   }\r
 \r

-  Token = UdpIoCreateRxToken (UdpIo, CallBack, Context, HeadLen);\r
+  RxToken = UdpIoCreateRxToken (UdpIo, CallBack, Context, HeadLen);\r

 \r
 \r

-  if (Token == NULL) {\r
+  if (RxToken == NULL) {\r

     return EFI_OUT_OF_RESOURCES;\r
   }\r
 \r
     return EFI_OUT_OF_RESOURCES;\r
   }\r
 \r

-  UdpIo->RecvRequest = Token;\r
-  Status = UdpIo->Udp->Receive (UdpIo->Udp, &Token->UdpToken);\r
+  UdpIo->RecvRequest = RxToken;\r
+  if (UdpIo->UdpVersion == UDP_IO_UDP4_VERSION) {\r
+    Status = UdpIo->Protocol.Udp4->Receive (UdpIo->Protocol.Udp4, &RxToken->Token.Udp4);\r
+  } else {\r
+    Status = UdpIo->Protocol.Udp6->Receive (UdpIo->Protocol.Udp6, &RxToken->Token.Udp6);\r
+  }\r

 \r
   if (EFI_ERROR (Status)) {\r
     UdpIo->RecvRequest = NULL;\r
 \r
   if (EFI_ERROR (Status)) {\r
     UdpIo->RecvRequest = NULL;\r

-    UdpIoFreeRxToken (Token);\r
+    UdpIoFreeRxToken (RxToken);\r

   }\r
 \r
   return Status;\r
   }\r
 \r
   return Status;\r