NetworkPkg/DnsDxe: Fix the potential memory leak issue.

Cc: Ye Ting <ting.ye@intel.com>
Cc: Fu Siyuan <siyuan.fu@intel.com>
Cc: Wang Fan <fan.wang@intel.com>
Contributed-under: TianoCore Contribution Agreement 1.0
Signed-off-by: Wu Jiaxin <jiaxin.wu@intel.com>
Reviewed-by: Fu Siyuan <siyuan.fu@intel.com>

1 files changed, 121 insertions, 10 deletions

diff --git a/NetworkPkg/DnsDxe/DnsImpl.c b/NetworkPkg/DnsDxe/DnsImpl.c
index 7c236a0015..7057bfbbec 100644
--- a/NetworkPkg/DnsDxe/DnsImpl.c
+++ b/NetworkPkg/DnsDxe/DnsImpl.c
@@ -792,6 +792,10 @@ UpdateDns4Cache (
         // Delete matching DNS Cache entry

         //

         RemoveEntryList (&Item->AllCacheLink);

+

+        FreePool (Item->DnsCache.HostName);

+        FreePool (Item->DnsCache.IpAddress);

+        FreePool (Item);

         

         return EFI_SUCCESS;

       } else if (Override) {

@@ -819,13 +823,16 @@ UpdateDns4Cache (
    

   NewDnsCache->DnsCache.HostName = AllocatePool (StrSize (DnsCacheEntry.HostName));

   if (NewDnsCache->DnsCache.HostName == NULL) { 

+    FreePool (NewDnsCache);

     return EFI_OUT_OF_RESOURCES;

   }

   

   CopyMem (NewDnsCache->DnsCache.HostName, DnsCacheEntry.HostName, StrSize (DnsCacheEntry.HostName));

 

   NewDnsCache->DnsCache.IpAddress = AllocatePool (sizeof (EFI_IPv4_ADDRESS));

-  if (NewDnsCache->DnsCache.IpAddress == NULL) { 

+  if (NewDnsCache->DnsCache.IpAddress == NULL) {

+    FreePool (NewDnsCache->DnsCache.HostName);

+    FreePool (NewDnsCache);

     return EFI_OUT_OF_RESOURCES;

   }

 

@@ -884,6 +891,10 @@ UpdateDns6Cache (
         //

         RemoveEntryList (&Item->AllCacheLink);

         

+        FreePool (Item->DnsCache.HostName);

+        FreePool (Item->DnsCache.IpAddress);

+        FreePool (Item);

+        

         return EFI_SUCCESS;

       } else if (Override) {

         //

@@ -910,13 +921,16 @@ UpdateDns6Cache (
    

   NewDnsCache->DnsCache.HostName = AllocatePool (StrSize (DnsCacheEntry.HostName));

   if (NewDnsCache->DnsCache.HostName == NULL) { 

+    FreePool (NewDnsCache);

     return EFI_OUT_OF_RESOURCES;

   }

   

   CopyMem (NewDnsCache->DnsCache.HostName, DnsCacheEntry.HostName, StrSize (DnsCacheEntry.HostName));

 

   NewDnsCache->DnsCache.IpAddress = AllocatePool (sizeof (EFI_IPv6_ADDRESS));

-  if (NewDnsCache->DnsCache.IpAddress == NULL) { 

+  if (NewDnsCache->DnsCache.IpAddress == NULL) {

+    FreePool (NewDnsCache->DnsCache.HostName);

+    FreePool (NewDnsCache);

     return EFI_OUT_OF_RESOURCES;

   }

   

@@ -1256,12 +1270,12 @@ ParseDnsResponse (
       //

       // It's the GeneralLookUp querying.

       //

-      Dns4TokenEntry->Token->RspData.GLookupData = AllocatePool (sizeof (DNS_RESOURCE_RECORD));

+      Dns4TokenEntry->Token->RspData.GLookupData = AllocateZeroPool (sizeof (DNS_RESOURCE_RECORD));

       if (Dns4TokenEntry->Token->RspData.GLookupData == NULL) {

         Status = EFI_OUT_OF_RESOURCES;

         goto ON_EXIT;

       }

-      Dns4TokenEntry->Token->RspData.GLookupData->RRList = AllocatePool (DnsHeader->AnswersNum * sizeof (DNS_RESOURCE_RECORD));

+      Dns4TokenEntry->Token->RspData.GLookupData->RRList = AllocateZeroPool (DnsHeader->AnswersNum * sizeof (DNS_RESOURCE_RECORD));

       if (Dns4TokenEntry->Token->RspData.GLookupData->RRList == NULL) {

         Status = EFI_OUT_OF_RESOURCES;

         goto ON_EXIT;

@@ -1271,12 +1285,12 @@ ParseDnsResponse (
       // It's not the GeneralLookUp querying. Check the Query type.

       //

       if (QuerySection->Type == DNS_TYPE_A) {

-        Dns4TokenEntry->Token->RspData.H2AData = AllocatePool (sizeof (DNS_HOST_TO_ADDR_DATA));

+        Dns4TokenEntry->Token->RspData.H2AData = AllocateZeroPool (sizeof (DNS_HOST_TO_ADDR_DATA));

         if (Dns4TokenEntry->Token->RspData.H2AData == NULL) {

           Status = EFI_OUT_OF_RESOURCES;

           goto ON_EXIT;

         }

-        Dns4TokenEntry->Token->RspData.H2AData->IpList = AllocatePool (DnsHeader->AnswersNum * sizeof (EFI_IPv4_ADDRESS));

+        Dns4TokenEntry->Token->RspData.H2AData->IpList = AllocateZeroPool (DnsHeader->AnswersNum * sizeof (EFI_IPv4_ADDRESS));

         if (Dns4TokenEntry->Token->RspData.H2AData->IpList == NULL) {

           Status = EFI_OUT_OF_RESOURCES;

           goto ON_EXIT;

@@ -1293,12 +1307,12 @@ ParseDnsResponse (
       //

       // It's the GeneralLookUp querying.

       //

-      Dns6TokenEntry->Token->RspData.GLookupData = AllocatePool (sizeof (DNS_RESOURCE_RECORD));

+      Dns6TokenEntry->Token->RspData.GLookupData = AllocateZeroPool (sizeof (DNS_RESOURCE_RECORD));

       if (Dns6TokenEntry->Token->RspData.GLookupData == NULL) {

         Status = EFI_OUT_OF_RESOURCES;

         goto ON_EXIT;

       }

-      Dns6TokenEntry->Token->RspData.GLookupData->RRList = AllocatePool (DnsHeader->AnswersNum * sizeof (DNS_RESOURCE_RECORD));

+      Dns6TokenEntry->Token->RspData.GLookupData->RRList = AllocateZeroPool (DnsHeader->AnswersNum * sizeof (DNS_RESOURCE_RECORD));

       if (Dns6TokenEntry->Token->RspData.GLookupData->RRList == NULL) {

         Status = EFI_OUT_OF_RESOURCES;

         goto ON_EXIT;

@@ -1308,12 +1322,12 @@ ParseDnsResponse (
       // It's not the GeneralLookUp querying. Check the Query type.

       //

       if (QuerySection->Type == DNS_TYPE_AAAA) {

-        Dns6TokenEntry->Token->RspData.H2AData = AllocatePool (sizeof (DNS6_HOST_TO_ADDR_DATA));

+        Dns6TokenEntry->Token->RspData.H2AData = AllocateZeroPool (sizeof (DNS6_HOST_TO_ADDR_DATA));

         if (Dns6TokenEntry->Token->RspData.H2AData == NULL) {

           Status = EFI_OUT_OF_RESOURCES;

           goto ON_EXIT;

         }

-        Dns6TokenEntry->Token->RspData.H2AData->IpList = AllocatePool (DnsHeader->AnswersNum * sizeof (EFI_IPv6_ADDRESS));

+        Dns6TokenEntry->Token->RspData.H2AData->IpList = AllocateZeroPool (DnsHeader->AnswersNum * sizeof (EFI_IPv6_ADDRESS));

         if (Dns6TokenEntry->Token->RspData.H2AData->IpList == NULL) {

           Status = EFI_OUT_OF_RESOURCES;

           goto ON_EXIT;

@@ -1603,6 +1617,97 @@ ON_COMPLETE:
   }

 

 ON_EXIT:

+  //

+  // Free the allocated buffer if error happen.

+  //

+  if (EFI_ERROR (Status)) {

+    if (Dns4TokenEntry != NULL) {

+      if (Dns4TokenEntry->GeneralLookUp) {

+        if (Dns4TokenEntry->Token->RspData.GLookupData != NULL) {

+          if (Dns4TokenEntry->Token->RspData.GLookupData->RRList != NULL) {

+            while (RRCount != 0) {

+              RRCount --;

+              if (Dns4TokenEntry->Token->RspData.GLookupData->RRList[RRCount].QName != NULL) {

+                FreePool (Dns4TokenEntry->Token->RspData.GLookupData->RRList[RRCount].QName);

+              }

+

+              if (Dns4TokenEntry->Token->RspData.GLookupData->RRList[RRCount].RData != NULL) {

+                FreePool (Dns4TokenEntry->Token->RspData.GLookupData->RRList[RRCount].RData);

+              }

+            }

+            

+            FreePool (Dns4TokenEntry->Token->RspData.GLookupData->RRList);

+          }

+          

+          FreePool (Dns4TokenEntry->Token->RspData.GLookupData);

+        }

+      } else {

+        if (QuerySection->Type == DNS_TYPE_A && Dns4TokenEntry->Token->RspData.H2AData != NULL) {

+          if (Dns4TokenEntry->Token->RspData.H2AData->IpList != NULL) {

+            FreePool (Dns4TokenEntry->Token->RspData.H2AData->IpList);

+          }

+          

+          FreePool (Dns4TokenEntry->Token->RspData.H2AData);

+        }

+      }

+    }

+

+    if (Dns6TokenEntry != NULL) {

+      if (Dns6TokenEntry->GeneralLookUp) {

+        if (Dns6TokenEntry->Token->RspData.GLookupData != NULL) {

+          if (Dns6TokenEntry->Token->RspData.GLookupData->RRList != NULL) {

+            while (RRCount != 0) {

+              RRCount --;

+              if (Dns6TokenEntry->Token->RspData.GLookupData->RRList[RRCount].QName != NULL) {

+                FreePool (Dns6TokenEntry->Token->RspData.GLookupData->RRList[RRCount].QName);

+              }

+

+              if (Dns6TokenEntry->Token->RspData.GLookupData->RRList[RRCount].RData != NULL) {

+                FreePool (Dns6TokenEntry->Token->RspData.GLookupData->RRList[RRCount].RData);

+              }

+            }

+            

+            FreePool (Dns6TokenEntry->Token->RspData.GLookupData->RRList);

+          }

+          

+          FreePool (Dns6TokenEntry->Token->RspData.GLookupData);

+        }

+      } else {

+        if (QuerySection->Type == DNS_TYPE_AAAA && Dns6TokenEntry->Token->RspData.H2AData != NULL) {

+          if (Dns6TokenEntry->Token->RspData.H2AData->IpList != NULL) {

+            FreePool (Dns6TokenEntry->Token->RspData.H2AData->IpList);

+          }

+          

+          FreePool (Dns6TokenEntry->Token->RspData.H2AData);

+        }

+      }

+    }

+

+    if (Dns4CacheEntry != NULL) {

+      if (Dns4CacheEntry->HostName != NULL) {

+        FreePool (Dns4CacheEntry->HostName);

+      }

+

+      if (Dns4CacheEntry->IpAddress != NULL) {

+        FreePool (Dns4CacheEntry->IpAddress);

+      }

+

+      FreePool (Dns4CacheEntry);

+    }

+

+    if (Dns6CacheEntry != NULL) {

+      if (Dns6CacheEntry->HostName != NULL) {

+        FreePool (Dns6CacheEntry->HostName);

+      }

+

+      if (Dns6CacheEntry->IpAddress != NULL) {

+        FreePool (Dns6CacheEntry->IpAddress);

+      }

+

+      FreePool (Dns6CacheEntry);

+    }    

+  }

+  

   gBS->RestoreTPL (OldTpl);

   return Status;

 }

@@ -2053,6 +2158,9 @@ DnsOnTimerUpdate (
     Item4 = NET_LIST_USER_STRUCT (Entry, DNS4_CACHE, AllCacheLink);

     if (Item4->DnsCache.Timeout == 0) {

       RemoveEntryList (&Item4->AllCacheLink);

+      FreePool (Item4->DnsCache.HostName);

+      FreePool (Item4->DnsCache.IpAddress);

+      FreePool (Item4);

       Entry = mDriverData->Dns4CacheList.ForwardLink;

     } else {

       Entry = Entry->ForwardLink;

@@ -2072,6 +2180,9 @@ DnsOnTimerUpdate (
     Item6 = NET_LIST_USER_STRUCT (Entry, DNS6_CACHE, AllCacheLink);

     if (Item6->DnsCache.Timeout == 0) {

       RemoveEntryList (&Item6->AllCacheLink);

+      FreePool (Item6->DnsCache.HostName);

+      FreePool (Item6->DnsCache.IpAddress);

+      FreePool (Item6);

       Entry = mDriverData->Dns6CacheList.ForwardLink;

     } else {

       Entry = Entry->ForwardLink;