[Patch 4/4] NetworkPkg: Add more parameter or return status check in UDP6 driver

[fanwang2 <fan.wang@intel.com>]

From: Wang Fan <fan.wang@intel.com>

In UDP6Dxe, there are several places that may be enhanced
to check input parameters and returned status. This patch
is to fix these issues.

Cc: Ye Ting <ting.ye@intel.com>
Cc: Jiaxin Wu <jiaxin.wu@intel.com>
Cc: Fu Siyuan <siyuan.fu@intel.com>
Contributed-under: TianoCore Contribution Agreement 1.0
Signed-off-by: Wang Fan <fan.wang@intel.com>
---
 NetworkPkg/Udp6Dxe/Udp6Driver.c | 48 +++++++++++++++++++++++------------------
 NetworkPkg/Udp6Dxe/Udp6Impl.c   | 17 +++++++++++++++
 NetworkPkg/Udp6Dxe/Udp6Main.c   |  2 +-
 3 files changed, 45 insertions(+), 22 deletions(-)

diff --git a/NetworkPkg/Udp6Dxe/Udp6Driver.c b/NetworkPkg/Udp6Dxe/Udp6Driver.c
index 6dde1fc..f9d528e 100644
--- a/NetworkPkg/Udp6Dxe/Udp6Driver.c
+++ b/NetworkPkg/Udp6Dxe/Udp6Driver.c
@@ -288,22 +288,19 @@ Udp6DriverBindingStop (
                Udp6DestroyChildEntryInHandleBuffer,
                &Context,
                NULL
                );
   } else if (IsListEmpty (&Udp6Service->ChildrenList)) {
-    gBS->UninstallMultipleProtocolInterfaces (
-           NicHandle,
-           &gEfiUdp6ServiceBindingProtocolGuid,
-           &Udp6Service->ServiceBinding,
-           NULL
-           );
+    Status = gBS->UninstallMultipleProtocolInterfaces (
+               NicHandle,
+               &gEfiUdp6ServiceBindingProtocolGuid,
+               &Udp6Service->ServiceBinding,
+               NULL
+               );
  
     Udp6CleanService (Udp6Service);
-
     FreePool (Udp6Service);
-
-    Status = EFI_SUCCESS;
   }
 
   return Status;
 }
 
