X-Git-Url: https://git.proxmox.com/?p=mirror_edk2.git;a=blobdiff_plain;f=MdeModulePkg%2FLibrary%2FDxeUdpIoLib%2FDxeUdpIoLib.c;h=a27f03aab91414c891918e777eb27bb0bf6f83aa;hp=84bc295bb14bc139a60853f2abfb8a92d62a63cf;hb=e798cd87ca9a3a30c4cea50c5f5de84e10a8bc5a;hpb=cbf316f20726bb31b7c37424601643790dbd02d9

diff --git a/MdeModulePkg/Library/DxeUdpIoLib/DxeUdpIoLib.c b/MdeModulePkg/Library/DxeUdpIoLib/DxeUdpIoLib.c
index 84bc295bb1..a27f03aab9 100644
--- a/MdeModulePkg/Library/DxeUdpIoLib/DxeUdpIoLib.c
+++ b/MdeModulePkg/Library/DxeUdpIoLib/DxeUdpIoLib.c
@@ -1,174 +1,332 @@
 /** @file
+  Help functions to access UDP service, it is used by both the DHCP and MTFTP.
 
-Copyright (c) 2006 - 2007, Intel Corporation
+Copyright (c) 2005 - 2009, Intel Corporation.<BR>
 All rights reserved. This program and the accompanying materials
 are licensed and made available under the terms and conditions of the BSD License
-which accompanies this distribution.  The full text of the license may be found at
+which accompanies this distribution.  The full text of the license may be found at<BR>
 http://opensource.org/licenses/bsd-license.php
 
 THE PROGRAM IS DISTRIBUTED UNDER THE BSD LICENSE ON AN "AS IS" BASIS,
 WITHOUT WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.
+**/
 
+#include <Uefi.h>
 
-Module Name:
+#include <Protocol/Udp4.h>
+#include <Protocol/Udp6.h>
 
-  Udp4Io.c
+#include <Library/UdpIoLib.h>
+#include <Library/BaseLib.h>
+#include <Library/DebugLib.h>
+#include <Library/UefiBootServicesTableLib.h>
+#include <Library/MemoryAllocationLib.h>
+#include <Library/BaseMemoryLib.h>
+#include <Library/DpcLib.h>
 
-Abstract:
 
-  Help functions to access UDP service, it is used by both the DHCP and MTFTP.
+/**
+  Free a UDP_TX_TOKEN. The TX event is closed.
 
+  @param[in]  TxToken                 The UDP_TX_TOKEN to release.
 
 **/
+VOID
+UdpIoFreeTxToken (
+  IN UDP_TX_TOKEN           *TxToken
+  )
+{
 
-#include <PiDxe.h>
+  if (TxToken->UdpIo->UdpVersion == UDP_IO_UDP4_VERSION) {
+    gBS->CloseEvent (TxToken->Token.Udp4.Event);
+  } else if (TxToken->UdpIo->UdpVersion == UDP_IO_UDP6_VERSION) {
+    gBS->CloseEvent (TxToken->Token.Udp6.Event);
+  } else {
+    ASSERT (FALSE);
+  }
 
-#include <Protocol/Udp4.h>
+  FreePool (TxToken);
+}
 
-#include <Library/UdpIoLib.h>
-#include <Library/BaseLib.h>
-#include <Library/DebugLib.h>
-#include <Library/UefiBootServicesTableLib.h>
-#include <Library/MemoryAllocationLib.h>
+/**
+  Free a UDP_RX_TOKEN. The RX event is closed.
 
-STATIC
+  @param[in]  RxToken                 The UDP_RX_TOKEN to release.
+
+**/
+VOID
+UdpIoFreeRxToken (
+  IN UDP_RX_TOKEN           *RxToken
+  )
+{
+  if (RxToken->UdpIo->UdpVersion == UDP_IO_UDP4_VERSION) {
+    gBS->CloseEvent (RxToken->Token.Udp4.Event);
+  } else if (RxToken->UdpIo->UdpVersion == UDP_IO_UDP6_VERSION) {
+    gBS->CloseEvent (RxToken->Token.Udp6.Event);
+  } else {
+    ASSERT (FALSE);
+  }
+
+  FreePool (RxToken);
+}
+
+/**
+  The callback function when the packet is sent by UDP.
+
+  It will remove the packet from the local list then call
+  the packet owner's callback function set by UdpIoSendDatagram.
+
+  @param[in]  Context               The UDP TX Token.
+
+**/
+VOID
+EFIAPI
+UdpIoOnDgramSentDpc (
+  IN VOID                   *Context
+  )
+{
+  UDP_TX_TOKEN              *TxToken;
+
+  TxToken = (UDP_TX_TOKEN *) Context;
+  ASSERT (TxToken->Signature == UDP_IO_TX_SIGNATURE);
+  ASSERT ((TxToken->UdpIo->UdpVersion == UDP_IO_UDP4_VERSION) ||
+          (TxToken->UdpIo->UdpVersion == UDP_IO_UDP6_VERSION));
+
+  RemoveEntryList (&TxToken->Link);
+
+  if (TxToken->UdpIo->UdpVersion == UDP_IO_UDP4_VERSION) {
+    TxToken->CallBack (TxToken->Packet, NULL, TxToken->Token.Udp4.Status, TxToken->Context);
+  } else {
+    TxToken->CallBack (TxToken->Packet, NULL, TxToken->Token.Udp6.Status, TxToken->Context);
+  }
+
+  UdpIoFreeTxToken (TxToken);
+}
+
+/**
+  Request UdpIoOnDgramSentDpc as a DPC at TPL_CALLBACK.
+
+  @param[in]  Event                 The event signaled.
+  @param[in]  Context               The UDP TX Token.
+
+**/
 VOID
 EFIAPI
 UdpIoOnDgramSent (
   IN EFI_EVENT              Event,
   IN VOID                   *Context
-  );
+  )
+{
+  //
+  // Request UdpIoOnDgramSentDpc as a DPC at TPL_CALLBACK
+  //
+  QueueDpc (TPL_CALLBACK, UdpIoOnDgramSentDpc, Context);
+}
+
+/**
+  Recycle the received UDP data.
+
+  @param[in]  Context               The UDP_RX_TOKEN.
 
-STATIC
+**/
 VOID
 EFIAPI
-UdpIoOnDgramRcvd (
-  IN EFI_EVENT              Event,
+UdpIoRecycleDgram (
   IN VOID                   *Context
-  );
+  )
+{
+  UDP_RX_TOKEN              *RxToken;
 
+  RxToken = (UDP_RX_TOKEN *) Context;
+
+  if (RxToken->UdpIo->UdpVersion == UDP_IO_UDP4_VERSION) {
+    gBS->SignalEvent (RxToken->Token.Udp4.Packet.RxData->RecycleSignal);
+  } else if (RxToken->UdpIo->UdpVersion == UDP_IO_UDP6_VERSION) {
+    gBS->SignalEvent (RxToken->Token.Udp6.Packet.RxData->RecycleSignal);
+  } else {
+    ASSERT (FALSE);
+  }
+
+  UdpIoFreeRxToken (RxToken);
+}
 
 /**
-  Wrap a transmit request into a UDP_TX_TOKEN.
+  The event handle for UDP receive request.
 
-  @param  UdpIo                 The UdpIo port to send packet to
-  @param  Packet                The user's packet
-  @param  EndPoint              The local and remote access point
-  @param  Gateway               The overrided next hop
-  @param  CallBack              The function to call when transmission completed.
-  @param  Context               The opaque parameter to the call back
+  It will build a NET_BUF from the recieved UDP data, then deliver it
+  to the receiver.
 
-  @return The wrapped transmission request or NULL if failed to allocate resources.
+  @param[in]  Context               The UDP RX token.
 
 **/
