[Patch v2] MdeModulePkg: Free NET_BUF data after it is sent out to avoid memory leak

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

V2:
* Since packet has already been referred by DhcpSb->LastPacket, and will be
freed when sending another packet or clean up, there is no need to add an
extra free function in NetbufFromExt.

Cc: Jiaxin Wu <jiaxin.wu@intel.com>
Cc: Ye Ting <ting.ye@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>
---
 MdeModulePkg/Universal/Network/Dhcp4Dxe/Dhcp4Impl.h | 12 ++++++++++++
 MdeModulePkg/Universal/Network/Dhcp4Dxe/Dhcp4Io.c   | 13 +++++--------
 2 files changed, 17 insertions(+), 8 deletions(-)

diff --git a/MdeModulePkg/Universal/Network/Dhcp4Dxe/Dhcp4Impl.h b/MdeModulePkg/Universal/Network/Dhcp4Dxe/Dhcp4Impl.h
index e546a08..57f6d5e 100644
--- a/MdeModulePkg/Universal/Network/Dhcp4Dxe/Dhcp4Impl.h
+++ b/MdeModulePkg/Universal/Network/Dhcp4Dxe/Dhcp4Impl.h
@@ -182,10 +182,22 @@ VOID
 DhcpCleanConfigure (
   IN OUT EFI_DHCP4_CONFIG_DATA  *Config
   );
 
 /**
+  Callback of Dhcp packet. Does nothing.
+
+  @param Arg           The context.
+
+**/
+VOID
+EFIAPI
+DhcpDummyExtFree (
+  IN VOID                   *Arg
+  );
+
+/**
   Set the elapsed time based on the given instance and the pointer to the
   elapsed time option.
 
   @param[in]      Elapsed       The pointer to the position to append.
   @param[in]      Instance      The pointer to the Dhcp4 instance.
diff --git a/MdeModulePkg/Universal/Network/Dhcp4Dxe/Dhcp4Io.c b/MdeModulePkg/Universal/Network/Dhcp4Dxe/Dhcp4Io.c
index 3898223..54a610a 100644
--- a/MdeModulePkg/Universal/Network/Dhcp4Dxe/Dhcp4Io.c
+++ b/MdeModulePkg/Universal/Network/Dhcp4Dxe/Dhcp4Io.c
@@ -1357,17 +1357,16 @@ DhcpSendMessage (
     &DhcpSb->ClientAddressSendOut[0],
     &Packet->Dhcp4.Header.ClientHwAddr[0],
     Packet->Dhcp4.Header.HwAddrLen
     );
 
-
   //
   // Wrap it into a netbuf then send it.
   //
   Frag.Bulk = (UINT8 *) &Packet->Dhcp4.Header;
   Frag.Len  = Packet->Length;
-  Wrap      = NetbufFromExt (&Frag, 1, 0, 0, DhcpReleasePacket, Packet);
+  Wrap      = NetbufFromExt (&Frag, 1, 0, 0, DhcpDummyExtFree, NULL);
 
   if (Wrap == NULL) {
     FreePool (Packet);
     return EFI_OUT_OF_RESOURCES;
   }
@@ -1397,11 +1396,10 @@ DhcpSendMessage (
     EndPoint.LocalAddr.Addr[0]  = DhcpSb->ClientAddr;
     UdpIo                       = DhcpSb->LeaseIoPort;
   }
 
   ASSERT (UdpIo != NULL);
-  NET_GET_REF (Wrap);
   
   Status = UdpIoSendDatagram (
              UdpIo, 
              Wrap, 
              &EndPoint, 
@@ -1409,11 +1407,11 @@ DhcpSendMessage (
              DhcpOnPacketSent, 
              DhcpSb
              );
 
   if (EFI_ERROR (Status)) {
-    NET_PUT_REF (Wrap);
+    NetbufFree (Wrap);
     return EFI_ACCESS_DENIED;
   }
 
   return EFI_SUCCESS;
 }
@@ -1452,16 +1450,16 @@ DhcpRetransmit (
   //
   // Wrap it into a netbuf then send it.
   //
   Frag.Bulk = (UINT8 *) &DhcpSb->LastPacket->Dhcp4.Header;
   Frag.Len  = DhcpSb->LastPacket->Length;
-  Wrap      = NetbufFromExt (&Frag, 1, 0, 0, DhcpReleasePacket, DhcpSb->LastPacket);
+  Wrap      = NetbufFromExt (&Frag, 1, 0, 0, DhcpDummyExtFree, NULL);
 
   if (Wrap == NULL) {
     return EFI_OUT_OF_RESOURCES;
   }
-  
+
   //
   // Broadcast the message, unless we know the server address.
   //
   EndPoint.RemotePort         = DHCP_SERVER_PORT;
   EndPoint.LocalPort          = DHCP_CLIENT_PORT;
@@ -1475,22 +1473,21 @@ DhcpRetransmit (
     UdpIo                       = DhcpSb->LeaseIoPort;
   }
 
   ASSERT (UdpIo != NULL);
 
-  NET_GET_REF (Wrap);
   Status = UdpIoSendDatagram (
              UdpIo,
              Wrap,
              &EndPoint,
              NULL,
              DhcpOnPacketSent,
              DhcpSb
              );
 
   if (EFI_ERROR (Status)) {
-    NET_PUT_REF (Wrap);
+    NetbufFree (Wrap);
     return EFI_ACCESS_DENIED;
   }
 
   return EFI_SUCCESS;
 }
-- 
1.9.5.msysgit.1