@@ -508,25 +505,34 @@ Udp6ServiceBindingDestroyChild (
   Instance->InDestroy = TRUE;
 
   //
   // Close the Ip6 protocol on the default IpIo.
   //
-  gBS->CloseProtocol (
-         Udp6Service->IpIo->ChildHandle,
-         &gEfiIp6ProtocolGuid,
-         gUdp6DriverBinding.DriverBindingHandle,
-         Instance->ChildHandle
-         );
+  Status = gBS->CloseProtocol (
+             Udp6Service->IpIo->ChildHandle,
+             &gEfiIp6ProtocolGuid,
+             gUdp6DriverBinding.DriverBindingHandle,
+             Instance->ChildHandle
+             );
+  if (EFI_ERROR (Status)) {
+    Instance->InDestroy = FALSE;
+    return Status;
+  }
+
   //
   // Close the Ip6 protocol on this instance's IpInfo.
   //
-  gBS->CloseProtocol (
-         Instance->IpInfo->ChildHandle,
-         &gEfiIp6ProtocolGuid,
-         gUdp6DriverBinding.DriverBindingHandle,
-         Instance->ChildHandle
-         );
+  Status = gBS->CloseProtocol (
+             Instance->IpInfo->ChildHandle,
+             &gEfiIp6ProtocolGuid,
+             gUdp6DriverBinding.DriverBindingHandle,
+             Instance->ChildHandle
+             );
+  if (EFI_ERROR (Status)) {
+    Instance->InDestroy = FALSE;
+    return Status;
+  }
 
   //
   // Uninstall the Udp6Protocol previously installed on the ChildHandle.
   //
   Status = gBS->UninstallMultipleProtocolInterfaces (
diff --git a/NetworkPkg/Udp6Dxe/Udp6Impl.c b/NetworkPkg/Udp6Dxe/Udp6Impl.c
index 458470c..d014e2d 100644
--- a/NetworkPkg/Udp6Dxe/Udp6Impl.c
+++ b/NetworkPkg/Udp6Dxe/Udp6Impl.c
@@ -55,10 +55,13 @@ Udp6FindInstanceByPort (
 /**
   This function is the packet transmitting notify function registered to the IpIo
   interface. It's called to signal the udp TxToken when the IpIo layer completes
   transmitting of the udp datagram.
 
+  If Context is NULL, then ASSERT().
+  If NotifyData is NULL, then ASSERT().
+
   @param[in]  Status            The completion status of the output udp datagram.
   @param[in]  Context           Pointer to the context data.
   @param[in]  Sender            Specify a EFI_IP6_PROTOCOL for sending.
   @param[in]  NotifyData        Pointer to the notify data.
 
@@ -73,10 +76,14 @@ Udp6DgramSent (
   );
 
 /**
   This function processes the received datagram passed up by the IpIo layer.
 
+  If NetSession is NULL, then ASSERT().
+  If Packet is NULL, then ASSERT().
+  If Context is NULL, then ASSERT().
+
   @param[in]  Status            The status of this udp datagram.
   @param[in]  IcmpError         The IcmpError code, only available when Status is
                                 EFI_ICMP_ERROR.
   @param[in]  NetSession        Pointer to the EFI_NET_SESSION_DATA.
   @param[in]  Packet            Pointer to the NET_BUF containing the received udp
@@ -975,10 +982,13 @@ Udp6RemoveToken (
 /**
   This function is the packet transmitting notify function registered to the IpIo
   interface. It's called to signal the udp TxToken when IpIo layer completes the
   transmitting of the udp datagram.
 
+  If Context is NULL, then ASSERT().
+  If NotifyData is NULL, then ASSERT().
+
   @param[in]  Status            The completion status of the output udp datagram.
   @param[in]  Context           Pointer to the context data.
   @param[in]  Sender            Specify a EFI_IP6_PROTOCOL for sending.
   @param[in]  NotifyData        Pointer to the notify data.
 
@@ -993,10 +1003,12 @@ Udp6DgramSent (
   )
 {
   UDP6_INSTANCE_DATA         *Instance;
   EFI_UDP6_COMPLETION_TOKEN  *Token;
 
+  ASSERT (Context != NULL && NotifyData != NULL);
+
   Instance = (UDP6_INSTANCE_DATA *) Context;
   Token    = (EFI_UDP6_COMPLETION_TOKEN *) NotifyData;
 
   if (Udp6RemoveToken (&Instance->TxTokens, Token) == EFI_SUCCESS) {
     //
@@ -1010,10 +1022,14 @@ Udp6DgramSent (
 
 
 /**
   This function processes the received datagram passed up by the IpIo layer.
 
+  If NetSession is NULL, then ASSERT().
+  If Packet is NULL, then ASSERT().
+  If Context is NULL, then ASSERT().
+
   @param[in]  Status            The status of this udp datagram.
   @param[in]  IcmpError         The IcmpError code, only available when Status is
                                 EFI_ICMP_ERROR.
   @param[in]  NetSession        Pointer to the EFI_NET_SESSION_DATA.
   @param[in]  Packet            Pointer to the NET_BUF containing the received udp
@@ -1029,10 +1045,11 @@ Udp6DgramRcvd (
   IN EFI_NET_SESSION_DATA  *NetSession,
   IN NET_BUF               *Packet,
   IN VOID                  *Context
   )
 {
+  ASSERT (NetSession != NULL && Packet != NULL && Context != NULL);
   NET_CHECK_SIGNATURE (Packet, NET_BUF_SIGNATURE);
 
   //
   // IpIo only passes received packets with Status EFI_SUCCESS or EFI_ICMP_ERROR.
   //
diff --git a/NetworkPkg/Udp6Dxe/Udp6Main.c b/NetworkPkg/Udp6Dxe/Udp6Main.c
index 53145c3..1d7f0ac 100644
--- a/NetworkPkg/Udp6Dxe/Udp6Main.c
+++ b/NetworkPkg/Udp6Dxe/Udp6Main.c
@@ -379,11 +379,11 @@ Udp6Groups (
   if (JoinFlag) {
 
     Status = NetMapInsertTail (&Instance->McastIps, (VOID *) McastIp, NULL);
   } else {
 
-    NetMapIterate (&Instance->McastIps, Udp6LeaveGroup, MulticastAddress);
+    Status = NetMapIterate (&Instance->McastIps, Udp6LeaveGroup, MulticastAddress);
   }
 
 ON_EXIT:
 
   gBS->RestoreTPL (OldTpl);
-- 
1.9.5.msysgit.1