-STATIC
-UDP_TX_TOKEN *
-UdpIoWrapTx (
-  IN UDP_IO_PORT            *UdpIo,
-  IN NET_BUF                *Packet,
-  IN UDP_POINTS             *EndPoint, OPTIONAL
-  IN IP4_ADDR               Gateway,
-  IN UDP_IO_CALLBACK        CallBack,
+VOID
+EFIAPI
+UdpIoOnDgramRcvdDpc (
   IN VOID                   *Context
   )
 {
-  UDP_TX_TOKEN              *Token;
-  EFI_UDP4_COMPLETION_TOKEN *UdpToken;
-  EFI_UDP4_TRANSMIT_DATA    *UdpTxData;
   EFI_STATUS                Status;
-  UINT32                    Count;
+  VOID                      *Token;
+  VOID                      *RxData;
+  VOID                      *Session;
+  UDP_RX_TOKEN              *RxToken;
+  UDP_END_POINT             EndPoint;
+  NET_BUF                   *Netbuf;
 
-  Token = NetAllocatePool (sizeof (UDP_TX_TOKEN) +
-                           sizeof (EFI_UDP4_FRAGMENT_DATA) * (Packet->BlockOpNum - 1));
+  RxToken = (UDP_RX_TOKEN *) Context;
 
-  if (Token == NULL) {
-    return NULL;
-  }
+  ZeroMem (&EndPoint, sizeof(UDP_END_POINT));
 
-  Token->Signature  = UDP_IO_TX_SIGNATURE;
-  NetListInit (&Token->Link);
+  ASSERT ((RxToken->Signature == UDP_IO_RX_SIGNATURE) &&
+          (RxToken == RxToken->UdpIo->RecvRequest));
 
-  Token->UdpIo      = UdpIo;
-  Token->CallBack   = CallBack;
-  Token->Packet     = Packet;
-  Token->Context    = Context;
+  ASSERT ((RxToken->UdpIo->UdpVersion == UDP_IO_UDP4_VERSION) ||
+          (RxToken->UdpIo->UdpVersion == UDP_IO_UDP6_VERSION));
 
-  UdpToken          = &(Token->UdpToken);
-  UdpToken->Status  = EFI_NOT_READY;
+  //
+  // Clear the receive request first in case that the caller
+  // wants to restart the receive in the callback.
+  //
+  RxToken->UdpIo->RecvRequest = NULL;
+
+  if (RxToken->UdpIo->UdpVersion == UDP_IO_UDP4_VERSION) {
+    Token  = &RxToken->Token.Udp4;
+    RxData = ((EFI_UDP4_COMPLETION_TOKEN *) Token)->Packet.RxData;
+    Status = ((EFI_UDP4_COMPLETION_TOKEN *) Token)->Status;
+  } else {
+    Token  = &RxToken->Token.Udp6;
+    RxData = ((EFI_UDP6_COMPLETION_TOKEN *) Token)->Packet.RxData;
+    Status = ((EFI_UDP6_COMPLETION_TOKEN *) Token)->Status;
+  }
 
-  Status = gBS->CreateEvent (
-                  EVT_NOTIFY_SIGNAL,
-                  TPL_CALLBACK,
-                  UdpIoOnDgramSent,
-                  Token,
-                  &UdpToken->Event
-                  );
+  if (EFI_ERROR (Status) || RxData == NULL) {
+    if (Status != EFI_ABORTED) {
+      //
+      // Invoke the CallBack only if the reception is not actively aborted.
+      //
+      RxToken->CallBack (NULL, NULL, Status, RxToken->Context);
+    }
 
-  if (EFI_ERROR (Status)) {
-    NetFreePool (Token);
-    return NULL;
+    UdpIoFreeRxToken (RxToken);
+    return;
   }
 
-  UdpTxData                 = &Token->UdpTxData;
-  UdpToken->Packet.TxData   = UdpTxData;
+  //
+  // Build a NET_BUF from the UDP receive data, then deliver it up.
+  //
+  if (RxToken->UdpIo->UdpVersion == UDP_IO_UDP4_VERSION) {
+
+    Netbuf = NetbufFromExt (
+               (NET_FRAGMENT *)((EFI_UDP4_RECEIVE_DATA *) RxData)->FragmentTable,
+               ((EFI_UDP4_RECEIVE_DATA *) RxData)->FragmentCount,
+               0,
+               (UINT32) RxToken->HeadLen,
+               UdpIoRecycleDgram,
+               RxToken
+               );
+
+    if (Netbuf == NULL) {
+      gBS->SignalEvent (((EFI_UDP4_RECEIVE_DATA *) RxData)->RecycleSignal);
+      RxToken->CallBack (NULL, NULL, EFI_OUT_OF_RESOURCES, RxToken->Context);
+
+      UdpIoFreeRxToken (RxToken);
+      return;
+    }
 
-  UdpTxData->UdpSessionData = NULL;
-  UdpTxData->GatewayAddress = NULL;
+    Session             = &((EFI_UDP4_RECEIVE_DATA *) RxData)->UdpSession;
+    EndPoint.LocalPort  = ((EFI_UDP4_SESSION_DATA *) Session)->DestinationPort;
+    EndPoint.RemotePort = ((EFI_UDP4_SESSION_DATA *) Session)->SourcePort;
+
+    CopyMem (
+      &EndPoint.LocalAddr,
+      &((EFI_UDP4_SESSION_DATA *) Session)->DestinationAddress,
+      sizeof (EFI_IPv4_ADDRESS)
+      );
+
+    CopyMem (
+      &EndPoint.RemoteAddr,
+      &((EFI_UDP4_SESSION_DATA *) Session)->SourceAddress,
+      sizeof (EFI_IPv4_ADDRESS)
+      );
+
+    EndPoint.LocalAddr.Addr[0]  = NTOHL (EndPoint.LocalAddr.Addr[0]);
+    EndPoint.RemoteAddr.Addr[0] = NTOHL (EndPoint.RemoteAddr.Addr[0]);
+  } else {
+
+    Netbuf = NetbufFromExt (
+               (NET_FRAGMENT *)((EFI_UDP6_RECEIVE_DATA *) RxData)->FragmentTable,
+               ((EFI_UDP6_RECEIVE_DATA *) RxData)->FragmentCount,
+               0,
+               (UINT32) RxToken->HeadLen,
+               UdpIoRecycleDgram,
+               RxToken
+               );
+
+    if (Netbuf == NULL) {
+      gBS->SignalEvent (((EFI_UDP6_RECEIVE_DATA *) RxData)->RecycleSignal);
+      RxToken->CallBack (NULL, NULL, EFI_OUT_OF_RESOURCES, RxToken->Context);
+
+      UdpIoFreeRxToken (RxToken);
+      return;
+    }
 
-  if (EndPoint != NULL) {
-    EFI_IP4 (Token->UdpSession.SourceAddress)      = HTONL (EndPoint->LocalAddr);
-    EFI_IP4 (Token->UdpSession.DestinationAddress) = HTONL (EndPoint->RemoteAddr);
-    Token->UdpSession.SourcePort                   = EndPoint->LocalPort;
-    Token->UdpSession.DestinationPort              = EndPoint->RemotePort;
-    UdpTxData->UdpSessionData                      = &Token->UdpSession;
-  }
+    Session             = &((EFI_UDP6_RECEIVE_DATA *) RxData)->UdpSession;
+    EndPoint.LocalPort  = ((EFI_UDP6_SESSION_DATA *) Session)->DestinationPort;
+    EndPoint.RemotePort = ((EFI_UDP6_SESSION_DATA *) Session)->SourcePort;
 
-  if (Gateway != 0) {
-    EFI_IP4 (Token->Gateway)  = HTONL (Gateway);
-    UdpTxData->GatewayAddress = &Token->Gateway;
-  }
+    CopyMem (
+      &EndPoint.LocalAddr,
+      &((EFI_UDP6_SESSION_DATA *) Session)->DestinationAddress,
+      sizeof (EFI_IPv6_ADDRESS)
+      );
 
-  UdpTxData->DataLength = Packet->TotalSize;
-  Count                 = Packet->BlockOpNum;
-  NetbufBuildExt (Packet, (NET_FRAGMENT *) UdpTxData->FragmentTable, &Count);
-  UdpTxData->FragmentCount = Count;
+    CopyMem (
+      &EndPoint.RemoteAddr,
+      &((EFI_UDP6_SESSION_DATA *) Session)->SourceAddress,
+      sizeof (EFI_IPv6_ADDRESS)
+      );
 
-  return Token;
-}
+    Ip6Swap128 (&EndPoint.LocalAddr.v6);
+    Ip6Swap128 (&EndPoint.RemoteAddr.v6);
+  }
 
+  RxToken->CallBack (Netbuf, &EndPoint, EFI_SUCCESS, RxToken->Context);
+}
 
 /**
-  Free a UDP_TX_TOKEN. The event is closed and memory released.
-
-  @param  Token                 The UDP_TX_TOKEN to release.
+  Request UdpIoOnDgramRcvdDpc() as a DPC at TPL_CALLBACK.
 
-  @return None
+  @param[in]  Event                 The UDP receive request event.
+  @param[in]  Context               The UDP RX token.
 
 **/
 VOID
-UdpIoFreeTxToken (
-  IN UDP_TX_TOKEN           *Token
+EFIAPI
+UdpIoOnDgramRcvd (
+  IN EFI_EVENT              Event,
+  IN VOID                   *Context
   )
 {
-  gBS->CloseEvent (Token->UdpToken.Event);
-  NetFreePool (Token);
+  //
+  // Request UdpIoOnDgramRcvdDpc as a DPC at TPL_CALLBACK
+  //
+  QueueDpc (TPL_CALLBACK, UdpIoOnDgramRcvdDpc, Context);
 }
 
-
 /**
   Create a UDP_RX_TOKEN to wrap the request.
 
-  @param  UdpIo                 The UdpIo to receive packets from
-  @param  CallBack              The function to call when receive finished.
-  @param  Context               The opaque parameter to the CallBack
-  @param  HeadLen               The head length to reserver for the packet.
+  @param[in]  UdpIo                 The UdpIo to receive packets from.
+  @param[in]  CallBack              The function to call when receive finished.
+  @param[in]  Context               The opaque parameter to the CallBack.
+  @param[in]  HeadLen               The head length to reserver for the packet.
 
-  @return The Wrapped request or NULL if failed to allocate resources.
+  @return The Wrapped request or NULL if failed to allocate resources or some errors happened.
 
 **/
 UDP_RX_TOKEN *
 UdpIoCreateRxToken (
-  IN UDP_IO_PORT            *UdpIo,
+  IN UDP_IO                 *UdpIo,
   IN UDP_IO_CALLBACK        CallBack,
   IN VOID                   *Context,
   IN UINT32                 HeadLen
@@ -177,7 +335,10 @@ UdpIoCreateRxToken (
   UDP_RX_TOKEN              *Token;
   EFI_STATUS                Status;
 
-  Token = NetAllocatePool (sizeof (UDP_RX_TOKEN));
+  ASSERT ((UdpIo->UdpVersion == UDP_IO_UDP4_VERSION) ||
+          (UdpIo->UdpVersion == UDP_IO_UDP6_VERSION));
+
+  Token = AllocatePool (sizeof (UDP_RX_TOKEN));
 
   if (Token == NULL) {
     return NULL;
@@ -189,197 +350,442 @@ UdpIoCreateRxToken (
   Token->Context                = Context;
   Token->HeadLen                = HeadLen;
 
-  Token->UdpToken.Status        = EFI_NOT_READY;
-  Token->UdpToken.Packet.RxData = NULL;
+  if (UdpIo->UdpVersion == UDP_IO_UDP4_VERSION) {
+
+    Token->Token.Udp4.Status        = EFI_NOT_READY;
+    Token->Token.Udp4.Packet.RxData = NULL;
+
+    Status = gBS->CreateEvent (
+                    EVT_NOTIFY_SIGNAL,
+                    TPL_NOTIFY,
+                    UdpIoOnDgramRcvd,
+                    Token,
+                    &Token->Token.Udp4.Event
+                    );
+    } else {
+
+    Token->Token.Udp6.Status        = EFI_NOT_READY;
+    Token->Token.Udp6.Packet.RxData = NULL;
+
+    Status = gBS->CreateEvent (
+                    EVT_NOTIFY_SIGNAL,
+                    TPL_NOTIFY,
+                    UdpIoOnDgramRcvd,
+                    Token,
+                    &Token->Token.Udp6.Event
+                    );
+  }
 
-  Status = gBS->CreateEvent (
-                  EVT_NOTIFY_SIGNAL,
-                  TPL_CALLBACK,
-                  UdpIoOnDgramRcvd,
-                  Token,
-                  &Token->UdpToken.Event
-                  );
 
   if (EFI_ERROR (Status)) {
-    NetFreePool (Token);
+    FreePool (Token);
     return NULL;
   }
 
   return Token;
 }
 
-
 /**
-  Free a receive request wrap.
+  Wrap a transmit request into a new created UDP_TX_TOKEN.
 
-  @param  Token                 The receive request to release.
+  @param[in]  UdpIo                 The UdpIo to send packet to.
+  @param[in]  Packet                The user's packet.
+  @param[in]  EndPoint              The local and remote access point.
+  @param[in]  Gateway               The overrided next hop.
+  @param[in]  CallBack              The function to call when transmission completed.
+  @param[in]  Context               The opaque parameter to the call back.
 
-  @return None
+  @return The wrapped transmission request or NULL if failed to allocate resources
+          or for some errors.
 
 **/
-VOID
-UdpIoFreeRxToken (
-  IN UDP_RX_TOKEN           *Token
+UDP_TX_TOKEN *
+UdpIoCreateTxToken (
+  IN UDP_IO                 *UdpIo,
+  IN NET_BUF                *Packet,
+  IN UDP_END_POINT          *EndPoint OPTIONAL,
+  IN EFI_IP_ADDRESS         *Gateway  OPTIONAL,
+  IN UDP_IO_CALLBACK        CallBack,
+  IN VOID                   *Context
   )
 {
-  gBS->CloseEvent (Token->UdpToken.Event);
-  NetFreePool (Token);
-}
+  UDP_TX_TOKEN              *TxToken;
+  VOID                      *Token;
+  VOID                      *Data;
+  EFI_STATUS                Status;
+  UINT32                    Count;
+  UINTN                     Size;
+  IP4_ADDR                  Ip;
+
+  ASSERT (Packet != NULL);
+  ASSERT ((UdpIo->UdpVersion == UDP_IO_UDP4_VERSION) ||
+          (UdpIo->UdpVersion == UDP_IO_UDP6_VERSION));
+
+  if (UdpIo->UdpVersion == UDP_IO_UDP4_VERSION) {
+    Size = sizeof (UDP_TX_TOKEN) + sizeof (EFI_UDP4_FRAGMENT_DATA) * (Packet->BlockOpNum - 1);
+  } else {
+    Size = sizeof (UDP_TX_TOKEN) + sizeof (EFI_UDP6_FRAGMENT_DATA) * (Packet->BlockOpNum - 1);
+  }
+
+  TxToken = AllocatePool (Size);
+
+  if (TxToken == NULL) {
+    return NULL;
+  }
+
+  TxToken->Signature = UDP_IO_TX_SIGNATURE;
+  InitializeListHead (&TxToken->Link);
+
+  TxToken->UdpIo     = UdpIo;
+  TxToken->CallBack  = CallBack;
+  TxToken->Packet    = Packet;
+  TxToken->Context   = Context;
+
+  Token              = &(TxToken->Token);
+  Count              = Packet->BlockOpNum;
+
+  if (UdpIo->UdpVersion == UDP_IO_UDP4_VERSION) {
+
+    ((EFI_UDP4_COMPLETION_TOKEN *) Token)->Status = EFI_NOT_READY;
+
+    Status = gBS->CreateEvent (
+                    EVT_NOTIFY_SIGNAL,
+                    TPL_NOTIFY,
+                    UdpIoOnDgramSent,
+                    TxToken,
+                    &((EFI_UDP4_COMPLETION_TOKEN *) Token)->Event
+                    );
+
+    if (EFI_ERROR (Status)) {
+      FreePool (TxToken);
+      return NULL;
+    }
+
+    Data = &(TxToken->Data.Udp4);
+    ((EFI_UDP4_COMPLETION_TOKEN *) Token)->Packet.TxData = Data;
+
+    ((EFI_UDP4_TRANSMIT_DATA *) Data)->UdpSessionData    = NULL;
+    ((EFI_UDP4_TRANSMIT_DATA *) Data)->GatewayAddress    = NULL;
+    ((EFI_UDP4_TRANSMIT_DATA *) Data)->DataLength        = Packet->TotalSize;
+
+    NetbufBuildExt (
+      Packet,
+      (NET_FRAGMENT *)((EFI_UDP4_TRANSMIT_DATA *) Data)->FragmentTable,
+      &Count
+      );
+
+    ((EFI_UDP4_TRANSMIT_DATA *) Data)->FragmentCount     = Count;
+
+    if (EndPoint != NULL) {
+      Ip = HTONL (EndPoint->LocalAddr.Addr[0]);
+      CopyMem (
+        &TxToken->Session.Udp4.SourceAddress,
+        &Ip,
+        sizeof (EFI_IPv4_ADDRESS)
+        );
+
+      Ip = HTONL (EndPoint->RemoteAddr.Addr[0]);
+      CopyMem (
+        &TxToken->Session.Udp4.DestinationAddress,
+        &Ip,
+        sizeof (EFI_IPv4_ADDRESS)
+        );
+
+      TxToken->Session.Udp4.SourcePort                   = EndPoint->LocalPort;
+      TxToken->Session.Udp4.DestinationPort              = EndPoint->RemotePort;
+      ((EFI_UDP4_TRANSMIT_DATA *) Data)->UdpSessionData  = &(TxToken->Session.Udp4);
+    }
+
+    if (Gateway != NULL && (Gateway->Addr[0] != 0)) {
+      Ip = HTONL (Gateway->Addr[0]);
+      CopyMem (&TxToken->Gateway, &Ip, sizeof (EFI_IPv4_ADDRESS));
+      ((EFI_UDP4_TRANSMIT_DATA *) Data)->GatewayAddress = &TxToken->Gateway;
+    }
+
+  } else {
+
+    ((EFI_UDP6_COMPLETION_TOKEN *) Token)->Status = EFI_NOT_READY;
+
+    Status = gBS->CreateEvent (
+                    EVT_NOTIFY_SIGNAL,
+                    TPL_NOTIFY,
+                    UdpIoOnDgramSent,
+                    TxToken,
+                    &((EFI_UDP6_COMPLETION_TOKEN *) Token)->Event
+                    );
+
+    if (EFI_ERROR (Status)) {
+      FreePool (TxToken);
+      return NULL;
+    }
 
+    Data = &(TxToken->Data.Udp6);
+    ((EFI_UDP6_COMPLETION_TOKEN *) Token)->Packet.TxData  = Data;
+    ((EFI_UDP6_TRANSMIT_DATA *) Data)->UdpSessionData     = NULL;
+    ((EFI_UDP6_TRANSMIT_DATA *) Data)->DataLength         = Packet->TotalSize;
+
+    NetbufBuildExt (
+      Packet,
+      (NET_FRAGMENT *)((EFI_UDP6_TRANSMIT_DATA *) Data)->FragmentTable,
+      &Count
+      );
+
+    ((EFI_UDP6_TRANSMIT_DATA *) Data)->FragmentCount      = Count;
+
+    if (EndPoint != NULL) {
+      CopyMem (
+        &TxToken->Session.Udp6.SourceAddress,
+        &EndPoint->LocalAddr.v6,
+        sizeof(EFI_IPv6_ADDRESS)
+        );
+
+      CopyMem (
+        &TxToken->Session.Udp6.DestinationAddress,
+        &EndPoint->RemoteAddr.v6,
+        sizeof(EFI_IPv6_ADDRESS)
+        );
+
+      TxToken->Session.Udp6.SourcePort                   = EndPoint->LocalPort;
+      TxToken->Session.Udp6.DestinationPort              = EndPoint->RemotePort;
+      ((EFI_UDP6_TRANSMIT_DATA *) Data)->UdpSessionData  = &(TxToken->Session.Udp6);
+    }
+  }
+
+  return TxToken;
+}
 
 /**
-  Create a UDP IO port to access the UDP service. It will
-  create and configure a UDP child.
+  Creates a UDP_IO to access the UDP service. It creates and configures
+  a UDP child.
+
+  It locates the UDP service binding prototype on the Controller parameter
+  uses the UDP service binding prototype to create a UDP child (also known as
+  a UDP instance) configures the UDP child by calling Configure function prototype.
+  Any failures in creating or configuring the UDP child return NULL for failure.
 
-  @param  Controller            The controller that has the UDP service binding
-                                protocol installed.
-  @param  Image                 The image handle for the driver.
-  @param  Configure             The function to configure the created UDP child
-  @param  Context               The opaque parameter for the Configure funtion.
+  @param[in]  Controller            The controller that has the UDP service binding.
+                                    protocol installed.
+  @param[in]  ImageHandle           The image handle for the driver.
+  @param[in]  Configure             The function to configure the created UDP child.
+  @param[in]  UdpVersion            The UDP protocol version, UDP4 or UDP6.
+  @param[in]  Context               The opaque parameter for the Configure funtion.
 
-  @return A point to just created UDP IO port or NULL if failed.
+  @return Newly-created UDP_IO or NULL if failed.
 
 **/
-UDP_IO_PORT *
-UdpIoCreatePort (
+UDP_IO *
+EFIAPI
+UdpIoCreateIo (
   IN  EFI_HANDLE            Controller,
-  IN  EFI_HANDLE            Image,
+  IN  EFI_HANDLE            ImageHandle,
   IN  UDP_IO_CONFIG         Configure,
+  IN  UINT8                 UdpVersion,
   IN  VOID                  *Context
   )
 {
-  UDP_IO_PORT               *UdpIo;
+  UDP_IO                    *UdpIo;
   EFI_STATUS                Status;
 
   ASSERT (Configure != NULL);
+  ASSERT ((UdpVersion == UDP_IO_UDP4_VERSION) || (UdpVersion == UDP_IO_UDP6_VERSION));
 
-  UdpIo = NetAllocatePool (sizeof (UDP_IO_PORT));
+  UdpIo = AllocatePool (sizeof (UDP_IO));
 
   if (UdpIo == NULL) {
     return NULL;
   }
 
+  UdpIo->UdpVersion   = UdpVersion;
   UdpIo->Signature    = UDP_IO_SIGNATURE;
-  NetListInit (&UdpIo->Link);
+  InitializeListHead (&UdpIo->Link);
   UdpIo->RefCnt       = 1;
 
   UdpIo->Controller   = Controller;
-  UdpIo->Image        = Image;
+  UdpIo->Image        = ImageHandle;
 
-  NetListInit (&UdpIo->SentDatagram);
+  InitializeListHead (&UdpIo->SentDatagram);
   UdpIo->RecvRequest  = NULL;
   UdpIo->UdpHandle    = NULL;
 
-  //
-  // Create a UDP child then open and configure it
-  //
-  Status = NetLibCreateServiceChild (
-             Controller,
-             Image,
-             &gEfiUdp4ServiceBindingProtocolGuid,
-             &UdpIo->UdpHandle
-             );
+  if (UdpVersion == UDP_IO_UDP4_VERSION) {
+    //
+    // Create a UDP child then open and configure it
+    //
+    Status = NetLibCreateServiceChild (
+               Controller,
+               ImageHandle,
+               &gEfiUdp4ServiceBindingProtocolGuid,
+               &UdpIo->UdpHandle
+               );
+
+    if (EFI_ERROR (Status)) {
+      goto FREE_MEM;
+    }
 
-  if (EFI_ERROR (Status)) {
-    goto FREE_MEM;
-  }
+    Status = gBS->OpenProtocol (
+                    UdpIo->UdpHandle,
+                    &gEfiUdp4ProtocolGuid,
+                    (VOID **) &UdpIo->Protocol.Udp4,
+                    ImageHandle,
+                    Controller,
+                    EFI_OPEN_PROTOCOL_BY_DRIVER
+                    );
+
+    if (EFI_ERROR (Status)) {
+      goto FREE_CHILD;
+    }
 
-  Status = gBS->OpenProtocol (
-                  UdpIo->UdpHandle,
-                  &gEfiUdp4ProtocolGuid,
-                  &UdpIo->Udp,
-                  Image,
-                  Controller,
-                  EFI_OPEN_PROTOCOL_BY_DRIVER
-                  );
+    if (EFI_ERROR (Configure (UdpIo, Context))) {
+      goto CLOSE_PROTOCOL;
+    }
 
-  if (EFI_ERROR (Status)) {
-    goto FREE_CHILD;
-  }
+    Status = UdpIo->Protocol.Udp4->GetModeData (
+                                     UdpIo->Protocol.Udp4,
+                                     NULL,
+                                     NULL,
+                                     NULL,
+                                     &UdpIo->SnpMode
+                                     );
 
-  if (EFI_ERROR (Configure (UdpIo, Context))) {
-    goto CLOSE_PROTOCOL;
-  }
+    if (EFI_ERROR (Status)) {
+      goto CLOSE_PROTOCOL;
+    }
 
-  Status = UdpIo->Udp->GetModeData (UdpIo->Udp, NULL, NULL, NULL, &UdpIo->SnpMode);
+  } else {
 
-  if (EFI_ERROR (Status)) {
-    goto CLOSE_PROTOCOL;
+    Status = NetLibCreateServiceChild (
+               Controller,
+               ImageHandle,
+               &gEfiUdp6ServiceBindingProtocolGuid,
+               &UdpIo->UdpHandle
+               );
+
+    if (EFI_ERROR (Status)) {
+      goto FREE_MEM;
+    }
+
+    Status = gBS->OpenProtocol (
+                    UdpIo->UdpHandle,
+                    &gEfiUdp6ProtocolGuid,
+                    (VOID **) &UdpIo->Protocol.Udp6,
+                    ImageHandle,
+                    Controller,
+                    EFI_OPEN_PROTOCOL_BY_DRIVER
+                    );
+
+    if (EFI_ERROR (Status)) {
+      goto FREE_CHILD;
+    }
+
+    if (EFI_ERROR (Configure (UdpIo, Context))) {
+      goto CLOSE_PROTOCOL;
+    }
+
+    Status = UdpIo->Protocol.Udp6->GetModeData (
+                                     UdpIo->Protocol.Udp6,
+                                     NULL,
+                                     NULL,
+                                     NULL,
+                                     &UdpIo->SnpMode
+                                     );
+
+    if (EFI_ERROR (Status)) {
+      goto CLOSE_PROTOCOL;
+    }
   }
 
   return UdpIo;
 
 CLOSE_PROTOCOL:
-  gBS->CloseProtocol (UdpIo->UdpHandle, &gEfiUdp4ProtocolGuid, Image, Controller);
+  if (UdpVersion == UDP_IO_UDP4_VERSION) {
+    gBS->CloseProtocol (UdpIo->UdpHandle, &gEfiUdp4ProtocolGuid, ImageHandle, Controller);
+  } else {
+    gBS->CloseProtocol (UdpIo->UdpHandle, &gEfiUdp6ProtocolGuid, ImageHandle, Controller);
+  }
 
 FREE_CHILD:
-  NetLibDestroyServiceChild (
-    Controller,
-    Image,
-    &gEfiUdp4ServiceBindingProtocolGuid,
-    UdpIo->UdpHandle
-    );
+  if (UdpVersion == UDP_IO_UDP4_VERSION) {
+    NetLibDestroyServiceChild (
+      Controller,
+      ImageHandle,
+      &gEfiUdp4ServiceBindingProtocolGuid,
+      UdpIo->UdpHandle
+      );
+  } else {
+    NetLibDestroyServiceChild (
+      Controller,
+      ImageHandle,
+      &gEfiUdp6ServiceBindingProtocolGuid,
+      UdpIo->UdpHandle
+      );
+  }
 
 FREE_MEM:
-  NetFreePool (UdpIo);
+  FreePool (UdpIo);
   return NULL;
 }
 
-
 /**
-  Cancel all the sent datagram that pass the selection of ToCancel.
+  Cancel all the sent datagram that pass the selection criteria of ToCancel.
   If ToCancel is NULL, all the datagrams are cancelled.
 
-  @param  UdpIo                 The UDP IO port to cancel packet
-  @param  IoStatus              The IoStatus to return to the packet owners.
-  @param  ToCancel              The select funtion to test whether to cancel this
-                                packet or not.
-  @param  Context               The opaque parameter to the ToCancel.
-
-  @return None
+  @param[in]  UdpIo                 The UDP_IO to cancel packet.
+  @param[in]  IoStatus              The IoStatus to return to the packet owners.
+  @param[in]  ToCancel              The select funtion to test whether to cancel this
+                                    packet or not.
+  @param[in]  Context               The opaque parameter to the ToCancel.
 
 **/
-STATIC
 VOID
+EFIAPI
 UdpIoCancelDgrams (
-  IN UDP_IO_PORT            *UdpIo,
+  IN UDP_IO                 *UdpIo,
   IN EFI_STATUS             IoStatus,
   IN UDP_IO_TO_CANCEL       ToCancel,        OPTIONAL
   IN VOID                   *Context
   )
 {
-  NET_LIST_ENTRY            *Entry;
-  NET_LIST_ENTRY            *Next;
-  UDP_TX_TOKEN              *Token;
+  LIST_ENTRY                *Entry;
+  LIST_ENTRY                *Next;
+  UDP_TX_TOKEN              *TxToken;
+
+  ASSERT ((UdpIo->UdpVersion == UDP_IO_UDP4_VERSION) ||
+          (UdpIo->UdpVersion == UDP_IO_UDP6_VERSION));
 
   NET_LIST_FOR_EACH_SAFE (Entry, Next, &UdpIo->SentDatagram) {
-    Token = NET_LIST_USER_STRUCT (Entry, UDP_TX_TOKEN, Link);
+    TxToken = NET_LIST_USER_STRUCT (Entry, UDP_TX_TOKEN, Link);
 
-    if ((ToCancel == NULL) || (ToCancel (Token, Context))) {
-      NetListRemoveEntry (Entry);
-      UdpIo->Udp->Cancel (UdpIo->Udp, &Token->UdpToken);
-      Token->CallBack (Token->Packet, NULL, IoStatus, Token->Context);
-      UdpIoFreeTxToken (Token);
+    if ((ToCancel == NULL) || (ToCancel (TxToken, Context))) {
+
+      if (UdpIo->UdpVersion == UDP_IO_UDP4_VERSION) {
+        UdpIo->Protocol.Udp4->Cancel (UdpIo->Protocol.Udp4, &TxToken->Token.Udp4);
+      } else {
+        UdpIo->Protocol.Udp6->Cancel (UdpIo->Protocol.Udp6, &TxToken->Token.Udp6);
+      }
     }
   }
 }
 
-
 /**
-  Free the UDP IO port and all its related resources including
-  all the transmitted packet.
+  Free the UDP_IO and all its related resources.
 
-  @param  UdpIo                 The UDP IO port to free.
+  The function will cancel all sent datagram and receive request.
 
-  @retval EFI_SUCCESS           The UDP IO port is freed.
+  @param[in]  UdpIo             The UDP_IO to free.
+
+  @retval EFI_SUCCESS           The UDP_IO is freed.
 
 **/
 EFI_STATUS
-UdpIoFreePort (
-  IN  UDP_IO_PORT           *UdpIo
+EFIAPI
+UdpIoFreeIo (
+  IN  UDP_IO           *UdpIo
   )
 {
-  UDP_RX_TOKEN  *RxToken;
+  UDP_RX_TOKEN         *RxToken;
+
+  ASSERT ((UdpIo->UdpVersion == UDP_IO_UDP4_VERSION) ||
+          (UdpIo->UdpVersion == UDP_IO_UDP6_VERSION));
 
   //
   // Cancel all the sent datagram and receive requests. The
@@ -391,156 +797,181 @@ UdpIoFreePort (
   //
   UdpIoCancelDgrams (UdpIo, EFI_ABORTED, NULL, NULL);
 
-  if ((RxToken = UdpIo->RecvRequest) != NULL) {
-    UdpIo->RecvRequest = NULL;
-    UdpIo->Udp->Cancel (UdpIo->Udp, &RxToken->UdpToken);
-    UdpIoFreeRxToken (RxToken);
+  if (UdpIo->UdpVersion == UDP_IO_UDP4_VERSION) {
+
+    if ((RxToken = UdpIo->RecvRequest) != NULL) {
+      UdpIo->Protocol.Udp4->Cancel (UdpIo->Protocol.Udp4, &RxToken->Token.Udp4);
+    }
+
+    //
+    // Close then destory the Udp4 child
+    //
+    gBS->CloseProtocol (
+           UdpIo->UdpHandle,
+           &gEfiUdp4ProtocolGuid,
+           UdpIo->Image,
+           UdpIo->Controller
+           );
+
+    NetLibDestroyServiceChild (
+      UdpIo->Controller,
+      UdpIo->Image,
+      &gEfiUdp4ServiceBindingProtocolGuid,
+      UdpIo->UdpHandle
+      );
+
+  } else {
+
+    if ((RxToken = UdpIo->RecvRequest) != NULL) {
+      UdpIo->Protocol.Udp6->Cancel (UdpIo->Protocol.Udp6, &RxToken->Token.Udp6);
+    }
+
+    //
+    // Close then destory the Udp6 child
+    //
+    gBS->CloseProtocol (
+           UdpIo->UdpHandle,
+           &gEfiUdp6ProtocolGuid,
+           UdpIo->Image,
+           UdpIo->Controller
+           );
+
+    NetLibDestroyServiceChild (
+      UdpIo->Controller,
+      UdpIo->Image,
+      &gEfiUdp6ServiceBindingProtocolGuid,
+      UdpIo->UdpHandle
+      );
+    }
+
+  if (!IsListEmpty(&UdpIo->Link)) {
+    RemoveEntryList (&UdpIo->Link);
   }
 
-  //
-  // Close then destory the UDP child
-  //
-  gBS->CloseProtocol (
-         UdpIo->UdpHandle,
-         &gEfiUdp4ProtocolGuid,
-         UdpIo->Image,
-         UdpIo->Controller
-         );
-
-  NetLibDestroyServiceChild (
-    UdpIo->Controller,
-    UdpIo->Image,
-    &gEfiUdp4ServiceBindingProtocolGuid,
-    UdpIo->UdpHandle
-    );
-
-  NetListRemoveEntry (&UdpIo->Link);
-  NetFreePool (UdpIo);
+  FreePool (UdpIo);
   return EFI_SUCCESS;
 }
 
 
 /**
-  Clean up the UDP IO port. It will release all the transmitted
-  datagrams and receive request. It will also configure NULL the
-  UDP child.
+  Clean up the UDP_IO without freeing it. The function is called when
+  user wants to re-use the UDP_IO later.
 
-  @param  UdpIo                 UDP IO port to clean up.
+  It will release all the transmitted datagrams and receive request. It will
+  also configure NULL for the UDP instance.
 
-  @return None
+  @param[in]  UdpIo                 The UDP_IO to clean up.
 
 **/
 VOID
-UdpIoCleanPort (
-  IN  UDP_IO_PORT           *UdpIo
+EFIAPI
+UdpIoCleanIo (
+  IN  UDP_IO                *UdpIo
   )
 {
   UDP_RX_TOKEN              *RxToken;
 
+  ASSERT ((UdpIo->UdpVersion == UDP_IO_UDP4_VERSION) ||
+          (UdpIo->UdpVersion == UDP_IO_UDP6_VERSION));
+
   //
   // Cancel all the sent datagram and receive requests.
   //
   UdpIoCancelDgrams (UdpIo, EFI_ABORTED, NULL, NULL);
 
-  if ((RxToken = UdpIo->RecvRequest) != NULL) {
-    UdpIo->RecvRequest = NULL;
-    UdpIo->Udp->Cancel (UdpIo->Udp, &RxToken->UdpToken);
-    UdpIoFreeRxToken (RxToken);
-  }
-
-  UdpIo->Udp->Configure (UdpIo->Udp, NULL);
-}
-
-
-/**
-  The callback function when the packet is sent by UDP.
-  It will remove the packet from the local list then call
-  the packet owner's callback function.
-
-  @param  Event                 The event signalled.
-  @param  Context               The UDP TX Token.
-
-  @return None
-
-**/
-STATIC
-VOID
-EFIAPI
-UdpIoOnDgramSent (
-  IN EFI_EVENT              Event,
-  IN VOID                   *Context
-  )
-{
-  UDP_TX_TOKEN              *Token;
+  if (UdpIo->UdpVersion == UDP_IO_UDP4_VERSION) {
+    if ((RxToken = UdpIo->RecvRequest) != NULL) {
+      UdpIo->Protocol.Udp4->Cancel (UdpIo->Protocol.Udp4, &RxToken->Token.Udp4);
+    }
 
-  Token   = (UDP_TX_TOKEN *) Context;
-  ASSERT (Token->Signature == UDP_IO_TX_SIGNATURE);
+    UdpIo->Protocol.Udp4->Configure (UdpIo->Protocol.Udp4, NULL);
 
-  NetListRemoveEntry (&Token->Link);
-  Token->CallBack (Token->Packet, NULL, Token->UdpToken.Status, Token->Context);
+  } else {
+    if ((RxToken = UdpIo->RecvRequest) != NULL) {
+      UdpIo->Protocol.Udp6->Cancel (UdpIo->Protocol.Udp6, &RxToken->Token.Udp6);
+    }
 
-  UdpIoFreeTxToken (Token);
+    UdpIo->Protocol.Udp6->Configure (UdpIo->Protocol.Udp6, NULL);
+  }
 }
 
-
 /**
-  Send a packet through the UDP IO port.
-
-  @param  UdpIo                 The UDP IO Port to send the packet through
-  @param  Packet                The packet to send
-  @param  EndPoint              The local and remote access point
-  @param  Gateway               The gateway to use
-  @param  CallBack              The call back function to call when packet is
-                                transmitted or failed.
-  @param  Context               The opque parameter to the CallBack
-
-  @retval EFI_OUT_OF_RESOURCES  Failed to allocate resource for the packet
+  Send a packet through the UDP_IO.
+
+  The packet will be wrapped in UDP_TX_TOKEN. Function Callback will be called
+  when the packet is sent. The optional parameter EndPoint overrides the default
+  address pair if specified.
+
+  @param[in]  UdpIo                 The UDP_IO to send the packet through.
+  @param[in]  Packet                The packet to send.
+  @param[in]  EndPoint              The local and remote access point. Override the
+                                    default address pair set during configuration.
+  @param[in]  Gateway               The gateway to use.
+  @param[in]  CallBack              The function being called when packet is
+                                    transmitted or failed.
+  @param[in]  Context               The opaque parameter passed to CallBack.
+
+  @retval EFI_OUT_OF_RESOURCES  Failed to allocate resource for the packet.
   @retval EFI_SUCCESS           The packet is successfully delivered to UDP  for
                                 transmission.
 
 **/
 EFI_STATUS
+EFIAPI
 UdpIoSendDatagram (
-  IN  UDP_IO_PORT           *UdpIo,
+  IN  UDP_IO                *UdpIo,
   IN  NET_BUF               *Packet,
-  IN  UDP_POINTS            *EndPoint, OPTIONAL
-  IN  IP4_ADDR              Gateway,
+  IN  UDP_END_POINT         *EndPoint OPTIONAL,
+  IN  EFI_IP_ADDRESS        *Gateway  OPTIONAL,
   IN  UDP_IO_CALLBACK       CallBack,
   IN  VOID                  *Context
   )
 {
-  UDP_TX_TOKEN              *Token;
+  UDP_TX_TOKEN              *TxToken;
   EFI_STATUS                Status;
 
-  Token = UdpIoWrapTx (UdpIo, Packet, EndPoint, Gateway, CallBack, Context);
+  ASSERT ((UdpIo->UdpVersion == UDP_IO_UDP4_VERSION) ||
+          (UdpIo->UdpVersion == UDP_IO_UDP6_VERSION));
 
-  if (Token == NULL) {
+  TxToken = UdpIoCreateTxToken (UdpIo, Packet, EndPoint, Gateway, CallBack, Context);
+
+  if (TxToken == NULL) {
     return EFI_OUT_OF_RESOURCES;
   }
 
-  Status = UdpIo->Udp->Transmit (UdpIo->Udp, &Token->UdpToken);
+  //
+  // Insert the tx token into SendDatagram list before transmitting it. Remove
+  // it from the list if the returned status is not EFI_SUCCESS.
+  //
+  InsertHeadList (&UdpIo->SentDatagram, &TxToken->Link);
+
+  if (UdpIo->UdpVersion == UDP_IO_UDP4_VERSION) {
+    Status = UdpIo->Protocol.Udp4->Transmit (UdpIo->Protocol.Udp4, &TxToken->Token.Udp4);
+  } else {
+    Status = UdpIo->Protocol.Udp6->Transmit (UdpIo->Protocol.Udp6, &TxToken->Token.Udp6);
+  }
 
   if (EFI_ERROR (Status)) {
-    UdpIoFreeTxToken (Token);
+    RemoveEntryList (&TxToken->Link);
+    UdpIoFreeTxToken (TxToken);
     return Status;
   }
 
-  NetListInsertHead (&UdpIo->SentDatagram, &Token->Link);
   return EFI_SUCCESS;
 }
 
 
 /**
-  The selection function to cancel a single sent datagram.
+  The select function to cancel a single sent datagram.
 
-  @param  Token                 The UDP TX token to test againist.
-  @param  Context               The context
-
-  @return TRUE if the packet is to be cancelled, otherwise FALSE.
+  @param[in]  Token                 The UDP_TX_TOKEN to test against
+  @param[in]  Context               The NET_BUF of the sent datagram
 
+  @retval TRUE              The packet is to be cancelled.
+  @retval FALSE             The packet is not to be cancelled.
 **/
-STATIC
 BOOLEAN
+EFIAPI
 UdpIoCancelSingleDgram (
   IN UDP_TX_TOKEN           *Token,
   IN VOID                   *Context
@@ -557,171 +988,80 @@ UdpIoCancelSingleDgram (
   return FALSE;
 }
 
-
 /**
   Cancel a single sent datagram.
 
-  @param  UdpIo                 The UDP IO port to cancel the packet from
-  @param  Packet                The packet to cancel
-
-  @return None
+  @param[in]  UdpIo                 The UDP_IO to cancel the packet from
+  @param[in]  Packet                The packet to cancel
 
 **/
 VOID
+EFIAPI
 UdpIoCancelSentDatagram (
-  IN  UDP_IO_PORT           *UdpIo,
+  IN  UDP_IO                *UdpIo,
   IN  NET_BUF               *Packet
   )
 {
   UdpIoCancelDgrams (UdpIo, EFI_ABORTED, UdpIoCancelSingleDgram, Packet);
 }
 
-
-/**
-  Recycle the received UDP data.
-
-  @param  Context               The UDP_RX_TOKEN
-
-  @return None
-
-**/
-STATIC
-VOID
-UdpIoRecycleDgram (
-  IN VOID                   *Context
-  )
-{
-  UDP_RX_TOKEN              *Token;
-
-  Token = (UDP_RX_TOKEN *) Context;
-  gBS->SignalEvent (Token->UdpToken.Packet.RxData->RecycleSignal);
-  UdpIoFreeRxToken (Token);
-}
-
-
 /**
-  The event handle for UDP receive request. It will build
-  a NET_BUF from the recieved UDP data, then deliver it
-  to the receiver.
-
-  @param  Event                 The UDP receive request event
-  @param  Context               The UDP RX token.
-
-  @return None
-
-**/
-STATIC
-VOID
-EFIAPI
-UdpIoOnDgramRcvd (
-  IN EFI_EVENT              Event,
-  IN VOID                   *Context
-  )
-{
-  EFI_UDP4_COMPLETION_TOKEN *UdpToken;
-  EFI_UDP4_RECEIVE_DATA     *UdpRxData;
-  EFI_UDP4_SESSION_DATA     *UdpSession;
-  UDP_RX_TOKEN              *Token;
-  UDP_POINTS                Points;
-  NET_BUF                   *Netbuf;
-
-  Token   = (UDP_RX_TOKEN *) Context;
-
-  ASSERT ((Token->Signature == UDP_IO_RX_SIGNATURE) &&
-          (Token == Token->UdpIo->RecvRequest));
-
-  //
-  // Clear the receive request first in case that the caller
-  // wants to restart the receive in the callback.
-  //
-  Token->UdpIo->RecvRequest = NULL;
-
-  UdpToken  = &Token->UdpToken;
-  UdpRxData = UdpToken->Packet.RxData;
-
-  if (EFI_ERROR (UdpToken->Status) || (UdpRxData == NULL)) {
-    Token->CallBack (NULL, NULL, UdpToken->Status, Token->Context);
-    UdpIoFreeRxToken (Token);
-
-    goto ON_EXIT;
-  }
-
-  //
-  // Build a NET_BUF from the UDP receive data, then deliver it up.
-  //
-  Netbuf = NetbufFromExt (
-             (NET_FRAGMENT *) UdpRxData->FragmentTable,
-             UdpRxData->FragmentCount,
-             0,
-             (UINT32) Token->HeadLen,
-             UdpIoRecycleDgram,
-             Token
-             );
-
-  if (Netbuf == NULL) {
-    gBS->SignalEvent (UdpRxData->RecycleSignal);
-    Token->CallBack (NULL, NULL, EFI_OUT_OF_RESOURCES, Token->Context);
-
-    UdpIoFreeRxToken (Token);
-    goto ON_EXIT;
-  }
-
-  UdpSession        = &UdpRxData->UdpSession;
-  Points.LocalAddr  = EFI_NTOHL (UdpSession->DestinationAddress);
-  Points.LocalPort  = UdpSession->DestinationPort;
-  Points.RemoteAddr = EFI_NTOHL (UdpSession->SourceAddress);
-  Points.RemotePort = UdpSession->SourcePort;
-
-  Token->CallBack (Netbuf, &Points, EFI_SUCCESS, Token->Context);
-
-ON_EXIT:
-  return;
-}
+  Issue a receive request to the UDP_IO.
 
+  This function is called when upper-layer needs packet from UDP for processing.
+  Only one receive request is acceptable at a time so a common usage model is
+  to invoke this function inside its Callback function when the former packet
+  is processed.
 
-/**
-  Issue a receive request to the UDP IO port.
-
-  @param  UdpIo                 The UDP IO port to recieve the packet from.
-  @param  CallBack              The call back function to execute when receive
-                                finished.
-  @param  Context               The opque context to the call back
-  @param  HeadLen               The lenght of the application's header
+  @param[in]  UdpIo                 The UDP_IO to receive the packet from.
+  @param[in]  CallBack              The call back function to execute when the packet
+                                    is received.
+  @param[in]  Context               The opaque context passed to Callback.
+  @param[in]  HeadLen               The length of the upper-layer's protocol header.
 
   @retval EFI_ALREADY_STARTED   There is already a pending receive request. Only
-                                one receive request is supported.
-  @retval EFI_OUT_OF_RESOURCES  Failed to allocate some resource.
+                                one receive request is supported at a time.
+  @retval EFI_OUT_OF_RESOURCES  Failed to allocate needed resources.
   @retval EFI_SUCCESS           The receive request is issued successfully.
+  @retval EFI_UNSUPPORTED       The UDP version in UDP_IO is not supported.
 
 **/
 EFI_STATUS
+EFIAPI
 UdpIoRecvDatagram (
-  IN  UDP_IO_PORT           *UdpIo,
+  IN  UDP_IO                *UdpIo,
   IN  UDP_IO_CALLBACK       CallBack,
   IN  VOID                  *Context,
   IN  UINT32                HeadLen
   )
 {
-  UDP_RX_TOKEN              *Token;
+  UDP_RX_TOKEN              *RxToken;
   EFI_STATUS                Status;
 
+  ASSERT ((UdpIo->UdpVersion == UDP_IO_UDP4_VERSION) ||
+          (UdpIo->UdpVersion == UDP_IO_UDP6_VERSION));
+
   if (UdpIo->RecvRequest != NULL) {
     return EFI_ALREADY_STARTED;
   }
 
-  Token = UdpIoCreateRxToken (UdpIo, CallBack, Context, HeadLen);
+  RxToken = UdpIoCreateRxToken (UdpIo, CallBack, Context, HeadLen);
 
-  if (Token == NULL) {
+  if (RxToken == NULL) {
     return EFI_OUT_OF_RESOURCES;
   }
 
-  Status = UdpIo->Udp->Receive (UdpIo->Udp, &Token->UdpToken);
+  UdpIo->RecvRequest = RxToken;
+  if (UdpIo->UdpVersion == UDP_IO_UDP4_VERSION) {
+    Status = UdpIo->Protocol.Udp4->Receive (UdpIo->Protocol.Udp4, &RxToken->Token.Udp4);
+  } else {
+    Status = UdpIo->Protocol.Udp6->Receive (UdpIo->Protocol.Udp6, &RxToken->Token.Udp6);
+  }
 
   if (EFI_ERROR (Status)) {
-    UdpIoFreeRxToken (Token);
-    return Status;
+    UdpIo->RecvRequest = NULL;
+    UdpIoFreeRxToken (RxToken);
   }
 
-  UdpIo->RecvRequest = Token;
-  return EFI_SUCCESS;
+  return Status;
 }