Re: [edk2] [PATCH v2] NetworkPkg: Replace ASSERT with error handling in DnsDxe

2016-06-21 Thread Zhang, Lubo
Reviewed-by: Zhang Lubo 

-Original Message-
From: Wu, Jiaxin 
Sent: Monday, June 20, 2016 9:46 AM
To: edk2-devel@lists.01.org
Cc: Ye, Ting ; Fu, Siyuan ; Zhang, Lubo 

Subject: [PATCH v2] NetworkPkg: Replace ASSERT with error handling in DnsDxe

v2:
*Use goto to simplify code logic.

This patch is used to replace ASSERT with error handling in DnsDxe driver.

Cc: Ye Ting 
Cc: Fu Siyuan 
Cc: Zhang Lubo 
Contributed-under: TianoCore Contribution Agreement 1.0
Signed-off-by: Jiaxin Wu 
---
 NetworkPkg/DnsDxe/DnsProtocol.c | 56 -
 1 file changed, 38 insertions(+), 18 deletions(-)

diff --git a/NetworkPkg/DnsDxe/DnsProtocol.c b/NetworkPkg/DnsDxe/DnsProtocol.c 
index e9101d6..64fca6a 100644
--- a/NetworkPkg/DnsDxe/DnsProtocol.c
+++ b/NetworkPkg/DnsDxe/DnsProtocol.c
@@ -86,23 +86,22 @@ Dns4GetModeData (
 
   OldTpl = gBS->RaiseTPL (TPL_CALLBACK);
 
   Instance = DNS_INSTANCE_FROM_THIS_PROTOCOL4 (This);
   if (Instance->State == DNS_STATE_UNCONFIGED) {
-gBS->RestoreTPL (OldTpl);
-return  EFI_NOT_STARTED;
+Status = EFI_NOT_STARTED;
+goto ON_EXIT;
   }
   
   ZeroMem (DnsModeData, sizeof (EFI_DNS4_MODE_DATA));
 
   //
   // Get the current configuration data of this instance. 
   //
   Status = Dns4CopyConfigure (>DnsConfigData, 
>Dns4CfgData);
   if (EFI_ERROR (Status)) {
-gBS->RestoreTPL (OldTpl);
-return Status;
+goto ON_EXIT;
   }
 
   //
   // Get the DnsServerCount and DnsServerList
   //
@@ -110,11 +109,16 @@ Dns4GetModeData (
   NET_LIST_FOR_EACH_SAFE (Entry, Next, >Dns4ServerList) {
 Index++;
   }
   DnsModeData->DnsServerCount = (UINT32) Index;
   ServerList = AllocatePool (sizeof (EFI_IPv4_ADDRESS) * 
DnsModeData->DnsServerCount);
-  ASSERT (ServerList != NULL);
+  if (ServerList == NULL) {
+Status = EFI_OUT_OF_RESOURCES;
+Dns4CleanConfigure (>DnsConfigData);
+goto ON_EXIT;
+  }
+  
   Index = 0;
   NET_LIST_FOR_EACH_SAFE (Entry, Next, >Dns4ServerList) {
 ServerItem = NET_LIST_USER_STRUCT (Entry, DNS4_SERVER_IP, AllServerLink);
 CopyMem (ServerList + Index, >Dns4ServerIp, sizeof 
(EFI_IPv4_ADDRESS));
 Index++;
@@ -128,22 +132,28 @@ Dns4GetModeData (
   NET_LIST_FOR_EACH_SAFE (Entry, Next, >Dns4CacheList) {
 Index++;
   }
   DnsModeData->DnsCacheCount = (UINT32) Index;
   CacheList = AllocatePool (sizeof (EFI_DNS4_CACHE_ENTRY) * 
DnsModeData->DnsCacheCount);
-  ASSERT (CacheList != NULL);
+  if (CacheList == NULL) {
+Status = EFI_OUT_OF_RESOURCES;
+Dns4CleanConfigure (>DnsConfigData);
+FreePool (ServerList);
+goto ON_EXIT;
+  }
+  
   Index =0;
   NET_LIST_FOR_EACH_SAFE (Entry, Next, >Dns4CacheList) {
 CacheItem = NET_LIST_USER_STRUCT (Entry, DNS4_CACHE, AllCacheLink);
 CopyMem (CacheList + Index, >DnsCache, sizeof 
(EFI_DNS4_CACHE_ENTRY));
 Index++;
   }
   DnsModeData->DnsCacheList = CacheList;
 
+ON_EXIT:
   gBS->RestoreTPL (OldTpl);
-
-  return EFI_SUCCESS;
+  return Status;
 }
 
 /**
   Configure this DNS instance.
 
@@ -907,23 +917,22 @@ Dns6GetModeData (
 
   OldTpl = gBS->RaiseTPL (TPL_CALLBACK);
 
   Instance  = DNS_INSTANCE_FROM_THIS_PROTOCOL6 (This);
   if (Instance->State == DNS_STATE_UNCONFIGED) {
-gBS->RestoreTPL (OldTpl);
-return  EFI_NOT_STARTED;
+Status =  EFI_NOT_STARTED;
+goto ON_EXIT;
   }
 
   ZeroMem (DnsModeData, sizeof (EFI_DNS6_MODE_DATA));
 
   //
   // Get the current configuration data of this instance. 
   //
-  Status = Dns6CopyConfigure(>DnsConfigData, 
>Dns6CfgData);
+  Status = Dns6CopyConfigure (>DnsConfigData, 
+ >Dns6CfgData);
   if (EFI_ERROR (Status)) {
-gBS->RestoreTPL (OldTpl);
-return Status;
+goto ON_EXIT;
   }
   
   //
   // Get the DnsServerCount and DnsServerList
   //
@@ -931,11 +940,16 @@ Dns6GetModeData (
   NET_LIST_FOR_EACH_SAFE (Entry, Next, >Dns6ServerList) {
 Index++;
   }
   DnsModeData->DnsServerCount = (UINT32) Index;
   ServerList = AllocatePool (sizeof(EFI_IPv6_ADDRESS) * 
DnsModeData->DnsServerCount);
-  ASSERT (ServerList != NULL);
+  if (ServerList == NULL) {
+Status = EFI_OUT_OF_RESOURCES;
+Dns6CleanConfigure (>DnsConfigData);
+goto ON_EXIT;
+  }
+  
   Index = 0;
   NET_LIST_FOR_EACH_SAFE (Entry, Next, >Dns6ServerList) {
 ServerItem = NET_LIST_USER_STRUCT (Entry, DNS6_SERVER_IP, AllServerLink);
 CopyMem (ServerList + Index, >Dns6ServerIp, sizeof 
(EFI_IPv6_ADDRESS));
 Index++;
@@ -949,22 +963,28 @@ Dns6GetModeData (
   NET_LIST_FOR_EACH_SAFE (Entry, Next, >Dns6CacheList) {
 Index++;
   }
   DnsModeData->DnsCacheCount = (UINT32) Index;
   CacheList = AllocatePool (sizeof(EFI_DNS6_CACHE_ENTRY) * 
DnsModeData->DnsCacheCount);
-  ASSERT (CacheList != NULL);
+  if (CacheList == NULL) {
+Status = EFI_OUT_OF_RESOURCES;
+Dns6CleanConfigure (>DnsConfigData);
+FreePool 

Re: [edk2] [PATCH v2] NetworkPkg: Replace ASSERT with error handling in DnsDxe

2016-06-20 Thread Ye, Ting
Reviewed-by: Ye Ting  

-Original Message-
From: Wu, Jiaxin 
Sent: Monday, June 20, 2016 9:46 AM
To: edk2-devel@lists.01.org
Cc: Ye, Ting ; Fu, Siyuan ; Zhang, Lubo 

Subject: [PATCH v2] NetworkPkg: Replace ASSERT with error handling in DnsDxe

v2:
*Use goto to simplify code logic.

This patch is used to replace ASSERT with error handling in DnsDxe driver.

Cc: Ye Ting 
Cc: Fu Siyuan 
Cc: Zhang Lubo 
Contributed-under: TianoCore Contribution Agreement 1.0
Signed-off-by: Jiaxin Wu 
---
 NetworkPkg/DnsDxe/DnsProtocol.c | 56 -
 1 file changed, 38 insertions(+), 18 deletions(-)

diff --git a/NetworkPkg/DnsDxe/DnsProtocol.c b/NetworkPkg/DnsDxe/DnsProtocol.c 
index e9101d6..64fca6a 100644
--- a/NetworkPkg/DnsDxe/DnsProtocol.c
+++ b/NetworkPkg/DnsDxe/DnsProtocol.c
@@ -86,23 +86,22 @@ Dns4GetModeData (
 
   OldTpl = gBS->RaiseTPL (TPL_CALLBACK);
 
   Instance = DNS_INSTANCE_FROM_THIS_PROTOCOL4 (This);
   if (Instance->State == DNS_STATE_UNCONFIGED) {
-gBS->RestoreTPL (OldTpl);
-return  EFI_NOT_STARTED;
+Status = EFI_NOT_STARTED;
+goto ON_EXIT;
   }
   
   ZeroMem (DnsModeData, sizeof (EFI_DNS4_MODE_DATA));
 
   //
   // Get the current configuration data of this instance. 
   //
   Status = Dns4CopyConfigure (>DnsConfigData, 
>Dns4CfgData);
   if (EFI_ERROR (Status)) {
-gBS->RestoreTPL (OldTpl);
-return Status;
+goto ON_EXIT;
   }
 
   //
   // Get the DnsServerCount and DnsServerList
   //
@@ -110,11 +109,16 @@ Dns4GetModeData (
   NET_LIST_FOR_EACH_SAFE (Entry, Next, >Dns4ServerList) {
 Index++;
   }
   DnsModeData->DnsServerCount = (UINT32) Index;
   ServerList = AllocatePool (sizeof (EFI_IPv4_ADDRESS) * 
DnsModeData->DnsServerCount);
-  ASSERT (ServerList != NULL);
+  if (ServerList == NULL) {
+Status = EFI_OUT_OF_RESOURCES;
+Dns4CleanConfigure (>DnsConfigData);
+goto ON_EXIT;
+  }
+  
   Index = 0;
   NET_LIST_FOR_EACH_SAFE (Entry, Next, >Dns4ServerList) {
 ServerItem = NET_LIST_USER_STRUCT (Entry, DNS4_SERVER_IP, AllServerLink);
 CopyMem (ServerList + Index, >Dns4ServerIp, sizeof 
(EFI_IPv4_ADDRESS));
 Index++;
@@ -128,22 +132,28 @@ Dns4GetModeData (
   NET_LIST_FOR_EACH_SAFE (Entry, Next, >Dns4CacheList) {
 Index++;
   }
   DnsModeData->DnsCacheCount = (UINT32) Index;
   CacheList = AllocatePool (sizeof (EFI_DNS4_CACHE_ENTRY) * 
DnsModeData->DnsCacheCount);
-  ASSERT (CacheList != NULL);
+  if (CacheList == NULL) {
+Status = EFI_OUT_OF_RESOURCES;
+Dns4CleanConfigure (>DnsConfigData);
+FreePool (ServerList);
+goto ON_EXIT;
+  }
+  
   Index =0;
   NET_LIST_FOR_EACH_SAFE (Entry, Next, >Dns4CacheList) {
 CacheItem = NET_LIST_USER_STRUCT (Entry, DNS4_CACHE, AllCacheLink);
 CopyMem (CacheList + Index, >DnsCache, sizeof 
(EFI_DNS4_CACHE_ENTRY));
 Index++;
   }
   DnsModeData->DnsCacheList = CacheList;
 
+ON_EXIT:
   gBS->RestoreTPL (OldTpl);
-
-  return EFI_SUCCESS;
+  return Status;
 }
 
 /**
   Configure this DNS instance.
 
@@ -907,23 +917,22 @@ Dns6GetModeData (
 
   OldTpl = gBS->RaiseTPL (TPL_CALLBACK);
 
   Instance  = DNS_INSTANCE_FROM_THIS_PROTOCOL6 (This);
   if (Instance->State == DNS_STATE_UNCONFIGED) {
-gBS->RestoreTPL (OldTpl);
-return  EFI_NOT_STARTED;
+Status =  EFI_NOT_STARTED;
+goto ON_EXIT;
   }
 
   ZeroMem (DnsModeData, sizeof (EFI_DNS6_MODE_DATA));
 
   //
   // Get the current configuration data of this instance. 
   //
-  Status = Dns6CopyConfigure(>DnsConfigData, 
>Dns6CfgData);
+  Status = Dns6CopyConfigure (>DnsConfigData, 
+ >Dns6CfgData);
   if (EFI_ERROR (Status)) {
-gBS->RestoreTPL (OldTpl);
-return Status;
+goto ON_EXIT;
   }
   
   //
   // Get the DnsServerCount and DnsServerList
   //
@@ -931,11 +940,16 @@ Dns6GetModeData (
   NET_LIST_FOR_EACH_SAFE (Entry, Next, >Dns6ServerList) {
 Index++;
   }
   DnsModeData->DnsServerCount = (UINT32) Index;
   ServerList = AllocatePool (sizeof(EFI_IPv6_ADDRESS) * 
DnsModeData->DnsServerCount);
-  ASSERT (ServerList != NULL);
+  if (ServerList == NULL) {
+Status = EFI_OUT_OF_RESOURCES;
+Dns6CleanConfigure (>DnsConfigData);
+goto ON_EXIT;
+  }
+  
   Index = 0;
   NET_LIST_FOR_EACH_SAFE (Entry, Next, >Dns6ServerList) {
 ServerItem = NET_LIST_USER_STRUCT (Entry, DNS6_SERVER_IP, AllServerLink);
 CopyMem (ServerList + Index, >Dns6ServerIp, sizeof 
(EFI_IPv6_ADDRESS));
 Index++;
@@ -949,22 +963,28 @@ Dns6GetModeData (
   NET_LIST_FOR_EACH_SAFE (Entry, Next, >Dns6CacheList) {
 Index++;
   }
   DnsModeData->DnsCacheCount = (UINT32) Index;
   CacheList = AllocatePool (sizeof(EFI_DNS6_CACHE_ENTRY) * 
DnsModeData->DnsCacheCount);
-  ASSERT (CacheList != NULL);
+  if (CacheList == NULL) {
+Status = EFI_OUT_OF_RESOURCES;
+Dns6CleanConfigure (>DnsConfigData);
+FreePool 

Re: [edk2] [PATCH v2] NetworkPkg: Replace ASSERT with error handling in DnsDxe

2016-06-19 Thread Fu, Siyuan
Reviewed-by: Fu Siyuan 



> -Original Message-
> From: Wu, Jiaxin
> Sent: Monday, June 20, 2016 9:46 AM
> To: edk2-devel@lists.01.org
> Cc: Ye, Ting ; Fu, Siyuan ; Zhang,
> Lubo 
> Subject: [PATCH v2] NetworkPkg: Replace ASSERT with error handling in
> DnsDxe
> 
> v2:
> *Use goto to simplify code logic.
> 
> This patch is used to replace ASSERT with error handling in
> DnsDxe driver.
> 
> Cc: Ye Ting 
> Cc: Fu Siyuan 
> Cc: Zhang Lubo 
> Contributed-under: TianoCore Contribution Agreement 1.0
> Signed-off-by: Jiaxin Wu 
> ---
>  NetworkPkg/DnsDxe/DnsProtocol.c | 56 ---
> --
>  1 file changed, 38 insertions(+), 18 deletions(-)
> 
> diff --git a/NetworkPkg/DnsDxe/DnsProtocol.c
> b/NetworkPkg/DnsDxe/DnsProtocol.c
> index e9101d6..64fca6a 100644
> --- a/NetworkPkg/DnsDxe/DnsProtocol.c
> +++ b/NetworkPkg/DnsDxe/DnsProtocol.c
> @@ -86,23 +86,22 @@ Dns4GetModeData (
> 
>OldTpl = gBS->RaiseTPL (TPL_CALLBACK);
> 
>Instance = DNS_INSTANCE_FROM_THIS_PROTOCOL4 (This);
>if (Instance->State == DNS_STATE_UNCONFIGED) {
> -gBS->RestoreTPL (OldTpl);
> -return  EFI_NOT_STARTED;
> +Status = EFI_NOT_STARTED;
> +goto ON_EXIT;
>}
> 
>ZeroMem (DnsModeData, sizeof (EFI_DNS4_MODE_DATA));
> 
>//
>// Get the current configuration data of this instance.
>//
>Status = Dns4CopyConfigure (>DnsConfigData, 
> >Dns4CfgData);
>if (EFI_ERROR (Status)) {
> -gBS->RestoreTPL (OldTpl);
> -return Status;
> +goto ON_EXIT;
>}
> 
>//
>// Get the DnsServerCount and DnsServerList
>//
> @@ -110,11 +109,16 @@ Dns4GetModeData (
>NET_LIST_FOR_EACH_SAFE (Entry, Next, >Dns4ServerList) {
>  Index++;
>}
>DnsModeData->DnsServerCount = (UINT32) Index;
>ServerList = AllocatePool (sizeof (EFI_IPv4_ADDRESS) * DnsModeData-
> >DnsServerCount);
> -  ASSERT (ServerList != NULL);
> +  if (ServerList == NULL) {
> +Status = EFI_OUT_OF_RESOURCES;
> +Dns4CleanConfigure (>DnsConfigData);
> +goto ON_EXIT;
> +  }
> +
>Index = 0;
>NET_LIST_FOR_EACH_SAFE (Entry, Next, >Dns4ServerList) {
>  ServerItem = NET_LIST_USER_STRUCT (Entry, DNS4_SERVER_IP,
> AllServerLink);
>  CopyMem (ServerList + Index, >Dns4ServerIp, sizeof
> (EFI_IPv4_ADDRESS));
>  Index++;
> @@ -128,22 +132,28 @@ Dns4GetModeData (
>NET_LIST_FOR_EACH_SAFE (Entry, Next, >Dns4CacheList) {
>  Index++;
>}
>DnsModeData->DnsCacheCount = (UINT32) Index;
>CacheList = AllocatePool (sizeof (EFI_DNS4_CACHE_ENTRY) *
> DnsModeData->DnsCacheCount);
> -  ASSERT (CacheList != NULL);
> +  if (CacheList == NULL) {
> +Status = EFI_OUT_OF_RESOURCES;
> +Dns4CleanConfigure (>DnsConfigData);
> +FreePool (ServerList);
> +goto ON_EXIT;
> +  }
> +
>Index =0;
>NET_LIST_FOR_EACH_SAFE (Entry, Next, >Dns4CacheList) {
>  CacheItem = NET_LIST_USER_STRUCT (Entry, DNS4_CACHE, AllCacheLink);
>  CopyMem (CacheList + Index, >DnsCache, sizeof
> (EFI_DNS4_CACHE_ENTRY));
>  Index++;
>}
>DnsModeData->DnsCacheList = CacheList;
> 
> +ON_EXIT:
>gBS->RestoreTPL (OldTpl);
> -
> -  return EFI_SUCCESS;
> +  return Status;
>  }
> 
>  /**
>Configure this DNS instance.
> 
> @@ -907,23 +917,22 @@ Dns6GetModeData (
> 
>OldTpl = gBS->RaiseTPL (TPL_CALLBACK);
> 
>Instance  = DNS_INSTANCE_FROM_THIS_PROTOCOL6 (This);
>if (Instance->State == DNS_STATE_UNCONFIGED) {
> -gBS->RestoreTPL (OldTpl);
> -return  EFI_NOT_STARTED;
> +Status =  EFI_NOT_STARTED;
> +goto ON_EXIT;
>}
> 
>ZeroMem (DnsModeData, sizeof (EFI_DNS6_MODE_DATA));
> 
>//
>// Get the current configuration data of this instance.
>//
> -  Status = Dns6CopyConfigure(>DnsConfigData, 
> >Dns6CfgData);
> +  Status = Dns6CopyConfigure (>DnsConfigData, 
> >Dns6CfgData);
>if (EFI_ERROR (Status)) {
> -gBS->RestoreTPL (OldTpl);
> -return Status;
> +goto ON_EXIT;
>}
> 
>//
>// Get the DnsServerCount and DnsServerList
>//
> @@ -931,11 +940,16 @@ Dns6GetModeData (
>NET_LIST_FOR_EACH_SAFE (Entry, Next, >Dns6ServerList) {
>  Index++;
>}
>DnsModeData->DnsServerCount = (UINT32) Index;
>ServerList = AllocatePool (sizeof(EFI_IPv6_ADDRESS) * DnsModeData-
> >DnsServerCount);
> -  ASSERT (ServerList != NULL);
> +  if (ServerList == NULL) {
> +Status = EFI_OUT_OF_RESOURCES;
> +Dns6CleanConfigure (>DnsConfigData);
> +goto ON_EXIT;
> +  }
> +
>Index = 0;
>NET_LIST_FOR_EACH_SAFE (Entry, Next, >Dns6ServerList) {
>  ServerItem = NET_LIST_USER_STRUCT (Entry, DNS6_SERVER_IP,
> AllServerLink);
>  CopyMem (ServerList + Index, >Dns6ServerIp, sizeof
> (EFI_IPv6_ADDRESS));
>  Index++;
> @@ -949,22 +963,28 @@ Dns6GetModeData (
>NET_LIST_FOR_EACH_SAFE (Entry, Next, >Dns6CacheList) {
>  Index++;
>}
> 

[edk2] [PATCH v2] NetworkPkg: Replace ASSERT with error handling in DnsDxe

2016-06-19 Thread Jiaxin Wu
v2:
*Use goto to simplify code logic.

This patch is used to replace ASSERT with error handling in
DnsDxe driver.

Cc: Ye Ting 
Cc: Fu Siyuan 
Cc: Zhang Lubo 
Contributed-under: TianoCore Contribution Agreement 1.0
Signed-off-by: Jiaxin Wu 
---
 NetworkPkg/DnsDxe/DnsProtocol.c | 56 -
 1 file changed, 38 insertions(+), 18 deletions(-)

diff --git a/NetworkPkg/DnsDxe/DnsProtocol.c b/NetworkPkg/DnsDxe/DnsProtocol.c
index e9101d6..64fca6a 100644
--- a/NetworkPkg/DnsDxe/DnsProtocol.c
+++ b/NetworkPkg/DnsDxe/DnsProtocol.c
@@ -86,23 +86,22 @@ Dns4GetModeData (
 
   OldTpl = gBS->RaiseTPL (TPL_CALLBACK);
 
   Instance = DNS_INSTANCE_FROM_THIS_PROTOCOL4 (This);
   if (Instance->State == DNS_STATE_UNCONFIGED) {
-gBS->RestoreTPL (OldTpl);
-return  EFI_NOT_STARTED;
+Status = EFI_NOT_STARTED;
+goto ON_EXIT;
   }
   
   ZeroMem (DnsModeData, sizeof (EFI_DNS4_MODE_DATA));
 
   //
   // Get the current configuration data of this instance. 
   //
   Status = Dns4CopyConfigure (>DnsConfigData, 
>Dns4CfgData);
   if (EFI_ERROR (Status)) {
-gBS->RestoreTPL (OldTpl);
-return Status;
+goto ON_EXIT;
   }
 
   //
   // Get the DnsServerCount and DnsServerList
   //
@@ -110,11 +109,16 @@ Dns4GetModeData (
   NET_LIST_FOR_EACH_SAFE (Entry, Next, >Dns4ServerList) {
 Index++;
   }
   DnsModeData->DnsServerCount = (UINT32) Index;
   ServerList = AllocatePool (sizeof (EFI_IPv4_ADDRESS) * 
DnsModeData->DnsServerCount);
-  ASSERT (ServerList != NULL);
+  if (ServerList == NULL) {
+Status = EFI_OUT_OF_RESOURCES;
+Dns4CleanConfigure (>DnsConfigData);
+goto ON_EXIT;
+  }
+  
   Index = 0;
   NET_LIST_FOR_EACH_SAFE (Entry, Next, >Dns4ServerList) {
 ServerItem = NET_LIST_USER_STRUCT (Entry, DNS4_SERVER_IP, AllServerLink);
 CopyMem (ServerList + Index, >Dns4ServerIp, sizeof 
(EFI_IPv4_ADDRESS));
 Index++;
@@ -128,22 +132,28 @@ Dns4GetModeData (
   NET_LIST_FOR_EACH_SAFE (Entry, Next, >Dns4CacheList) {
 Index++;
   }
   DnsModeData->DnsCacheCount = (UINT32) Index;
   CacheList = AllocatePool (sizeof (EFI_DNS4_CACHE_ENTRY) * 
DnsModeData->DnsCacheCount);
-  ASSERT (CacheList != NULL);
+  if (CacheList == NULL) {
+Status = EFI_OUT_OF_RESOURCES;
+Dns4CleanConfigure (>DnsConfigData);
+FreePool (ServerList);
+goto ON_EXIT;
+  }
+  
   Index =0;
   NET_LIST_FOR_EACH_SAFE (Entry, Next, >Dns4CacheList) {
 CacheItem = NET_LIST_USER_STRUCT (Entry, DNS4_CACHE, AllCacheLink);
 CopyMem (CacheList + Index, >DnsCache, sizeof 
(EFI_DNS4_CACHE_ENTRY));
 Index++;
   }
   DnsModeData->DnsCacheList = CacheList;
 
+ON_EXIT:
   gBS->RestoreTPL (OldTpl);
-  
-  return EFI_SUCCESS;
+  return Status;
 }
 
 /**
   Configure this DNS instance.
 
@@ -907,23 +917,22 @@ Dns6GetModeData (
 
   OldTpl = gBS->RaiseTPL (TPL_CALLBACK);
 
   Instance  = DNS_INSTANCE_FROM_THIS_PROTOCOL6 (This);
   if (Instance->State == DNS_STATE_UNCONFIGED) {
-gBS->RestoreTPL (OldTpl);
-return  EFI_NOT_STARTED;
+Status =  EFI_NOT_STARTED;
+goto ON_EXIT;
   }
 
   ZeroMem (DnsModeData, sizeof (EFI_DNS6_MODE_DATA));
 
   //
   // Get the current configuration data of this instance. 
   //
-  Status = Dns6CopyConfigure(>DnsConfigData, 
>Dns6CfgData);
+  Status = Dns6CopyConfigure (>DnsConfigData, 
>Dns6CfgData);
   if (EFI_ERROR (Status)) {
-gBS->RestoreTPL (OldTpl);
-return Status;
+goto ON_EXIT;
   }
   
   //
   // Get the DnsServerCount and DnsServerList
   //
@@ -931,11 +940,16 @@ Dns6GetModeData (
   NET_LIST_FOR_EACH_SAFE (Entry, Next, >Dns6ServerList) {
 Index++;
   }
   DnsModeData->DnsServerCount = (UINT32) Index;
   ServerList = AllocatePool (sizeof(EFI_IPv6_ADDRESS) * 
DnsModeData->DnsServerCount);
-  ASSERT (ServerList != NULL);
+  if (ServerList == NULL) {
+Status = EFI_OUT_OF_RESOURCES;
+Dns6CleanConfigure (>DnsConfigData);
+goto ON_EXIT;
+  }
+  
   Index = 0;
   NET_LIST_FOR_EACH_SAFE (Entry, Next, >Dns6ServerList) {
 ServerItem = NET_LIST_USER_STRUCT (Entry, DNS6_SERVER_IP, AllServerLink);
 CopyMem (ServerList + Index, >Dns6ServerIp, sizeof 
(EFI_IPv6_ADDRESS));
 Index++;
@@ -949,22 +963,28 @@ Dns6GetModeData (
   NET_LIST_FOR_EACH_SAFE (Entry, Next, >Dns6CacheList) {
 Index++;
   }
   DnsModeData->DnsCacheCount = (UINT32) Index;
   CacheList = AllocatePool (sizeof(EFI_DNS6_CACHE_ENTRY) * 
DnsModeData->DnsCacheCount);
-  ASSERT (CacheList != NULL);
+  if (CacheList == NULL) {
+Status = EFI_OUT_OF_RESOURCES;
+Dns6CleanConfigure (>DnsConfigData);
+FreePool (ServerList);
+goto ON_EXIT;
+  }
+  
   Index =0;
   NET_LIST_FOR_EACH_SAFE (Entry, Next, >Dns6CacheList) {
 CacheItem = NET_LIST_USER_STRUCT (Entry, DNS6_CACHE, AllCacheLink);
 CopyMem (CacheList + Index, >DnsCache, sizeof 
(EFI_DNS6_CACHE_ENTRY));
 Index++;
   }
   DnsModeData->DnsCacheList = CacheList;
-
-  gBS->RestoreTPL