]> git.proxmox.com Git - mirror_edk2.git/blobdiff - MdeModulePkg/Library/DxeUdpIoLib/DxeUdpIoLib.c
1. Update the UdpIo to a combined UdpIo to support both v4 and v6 stack.
[mirror_edk2.git] / MdeModulePkg / Library / DxeUdpIoLib / DxeUdpIoLib.c
index 91e211a9bdca379fa36d6719918c08da6fa5616a..bb85aebc282a13f76bb603ded45eea784fa5b083 100644 (file)
@@ -1,7 +1,7 @@
 /** @file\r
   Help functions to access UDP service, it is used by both the DHCP and MTFTP.\r
   \r
-Copyright (c) 2005 - 2007, Intel Corporation.<BR>\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
 which accompanies this distribution.  The full text of the license may be found at<BR>\r
@@ -14,6 +14,7 @@ WITHOUT WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.
 #include <Uefi.h>\r
 \r
 #include <Protocol/Udp4.h>\r
+#include <Protocol/Udp6.h>\r
 \r
 #include <Library/UdpIoLib.h>\r
 #include <Library/BaseLib.h>\r
@@ -27,31 +28,46 @@ WITHOUT WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.
 /**\r
   Free a UDP_TX_TOKEN. The TX event is closed.\r
 \r
-  @param[in]  Token                 The UDP_TX_TOKEN to release.\r
+  @param[in]  TxToken                 The UDP_TX_TOKEN to release.\r
 \r
 **/\r
 VOID\r
 UdpIoFreeTxToken (\r
-  IN UDP_TX_TOKEN           *Token\r
+  IN UDP_TX_TOKEN           *TxToken\r
   )\r
 {\r
-  gBS->CloseEvent (Token->UdpToken.Event);\r
-  gBS->FreePool (Token);\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]  Token                 The UDP_RX_TOKEN to release.\r
+  @param[in]  RxToken                 The UDP_RX_TOKEN to release.\r
 \r
 **/\r
 VOID\r
 UdpIoFreeRxToken (\r
-  IN UDP_RX_TOKEN           *Token\r
+  IN UDP_RX_TOKEN           *RxToken\r
   )\r
 {\r
-  gBS->CloseEvent (Token->UdpToken.Event);\r
-  gBS->FreePool (Token);\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
@@ -69,15 +85,22 @@ UdpIoOnDgramSentDpc (
   IN VOID                   *Context\r
   )\r
 {\r
-  UDP_TX_TOKEN              *Token;\r
+  UDP_TX_TOKEN              *TxToken;\r
 \r
-  Token   = (UDP_TX_TOKEN *) Context;\r
-  ASSERT (Token->Signature == UDP_IO_TX_SIGNATURE);\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
-  RemoveEntryList (&Token->Link);\r
-  Token->CallBack (Token->Packet, NULL, Token->UdpToken.Status, Token->Context);\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 (Token);\r
+  UdpIoFreeTxToken (TxToken);\r
 }\r
 \r
 /**\r
@@ -110,12 +133,20 @@ VOID
 UdpIoRecycleDgram (\r
   IN VOID                   *Context\r
   )\r
-{\r
-  UDP_RX_TOKEN              *Token;\r
+{ \r
+  UDP_RX_TOKEN              *RxToken;\r
+\r
+  RxToken = (UDP_RX_TOKEN *) Context;\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
-  Token = (UDP_RX_TOKEN *) Context;\r
-  gBS->SignalEvent (Token->UdpToken.Packet.RxData->RecycleSignal);\r
-  UdpIoFreeRxToken (Token);\r
+  UdpIoFreeRxToken (RxToken);\r
 }\r
 \r
 /**\r
@@ -133,69 +164,132 @@ UdpIoOnDgramRcvdDpc (
   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
+  EFI_STATUS                Status;\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
-  Token   = (UDP_RX_TOKEN *) Context;\r
+  RxToken = (UDP_RX_TOKEN *) Context;\r
 \r
-  ASSERT ((Token->Signature == UDP_IO_RX_SIGNATURE) &&\r
-          (Token == Token->UdpIo->RecvRequest));\r
+  ZeroMem (&EndPoint, sizeof(UDP_END_POINT));\r
+\r
+  ASSERT ((RxToken->Signature == UDP_IO_RX_SIGNATURE) && \r
+          (RxToken == RxToken->UdpIo->RecvRequest));\r
+\r
+  ASSERT ((RxToken->UdpIo->UdpVersion == UDP_IO_UDP4_VERSION) ||\r
+          (RxToken->UdpIo->UdpVersion == UDP_IO_UDP6_VERSION));\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
+  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
-  if (EFI_ERROR (UdpToken->Status) || (UdpRxData == NULL)) {\r
-    if (UdpToken->Status != EFI_ABORTED) {\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
-      Token->CallBack (NULL, NULL, UdpToken->Status, Token->Context);\r
+      RxToken->CallBack (NULL, NULL, Status, RxToken->Context);\r
     }\r
 \r
-    UdpIoFreeRxToken (Token);\r
+    UdpIoFreeRxToken (RxToken);\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
+  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
-  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
+    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
+    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
+    CopyMem (\r
+      &EndPoint.LocalAddr,\r
+      &((EFI_UDP6_SESSION_DATA *) Session)->DestinationAddress,\r
+      sizeof (EFI_IPv6_ADDRESS)\r
+      );\r
+\r
+    CopyMem (\r
+      &EndPoint.RemoteAddr,\r
+      &((EFI_UDP6_SESSION_DATA *) Session)->SourceAddress,\r
+      sizeof (EFI_IPv6_ADDRESS)\r
+      );\r
+\r
+    Ip6Swap128 (&EndPoint.LocalAddr.v6);\r
+    Ip6Swap128 (&EndPoint.RemoteAddr.v6);\r
+  }\r
 \r
-  Token->CallBack (Netbuf, &Points, EFI_SUCCESS, Token->Context);\r
+  RxToken->CallBack (Netbuf, &EndPoint, EFI_SUCCESS, RxToken->Context);\r
 }\r
 \r
 /**\r
@@ -231,7 +325,7 @@ UdpIoOnDgramRcvd (
 **/\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
@@ -240,6 +334,9 @@ UdpIoCreateRxToken (
   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
@@ -252,19 +349,35 @@ UdpIoCreateRxToken (
   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
-  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
-    gBS->FreePool (Token);\r
+    FreePool (Token);\r
     return NULL;\r
   }\r
 \r
@@ -272,9 +385,9 @@ UdpIoCreateRxToken (
 }\r
 \r
 /**\r
-  Wrap a transmit request into a UDP_TX_TOKEN.\r
+  Wrap a transmit request into a new created UDP_TX_TOKEN.\r
 \r
-  @param[in]  UdpIo                 The UdpIo port to send packet to.\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
@@ -286,188 +399,329 @@ UdpIoCreateRxToken (
 \r
 **/\r
 UDP_TX_TOKEN *\r
-UdpIoWrapTx (\r
-  IN UDP_IO_PORT            *UdpIo,\r
+UdpIoCreateTxToken (\r
+  IN UDP_IO                 *UdpIo,\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
-  UDP_TX_TOKEN              *Token;\r
-  EFI_UDP4_COMPLETION_TOKEN *UdpToken;\r
-  EFI_UDP4_TRANSMIT_DATA    *UdpTxData;\r
+  UDP_TX_TOKEN              *TxToken;\r
+  VOID                      *Token;\r
+  VOID                      *Data;\r
   EFI_STATUS                Status;\r
   UINT32                    Count;\r
-  IP4_ADDR                  Ip;\r
+  UINTN                     Size;\r
+  IP4_ADDR                  Ip;  \r
 \r
-  Token = AllocatePool (sizeof (UDP_TX_TOKEN) +\r
-                           sizeof (EFI_UDP4_FRAGMENT_DATA) * (Packet->BlockOpNum - 1));\r
+  ASSERT (Packet != NULL);\r
+  ASSERT ((UdpIo->UdpVersion == UDP_IO_UDP4_VERSION) ||\r
+          (UdpIo->UdpVersion == UDP_IO_UDP6_VERSION));\r
 \r
-  if (Token == NULL) {\r
-    return NULL;\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
-  Token->Signature  = UDP_IO_TX_SIGNATURE;\r
-  InitializeListHead (&Token->Link);\r
+  TxToken = AllocatePool (Size);\r
 \r
-  Token->UdpIo      = UdpIo;\r
-  Token->CallBack   = CallBack;\r
-  Token->Packet     = Packet;\r
-  Token->Context    = Context;\r
+  if (TxToken == NULL) {\r
+    return NULL;\r
+  }\r
 \r
-  UdpToken          = &(Token->UdpToken);\r
-  UdpToken->Status  = EFI_NOT_READY;\r
+  TxToken->Signature = UDP_IO_TX_SIGNATURE;\r
+  InitializeListHead (&TxToken->Link);\r
 \r
-  Status = gBS->CreateEvent (\r
-                  EVT_NOTIFY_SIGNAL,\r
-                  TPL_NOTIFY,\r
-                  UdpIoOnDgramSent,\r
-                  Token,\r
-                  &UdpToken->Event\r
-                  );\r
+  TxToken->UdpIo     = UdpIo;\r
+  TxToken->CallBack  = CallBack;\r
+  TxToken->Packet    = Packet;\r
+  TxToken->Context   = Context;\r
 \r
-  if (EFI_ERROR (Status)) {\r
-    gBS->FreePool (Token);\r
-    return NULL;\r
-  }\r
+  Token              = &(TxToken->Token);\r
+  Count              = Packet->BlockOpNum;\r
 \r
-  UdpTxData                 = &Token->UdpTxData;\r
-  UdpToken->Packet.TxData   = UdpTxData;\r
+  if (UdpIo->UdpVersion == UDP_IO_UDP4_VERSION) {\r
 \r
-  UdpTxData->UdpSessionData = NULL;\r
-  UdpTxData->GatewayAddress = NULL;\r
+    ((EFI_UDP4_COMPLETION_TOKEN *) Token)->Status = EFI_NOT_READY;\r
 \r
-  if (EndPoint != NULL) {\r
-    Ip = HTONL (EndPoint->LocalAddr);\r
-    CopyMem (&Token->UdpSession.SourceAddress, &Ip, sizeof (EFI_IPv4_ADDRESS));\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
-    Ip = HTONL (EndPoint->RemoteAddr);\r
-    CopyMem (&Token->UdpSession.DestinationAddress, &Ip, sizeof (EFI_IPv4_ADDRESS));\r
+    if (EFI_ERROR (Status)) {\r
+      FreePool (TxToken);\r
+      return NULL;\r
+    }\r
 \r
-    Token->UdpSession.SourcePort      = EndPoint->LocalPort;\r
-    Token->UdpSession.DestinationPort = EndPoint->RemotePort;\r
-    UdpTxData->UdpSessionData         = &Token->UdpSession;\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 != 0) {\r
-    Ip = HTONL (Gateway);\r
-    CopyMem (&Token->Gateway, &Ip, sizeof (EFI_IPv4_ADDRESS));\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
-    UdpTxData->GatewayAddress = &Token->Gateway;\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
-  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
+  return TxToken;\r
 }\r
 \r
 /**\r
-  Create a UDP_IO_PORT to access the UDP service. It will create and configure\r
+  Creates a UDP_IO to access the UDP service. It creates and configures\r
   a UDP child.\r
   \r
-  The function will locate the UDP service binding prototype on the Controller\r
-  parameter and use it to create a UDP child (aka Udp instance). Then the UDP\r
-  child will be configured by calling Configure function prototype. Any failures\r
-  in creating or configure the UDP child will lead to the failure of UDP_IO_PORT\r
-  creation.\r
+  This function:\r
+  # 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 a UDP instance)\r
+  # 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
   @param[in]  Controller            The controller that has the UDP service binding.\r
                                     protocol installed.\r
-  @param[in]  Image                 The image handle for the driver.\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
-  @return Newly-created UDP_IO_PORT or NULL if failed.\r
+  @return Newly-created UDP_IO or NULL if failed.\r
 \r
 **/\r
-UDP_IO_PORT *\r
+UDP_IO *\r
 EFIAPI\r
-UdpIoCreatePort (\r
+UdpIoCreateIo (\r
   IN  EFI_HANDLE            Controller,\r
-  IN  EFI_HANDLE            Image,\r
+  IN  EFI_HANDLE            ImageHandle,\r
   IN  UDP_IO_CONFIG         Configure,\r
+  IN  UINT8                 UdpVersion,\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
+  ASSERT ((UdpVersion == UDP_IO_UDP4_VERSION) || (UdpVersion == UDP_IO_UDP6_VERSION));\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
+  UdpIo->UdpVersion   = UdpVersion;\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
-  // Create a UDP child then open and configure it\r
-  //\r
-  Status = NetLibCreateServiceChild (\r
-             Controller,\r
-             Image,\r
-             &gEfiUdp4ServiceBindingProtocolGuid,\r
-             &UdpIo->UdpHandle\r
-             );\r
-\r
-  if (EFI_ERROR (Status)) {\r
-    goto FREE_MEM;\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
-\r
-  if (EFI_ERROR (Status)) {\r
-    goto FREE_CHILD;\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
-  if (EFI_ERROR (Configure (UdpIo, Context))) {\r
-    goto CLOSE_PROTOCOL;\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
-  Status = UdpIo->Udp->GetModeData (UdpIo->Udp, NULL, NULL, NULL, &UdpIo->SnpMode);\r
+    if (EFI_ERROR (Configure (UdpIo, Context))) {\r
+      goto CLOSE_PROTOCOL;\r
+    }\r
+  \r
+    Status = UdpIo->Protocol.Udp4->GetModeData (\r
+                                     UdpIo->Protocol.Udp4, \r
+                                     NULL, \r
+                                     NULL, \r
+                                     NULL, \r
+                                     &UdpIo->SnpMode\r
+                                     );\r
+  \r
+    if (EFI_ERROR (Status)) {\r
+      goto CLOSE_PROTOCOL;\r
+    }\r
 \r
-  if (EFI_ERROR (Status)) {\r
-    goto CLOSE_PROTOCOL;\r
+  } else {\r
+    \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
-  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
-  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
-  gBS->FreePool (UdpIo);\r
+  FreePool (UdpIo);\r
   return NULL;\r
 }\r
 \r
@@ -475,7 +729,7 @@ FREE_MEM:
   Cancel all the sent datagram that pass the selection criteria of ToCancel.\r
   If ToCancel is NULL, all the datagrams are cancelled.\r
 \r
-  @param[in]  UdpIo                 The UDP_IO_PORT to cancel packet.\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
@@ -485,7 +739,7 @@ FREE_MEM:
 VOID\r
 EFIAPI\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
@@ -493,34 +747,45 @@ UdpIoCancelDgrams (
 {\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
-    Token = NET_LIST_USER_STRUCT (Entry, UDP_TX_TOKEN, Link);\r
+    TxToken = NET_LIST_USER_STRUCT (Entry, UDP_TX_TOKEN, Link);\r
+\r
+    if ((ToCancel == NULL) || (ToCancel (TxToken, Context))) {\r
 \r
-    if ((ToCancel == NULL) || (ToCancel (Token, Context))) {\r
-      UdpIo->Udp->Cancel (UdpIo->Udp, &Token->UdpToken);\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
-  Free the UDP_IO_PORT and all its related resources.\r
+  Free the UDP_IO and all its related resources.\r
   \r
   The function will cancel all sent datagram and receive request.\r
 \r
-  @param[in]  UdpIo                 The UDP_IO_PORT to free.\r
+  @param[in]  UdpIo             The UDP_IO to free.\r
 \r
-  @retval EFI_SUCCESS           The UDP_IO_PORT is freed.\r
+  @retval EFI_SUCCESS           The UDP_IO is freed.\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
-  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
@@ -532,78 +797,116 @@ UdpIoFreePort (
   //\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
-  // 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
-  gBS->FreePool (UdpIo);\r
+  FreePool (UdpIo);\r
   return EFI_SUCCESS;\r
 }\r
 \r
 \r
 /**\r
-  Clean up the UDP_IO_PORT without freeing it. The function is called when\r
-  user wants to re-use the UDP_IO_PORT later.\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
   It will release all the transmitted datagrams and receive request. It will\r
   also configure NULL for the UDP instance.\r
 \r
-  @param[in]  UdpIo                 The UDP_IO_PORT to clean up.\r
+  @param[in]  UdpIo                 The UDP_IO to clean up.\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
+  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
-  if ((RxToken = UdpIo->RecvRequest) != NULL) {\r
-    UdpIo->Udp->Cancel (UdpIo->Udp, &RxToken->UdpToken);\r
-  }\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
+    UdpIo->Protocol.Udp4->Configure (UdpIo->Protocol.Udp4, NULL);\r
 \r
-  UdpIo->Udp->Configure (UdpIo->Udp, NULL);\r
+  } else {\r
+    if ((RxToken = UdpIo->RecvRequest) != NULL) {\r
+      UdpIo->Protocol.Udp6->Cancel (UdpIo->Protocol.Udp6, &RxToken->Token.Udp6);\r
+    }\r
+\r
+    UdpIo->Protocol.Udp6->Configure (UdpIo->Protocol.Udp6, NULL);\r
+  }\r
 }\r
 \r
 /**\r
-  Send a packet through the UDP_IO_PORT.\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_PORT to send the packet through.\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]  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
@@ -616,20 +919,23 @@ UdpIoCleanPort (
 EFI_STATUS\r
 EFIAPI\r
 UdpIoSendDatagram (\r
-  IN  UDP_IO_PORT           *UdpIo,\r
+  IN  UDP_IO                *UdpIo,\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
-  UDP_TX_TOKEN              *Token;\r
+  UDP_TX_TOKEN              *TxToken;\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
-  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
@@ -637,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
-  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
-    RemoveEntryList (&Token->Link);\r
-    UdpIoFreeTxToken (Token);\r
+    RemoveEntryList (&TxToken->Link);\r
+    UdpIoFreeTxToken (TxToken);\r
     return Status;\r
   }\r
 \r
@@ -678,14 +990,14 @@ UdpIoCancelSingleDgram (
 /**\r
   Cancel a single sent datagram.\r
 \r
-  @param[in]  UdpIo                 The UDP_IO_PORT to cancel the packet from\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
-  IN  UDP_IO_PORT           *UdpIo,\r
+  IN  UDP_IO                *UdpIo,\r
   IN  NET_BUF               *Packet\r
   )\r
 {\r
@@ -693,53 +1005,61 @@ UdpIoCancelSentDatagram (
 }\r
 \r
 /**\r
-  Issue a receive request to the UDP_IO_PORT.\r
+  Issue a receive request to the UDP_IO.\r
   \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
-  @param[in]  UdpIo                 The UDP_IO_PORT to receive the packet from.\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
+  @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
                                 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_UNSUPPORTED       The UDP version in UDP_IO is not supported.\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
-  UDP_RX_TOKEN              *Token;\r
+  UDP_RX_TOKEN              *RxToken;\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
-  Token = UdpIoCreateRxToken (UdpIo, CallBack, Context, HeadLen);\r
+  RxToken = UdpIoCreateRxToken (UdpIo, CallBack, Context, HeadLen);\r
 \r
-  if (Token == NULL) {\r
+  if (RxToken == NULL) {\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
-    UdpIoFreeRxToken (Token);\r
+    UdpIoFreeRxToken (RxToken);\r
   }\r
 \r
   return Status;\r