NetworkPkg: Remove TokenEntry from Token list before freed

[mirror_edk2.git] / NetworkPkg / DnsDxe / DnsImpl.c

diff --git a/NetworkPkg/DnsDxe/DnsImpl.c b/NetworkPkg/DnsDxe/DnsImpl.c
index 8725670a8e8a4a51c098dcd7b1e61f067039c021..4f10e17cf30ffca63646b355da0857c599dc176b 100644 (file)
--- a/NetworkPkg/DnsDxe/DnsImpl.c
+++ b/NetworkPkg/DnsDxe/DnsImpl.c
@@ -1,7 +1,7 @@
 /** @file\r
 DnsDxe support functions implementation.\r
   \r
-Copyright (c) 2015, Intel Corporation. All rights reserved.<BR>\r
+Copyright (c) 2016, Intel Corporation. All rights reserved.<BR>\r
 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\r
@@ -618,11 +618,35 @@ Dns6GetMapping (
   while (!EFI_ERROR (gBS->CheckEvent (Service->TimerToGetMap))) {\r
     Udp->Poll (Udp);\r
 \r
-    if (!EFI_ERROR (Udp->GetModeData (Udp, NULL, &Ip6Mode, NULL, NULL)) &&\r
-        Ip6Mode.IsConfigured) {\r
+    if (!EFI_ERROR (Udp->GetModeData (Udp, NULL, &Ip6Mode, NULL, NULL))) {\r
+      if (Ip6Mode.AddressList != NULL) {\r
+        FreePool (Ip6Mode.AddressList);\r
+      }\r
 \r
-      Udp->Configure (Udp, NULL);\r
-      return (BOOLEAN) (Udp->Configure (Udp, UdpCfgData) == EFI_SUCCESS);\r
+      if (Ip6Mode.GroupTable != NULL) {\r
+        FreePool (Ip6Mode.GroupTable);\r
+      }\r
+\r
+      if (Ip6Mode.RouteTable != NULL) {\r
+        FreePool (Ip6Mode.RouteTable);\r
+      }\r
+\r
+      if (Ip6Mode.NeighborCache != NULL) {\r
+        FreePool (Ip6Mode.NeighborCache);\r
+      }\r
+\r
+      if (Ip6Mode.PrefixTable != NULL) {\r
+        FreePool (Ip6Mode.PrefixTable);\r
+      }\r
+\r
+      if (Ip6Mode.IcmpTypeList != NULL) {\r
+        FreePool (Ip6Mode.IcmpTypeList);\r
+      }\r
+\r
+      if (Ip6Mode.IsConfigured) {\r
+        Udp->Configure (Udp, NULL);\r
+        return (BOOLEAN) (Udp->Configure (Udp, UdpCfgData) == EFI_SUCCESS);\r
+      }\r
     }\r
   }\r
 \r
@@ -1013,61 +1037,6 @@ AddDns6ServerIp (
   return EFI_SUCCESS;\r
 }\r
 \r
-/**\r
-  Fill QName for IP querying. QName is a domain name represented as \r
-  a sequence of labels, where each label consists of a length octet \r
-  followed by that number of octets. The domain name terminates with \r
-  the zero length octet for the null label of the root. Caller should \r
-  take responsibility to the buffer in QName.\r
-\r
-  @param  HostName          Queried HostName    \r
-\r
-  @retval NULL      Failed to fill QName.\r
-  @return           QName filled successfully.\r
-  \r
-**/ \r
-UINT8 *\r
-EFIAPI\r
-DnsFillinQNameForQueryIp (\r
-  IN  CHAR16              *HostName\r
-  )\r
-{\r
-  CHAR8                 *QueryName;\r
-  CHAR8                 *Header;\r
-  CHAR8                 *Tail;\r
-  UINTN                 Len;\r
-  UINTN                 Index;\r
-\r
-  QueryName  = NULL;\r
-  Header     = NULL;\r
-  Tail       = NULL;\r
-\r
-  QueryName = AllocateZeroPool (DNS_DEFAULT_BLKSIZE);\r
-  if (QueryName == NULL) {\r
-    return NULL;\r
-  }\r
-  \r
-  Header = QueryName;\r
-  Tail = Header + 1;\r
-  Len = 0;\r
-  for (Index = 0; HostName[Index] != 0; Index++) {\r
-    *Tail = (CHAR8) HostName[Index];\r
-    if (*Tail == '.') {\r
-      *Header = (CHAR8) Len;\r
-      Header = Tail;\r
-      Tail ++;\r
-      Len = 0;\r
-    } else {\r
-      Tail++;\r
-      Len++;\r
-    }\r
-  }\r
-  *Header = (CHAR8) Len;\r
-  *Tail = 0;\r
-\r
-  return QueryName;\r
-}\r
-\r
 /**\r
   Find out whether the response is valid or invalid.\r
 \r
@@ -1360,8 +1329,6 @@ ParseDnsResponse (
     // Check whether it's the GeneralLookUp querying.\r
     //\r
     if (Instance->Service->IpVersion == IP_VERSION_4 && Dns4TokenEntry->GeneralLookUp) {\r
-      ASSERT (Dns4TokenEntry != NULL);\r
-      \r
       Dns4RR = Dns4TokenEntry->Token->RspData.GLookupData->RRList;\r
       AnswerData = (UINT8 *) AnswerSection + sizeof (*AnswerSection);\r
 \r
@@ -1387,8 +1354,6 @@ ParseDnsResponse (
       \r
       RRCount ++;\r
     } else if (Instance->Service->IpVersion == IP_VERSION_6 && Dns6TokenEntry->GeneralLookUp) {\r
-      ASSERT (Dns6TokenEntry != NULL);\r
-      \r
       Dns6RR = Dns6TokenEntry->Token->RspData.GLookupData->RRList;\r
       AnswerData = (UINT8 *) AnswerSection + sizeof (*AnswerSection);\r
 \r
@@ -1650,6 +1615,10 @@ DnsOnPacketReceived (
   }\r
 \r
   ASSERT (Packet != NULL);\r
+\r
+  if (Packet->TotalSize <= sizeof (DNS_HEADER)) {\r
+    goto ON_EXIT;\r
+  }\r
   \r
   RcvString = NetbufGetByte (Packet, 0, NULL);\r
   ASSERT (RcvString != NULL);\r
@@ -1659,15 +1628,15 @@ DnsOnPacketReceived (
   //\r
   ParseDnsResponse (Instance, RcvString, &Completed);\r
 \r
-  ON_EXIT:\r
+ON_EXIT:\r
 \r
-    if (Packet != NULL) {\r
-      NetbufFree (Packet);\r
-    }\r
+  if (Packet != NULL) {\r
+    NetbufFree (Packet);\r
+  }\r
 \r
-    if (!Completed) {\r
-      UdpIoRecvDatagram (Instance->UdpIo, DnsOnPacketReceived, Instance, 0);\r
-    }\r
+  if (!Completed) {\r
+    UdpIoRecvDatagram (Instance->UdpIo, DnsOnPacketReceived, Instance, 0);\r
+  }\r
 }\r
 \r
 /**\r
@@ -1786,8 +1755,12 @@ ConstructDNSQuery (
   NET_FRAGMENT        Frag;\r
   DNS_HEADER          *DnsHeader;\r
   DNS_QUERY_SECTION   *DnsQuery;\r
-  \r
-  Frag.Bulk = AllocatePool (DNS_DEFAULT_BLKSIZE * sizeof (UINT8));\r
+\r
+  //\r
+  // Messages carried by UDP are restricted to 512 bytes (not counting the IP\r
+  // or UDP headers).\r
+  //\r
+  Frag.Bulk = AllocatePool (DNS_MAX_MESSAGE_SIZE * sizeof (UINT8));\r
   if (Frag.Bulk == NULL) {\r
     return EFI_OUT_OF_RESOURCES;\r
   }\r