NetworkPkg: Remove ASSERT and use error handling in IpSecDxe

This patch is used to refine the code by removing ASSERT and
using error handling in IpSecDxe driver.

Cc: Ye Ting <ting.ye@intel.com>
Cc: Fu Siyuan <siyuan.fu@intel.com>
Cc: Zhang Lubo <lubo.zhang@intel.com>
Cc: Yao Jiewen <jiewen.yao@intel.com>
Contributed-under: TianoCore Contribution Agreement 1.0
Signed-off-by: Jiaxin Wu <jiaxin.wu@intel.com>
Reviewed-by: Ye Ting <ting.ye@intel.com>

8 files changed, 379 insertions, 97 deletions

diff --git a/NetworkPkg/IpSecDxe/Ikev2/Info.c b/NetworkPkg/IpSecDxe/Ikev2/Info.c
index fe75d6c596..23e47ceea8 100644
--- a/NetworkPkg/IpSecDxe/Ikev2/Info.c
+++ b/NetworkPkg/IpSecDxe/Ikev2/Info.c
@@ -2,7 +2,7 @@
   The Implementations for Information Exchange.

 

   (C) Copyright 2015 Hewlett-Packard Development Company, L.P.<BR>

-  Copyright (c) 2010, Intel Corporation. All rights reserved.<BR>

+  Copyright (c) 2010 - 2016, Intel Corporation. All rights reserved.<BR>

 

   This program and the accompanying materials

   are licensed and made available under the terms and conditions of the BSD License

@@ -47,7 +47,9 @@ Ikev2InfoGenerator (
   InfoContext  = NULL;

   IkeSaSession = (IKEV2_SA_SESSION *) SaSession;

   IkePacket    = IkePacketAlloc ();

-  ASSERT (IkePacket != NULL);

+  if (IkePacket == NULL) {

+    return NULL;

+  }

 

   //

   // Fill IkePacket Header.

diff --git a/NetworkPkg/IpSecDxe/Ikev2/Payload.c b/NetworkPkg/IpSecDxe/Ikev2/Payload.c
index b92ae3a599..d5fe1abb55 100644
--- a/NetworkPkg/IpSecDxe/Ikev2/Payload.c
+++ b/NetworkPkg/IpSecDxe/Ikev2/Payload.c
@@ -2,7 +2,7 @@
   The implementation of Payloads Creation.

 

   (C) Copyright 2015 Hewlett-Packard Development Company, L.P.<BR>

-  Copyright (c) 2010 - 2015, Intel Corporation. All rights reserved.<BR>

+  Copyright (c) 2010 - 2016, Intel Corporation. All rights reserved.<BR>

 

   This program and the accompanying materials

   are licensed and made available under the terms and conditions of the BSD License

@@ -52,7 +52,10 @@ Ikev2GenerateSaPayload (
   UINTN         SaDataSize;

 

   SaPayload = IkePayloadAlloc ();

-  ASSERT (SaPayload != NULL);

+  if (SaPayload == NULL) {

+    return NULL;

+  }

+  

   //

   // TODO: Get the Proposal Number and Transform Number from IPsec Config,

   // after the Ipsecconfig Application is support it.

@@ -70,7 +73,10 @@ Ikev2GenerateSaPayload (
   }

 

   SaData = AllocateZeroPool (SaDataSize);

-  ASSERT (SaData != NULL);

+  if (SaData == NULL) {

+    IkePayloadFree (SaPayload);

+    return NULL;

+  }

 

   CopyMem (SaData, SessionSaData, SaDataSize);

   SaData->SaHeader.Header.NextPayload = NextPayload;

@@ -118,14 +124,20 @@ Ikev2GenerateNoncePayload (
   NonceBlock  = NonceBuf;

 

   Nonce       = AllocateZeroPool (Size);

-  ASSERT (Nonce != NULL);

+  if (Nonce == NULL) {

+    return NULL;

+  }

+  

   CopyMem (Nonce + 1, NonceBlock, Size - sizeof (IKEV2_NONCE));

 

   Nonce->Header.NextPayload   = NextPayload;

   Nonce->Header.PayloadLength = (UINT16) Size;

   NoncePayload                = IkePayloadAlloc ();

-

-  ASSERT (NoncePayload != NULL);

+  if (NoncePayload == NULL) {

+    FreePool (Nonce);

+    return NULL;

+  }

+  

   NoncePayload->PayloadType = IKEV2_PAYLOAD_TYPE_NONCE;

   NoncePayload->PayloadBuf  = (UINT8 *) Nonce;

   NoncePayload->PayloadSize = Size;

@@ -180,7 +192,9 @@ Ikev2GenerateKePayload (
   // Allocate buffer for Key Exchange

   //

   Ke = AllocateZeroPool (KeSize);

-  ASSERT (Ke != NULL);

+  if (Ke == NULL) {

+    return NULL;

+  }

 

   Ke->Header.NextPayload    = NextPayload;

   Ke->Header.PayloadLength  = (UINT16) KeSize;

@@ -192,7 +206,10 @@ Ikev2GenerateKePayload (
   // Create IKE_PAYLOAD to point to Key Exchange payload

   //

   KePayload = IkePayloadAlloc ();

-  ASSERT (KePayload != NULL);

+  if (KePayload == NULL) {

+    FreePool (Ke);

+    return NULL;

+  }

 

   KePayload->PayloadType = IKEV2_PAYLOAD_TYPE_KE;

   KePayload->PayloadBuf  = (UINT8 *) Ke;

@@ -241,10 +258,15 @@ Ikev2GenerateIdPayload (
   IdSize    = sizeof (IKEV2_ID) + AddrSize;

 

   Id = (IKEV2_ID *) AllocateZeroPool (IdSize);

-  ASSERT (Id != NULL);

+  if (Id == NULL) {

+    return NULL;

+  }

 

   IdPayload = IkePayloadAlloc ();

-  ASSERT (IdPayload != NULL);

+  if (IdPayload == NULL) {

+    FreePool (Id);

+    return NULL;

+  }

 

   IdPayload->PayloadType  = (UINT8) ((CommonSession->IsInitiator) ? IKEV2_PAYLOAD_TYPE_ID_INIT : IKEV2_PAYLOAD_TYPE_ID_RSP);

   IdPayload->PayloadBuf   = (UINT8 *) Id;

@@ -317,10 +339,15 @@ Ikev2GenerateCertIdPayload (
   IdSize = sizeof (IKEV2_ID) + SubjectSize;

 

   Id = (IKEV2_ID *) AllocateZeroPool (IdSize);

-  ASSERT (Id != NULL);

+  if (Id == NULL) {

+    return NULL;

+  }

 

   IdPayload = IkePayloadAlloc ();

-  ASSERT (IdPayload != NULL);

+  if (IdPayload == NULL) {

+    FreePool (Id);

+    return NULL;

+  }

 

   IdPayload->PayloadType  = (UINT8) ((CommonSession->IsInitiator) ? IKEV2_PAYLOAD_TYPE_ID_INIT : IKEV2_PAYLOAD_TYPE_ID_RSP);

   IdPayload->PayloadBuf   = (UINT8 *) Id;

@@ -398,13 +425,14 @@ Ikev2PskGenerateAuthPayload (
 

   DigestSize = IpSecGetHmacDigestLength ((UINT8)IkeSaSession->SessionCommon.SaParams->Prf);

   Digest     = AllocateZeroPool (DigestSize);

-

   if (Digest == NULL) {

     return NULL;

   }

+  

   if (IdPayload == NULL) {

     return NULL;

   }

+  

   //

   // Calcualte Prf(Seceret, "Key Pad for IKEv2");

   //

@@ -428,7 +456,11 @@ Ikev2PskGenerateAuthPayload (
   // Store the AuthKey into KeyBuf

   //

   KeyBuf = AllocateZeroPool (DigestSize);

-  ASSERT (KeyBuf != NULL);

+  if (KeyBuf == NULL) {

+    Status = EFI_OUT_OF_RESOURCES;

+    goto EXIT;

+  }

+  

   CopyMem (KeyBuf, Digest, DigestSize);

   KeySize = DigestSize;

 

@@ -486,6 +518,11 @@ Ikev2PskGenerateAuthPayload (
   // Copy the result of Prf(SK_Pr, IDi/r) to Fragments[2].

   //

   Fragments[2].Data     = AllocateZeroPool (DigestSize);

+  if (Fragments[2].Data == NULL) {

+    Status = EFI_OUT_OF_RESOURCES;

+    goto EXIT;

+  }

+  

   Fragments[2].DataSize = DigestSize;

   CopyMem (Fragments[2].Data, Digest, DigestSize);

 

@@ -509,11 +546,18 @@ Ikev2PskGenerateAuthPayload (
   // Allocate buffer for Auth Payload

   //

   AuthPayload               = IkePayloadAlloc ();

-  ASSERT (AuthPayload != NULL);

+  if (AuthPayload == NULL) {

+    Status = EFI_OUT_OF_RESOURCES;

+    goto EXIT;

+  }

 

   AuthPayload->PayloadSize  = sizeof (IKEV2_AUTH) + DigestSize;

   PayloadBuf                = (IKEV2_AUTH *) AllocateZeroPool (AuthPayload->PayloadSize);

-  ASSERT (PayloadBuf != NULL);

+  if (PayloadBuf == NULL) {

+    Status = EFI_OUT_OF_RESOURCES;

+    goto EXIT;

+  }

+  

   //

   // Fill in Auth payload.

   //

@@ -649,7 +693,6 @@ Ikev2CertGenerateAuthPayload (
   }

   DigestSize = IpSecGetHmacDigestLength ((UINT8)IkeSaSession->SessionCommon.SaParams->Prf);

   Digest     = AllocateZeroPool (DigestSize);

-

   if (Digest == NULL) {

     return NULL;

   }

@@ -658,8 +701,11 @@ Ikev2CertGenerateAuthPayload (
   // Store the AuthKey into KeyBuf

   //

   KeyBuf  = AllocateZeroPool (DigestSize);

-  ASSERT (KeyBuf != NULL);

-

+  if (KeyBuf == NULL) {

+    Status = EFI_OUT_OF_RESOURCES;

+    goto EXIT;

+  }

+  

   CopyMem (KeyBuf, Digest, DigestSize);

 

   //

@@ -724,6 +770,11 @@ Ikev2CertGenerateAuthPayload (
   // Copy the result of Prf(SK_Pr, IDi/r) to Fragments[2].

   //

   Fragments[2].Data     = AllocateZeroPool (DigestSize);

+  if (Fragments[2].Data == NULL) {

+    Status = EFI_OUT_OF_RESOURCES;

+    goto EXIT;

+  }

+  

   Fragments[2].DataSize = DigestSize;

   CopyMem (Fragments[2].Data, Digest, DigestSize);

 

@@ -766,7 +817,10 @@ Ikev2CertGenerateAuthPayload (
   // Allocate buffer for Auth Payload

   //

   AuthPayload = IkePayloadAlloc ();

-  ASSERT (AuthPayload != NULL);

+  if (AuthPayload == NULL) {

+    Status = EFI_OUT_OF_RESOURCES;

+    goto EXIT;

+  }

 

   if (!IsVerify) {

     AuthPayload->PayloadSize  = sizeof (IKEV2_AUTH) + SigSize;

@@ -775,7 +829,11 @@ Ikev2CertGenerateAuthPayload (
   }

 

   PayloadBuf = (IKEV2_AUTH *) AllocateZeroPool (AuthPayload->PayloadSize);

-  ASSERT (PayloadBuf != NULL);

+  if (PayloadBuf == NULL) {

+    Status = EFI_OUT_OF_RESOURCES;

+    goto EXIT;

+  }

+  

   //

   // Fill in Auth payload.

   //

@@ -879,7 +937,9 @@ Ikev2GenerateTsPayload (
   //

 

   TsPayload    = IkePayloadAlloc();

-  ASSERT (TsPayload != NULL);

+  if (TsPayload == NULL) {

+    return NULL;

+  }

 

   IpVersion    = ChildSa->SessionCommon.UdpService->IpVersion;

   //

@@ -890,7 +950,9 @@ Ikev2GenerateTsPayload (
   SelectorSize  = sizeof (TRAFFIC_SELECTOR) + 2 * AddrSize;

   TsPayloadSize = sizeof (IKEV2_TS) + SelectorSize;

   TsPayloadBuf = AllocateZeroPool (TsPayloadSize);

-  ASSERT (TsPayloadBuf != NULL);

+  if (TsPayloadBuf == NULL) {

+    goto ON_ERROR;

+  }

 

   TsPayload->PayloadBuf = (UINT8 *) TsPayloadBuf;

   TsSelector            = (TRAFFIC_SELECTOR*)(TsPayloadBuf + 1);

@@ -1146,7 +1208,9 @@ Ikev2GenerateNotifyPayload (
   //

   NotifyPayloadLen  = (UINT16) (sizeof (IKEV2_NOTIFY) + NotifyDataSize + SpiSize);

   Notify            = (IKEV2_NOTIFY *) AllocateZeroPool (NotifyPayloadLen);

-  ASSERT (Notify != NULL);

+  if (Notify == NULL) {

+    return NULL;

+  }

 

   //

   // Set Delete Payload's Generic Header

@@ -1177,7 +1241,11 @@ Ikev2GenerateNotifyPayload (
   // Create Payload for and set type as IKEV2_PAYLOAD_TYPE_NOTIFY

   //

   NotifyPayload = IkePayloadAlloc ();

-  ASSERT (NotifyPayload != NULL);

+  if (NotifyPayload == NULL) {

+    FreePool (Notify);

+    return NULL;

+  }

+  

   NotifyPayload->PayloadType  = IKEV2_PAYLOAD_TYPE_NOTIFY;

   NotifyPayload->PayloadBuf   = (UINT8 *) Notify;

   NotifyPayload->PayloadSize  = NotifyPayloadLen;

@@ -1238,7 +1306,9 @@ Ikev2GenerateDeletePayload (
   DelPayloadLen = (UINT16) (sizeof (IKEV2_DELETE) + SpiBufSize);

 

   Del           = AllocateZeroPool (DelPayloadLen);

-  ASSERT (Del != NULL);

+  if (Del == NULL) {

+    return NULL;

+  }

 

   //

   // Set Delete Payload's Generic Header

@@ -1262,7 +1332,11 @@ Ikev2GenerateDeletePayload (
   //

   CopyMem (Del + 1, SpiBuf, SpiBufSize);

   DelPayload = IkePayloadAlloc ();

-  ASSERT (DelPayload != NULL);

+  if (DelPayload == NULL) {

+    FreePool (Del);

+    return NULL;

+  }

+  

   DelPayload->PayloadType = IKEV2_PAYLOAD_TYPE_DELETE;

   DelPayload->PayloadBuf  = (UINT8 *) Del;

   DelPayload->PayloadSize = DelPayloadLen;

@@ -1626,7 +1700,10 @@ Ikev2EncodeSa (
   // Allocate buffer for IKE_SA.

   //

   Sa = AllocateZeroPool (SaSize);

-  ASSERT (Sa != NULL);

+  if (Sa == NULL) {

+    return NULL;

+  }

+  

   CopyMem (Sa, SaData, sizeof (IKEV2_SA));

   Sa->Header.PayloadLength  = (UINT16) sizeof (IKEV2_SA);

   ProposalsSize             = 0;

@@ -1819,7 +1896,11 @@ Ikev2DecodeSa (
                                TotalProposals * sizeof (IKEV2_PROPOSAL_DATA) +

                                TotalTransforms * sizeof (IKEV2_TRANSFORM_DATA)

                                );

-  ASSERT (SaData != NULL);

+  if (SaData == NULL) {

+    Status = EFI_OUT_OF_RESOURCES;

+    goto Exit;

+  }

+  

   CopyMem (SaData, Sa, sizeof (IKEV2_SA));

   SaData->NumProposals        = TotalProposals;

   ProposalData                = (IKEV2_PROPOSAL_DATA *) (SaData + 1);

@@ -1852,7 +1933,11 @@ Ikev2DecodeSa (
       // SpiSize == 4

       //

       Spi = AllocateZeroPool (Proposal->SpiSize);

-      ASSERT (Spi != NULL);

+      if (Spi == NULL) {

+        Status = EFI_OUT_OF_RESOURCES;

+        goto Exit;

+      }

+      

       CopyMem (Spi, (UINT32 *) (Proposal + 1), Proposal->SpiSize);

       *((UINT32*) Spi) = NTOHL (*((UINT32*) Spi));

       ProposalData->Spi = Spi;

@@ -2284,7 +2369,11 @@ Ikev2DecodePacket (
   //

   if (IkePacket->Header->ExchangeType == IKEV2_EXCHANGE_TYPE_INIT) {

     IkeHeader = AllocateZeroPool (sizeof (IKE_HEADER));

-    ASSERT (IkeHeader != NULL);

+    if (IkeHeader == NULL) {

+      Status = EFI_OUT_OF_RESOURCES;

+      goto Exit;

+    }

+    

     CopyMem (IkeHeader, IkePacket->Header, sizeof (IKE_HEADER));

 

     //

@@ -2358,7 +2447,10 @@ Ikev2DecodePacket (
     // Initial IkePayload

     //

     IkePayload = IkePayloadAlloc ();

-    ASSERT (IkePayload != NULL);

+    if (IkePayload == NULL) {

+      Status = EFI_OUT_OF_RESOURCES;

+      goto Exit;

+    }

 

     IkePayload->PayloadType     = PayloadType;

     IkePayload->PayloadBuf      = (UINT8 *) PayloadHdr;

@@ -2483,7 +2575,10 @@ Ikev2EncodePacket (
     if (SessionCommon->IsInitiator) {

       IkeSaSession->InitPacketSize = IkePacket->PayloadTotalSize + sizeof (IKE_HEADER);

       IkeSaSession->InitPacket     = AllocateZeroPool (IkeSaSession->InitPacketSize);

-      ASSERT (IkeSaSession->InitPacket != NULL);

+      if (IkeSaSession->InitPacket == NULL) {

+        return EFI_OUT_OF_RESOURCES;

+      }

+      

       CopyMem (IkeSaSession->InitPacket, IkePacket->Header, sizeof (IKE_HEADER));

       PayloadTotalSize = 0;

       for (Entry = IkePacket->PayloadList.ForwardLink; Entry != &(IkePacket->PayloadList);) {

@@ -2499,7 +2594,10 @@ Ikev2EncodePacket (
     } else {

       IkeSaSession->RespPacketSize = IkePacket->PayloadTotalSize + sizeof(IKE_HEADER);

       IkeSaSession->RespPacket     = AllocateZeroPool (IkeSaSession->RespPacketSize);

-      ASSERT (IkeSaSession->RespPacket != NULL);

+      if (IkeSaSession->RespPacket == NULL) {

+        return EFI_OUT_OF_RESOURCES;

+      }

+      

       CopyMem (IkeSaSession->RespPacket, IkePacket->Header, sizeof (IKE_HEADER));

       PayloadTotalSize = 0;

       for (Entry = IkePacket->PayloadList.ForwardLink; Entry != &(IkePacket->PayloadList);) {

@@ -2596,14 +2694,21 @@ Ikev2DecryptPacket (
   }

 

   CheckSumData = AllocateZeroPool (CheckSumSize);

-  ASSERT (CheckSumData != NULL);

+  if (CheckSumData == NULL) {

+    Status = EFI_OUT_OF_RESOURCES;

+    goto ON_EXIT;

+  }

 

   //

   // Fill in the Integrity buffer

   //

   IntegritySize   = IkePacket->PayloadTotalSize + sizeof (IKE_HEADER);

   IntegrityBuffer = AllocateZeroPool (IntegritySize);

-  ASSERT (IntegrityBuffer != NULL);

+  if (IntegrityBuffer == NULL) {

+    Status = EFI_OUT_OF_RESOURCES;

+    goto ON_EXIT;

+  }

+  

   CopyMem (IntegrityBuffer, IkePacket->Header, sizeof(IKE_HEADER));

   CopyMem (IntegrityBuffer + sizeof (IKE_HEADER), IkePacket->PayloadsBuf, IkePacket->PayloadTotalSize);

 

@@ -2664,7 +2769,10 @@ Ikev2DecryptPacket (
   //

   DecryptedSize = IkePacket->PayloadTotalSize - sizeof (IKEV2_COMMON_PAYLOAD_HEADER) - IvSize - CheckSumSize;

   DecryptedBuf  = AllocateZeroPool (DecryptedSize);

-  ASSERT (DecryptedBuf != NULL);

+  if (DecryptedBuf == NULL) {

+    Status = EFI_OUT_OF_RESOURCES;

+    goto ON_EXIT;

+  }

 

   CopyMem (

     DecryptedBuf,

@@ -2811,8 +2919,11 @@ Ikev2EncryptPacket (
   CryptBlockSizeMask  = (UINT8) (CryptBlockSize - 1);

   EncryptedSize       = (IkePacket->PayloadTotalSize + sizeof (IKEV2_PAD_LEN) + CryptBlockSizeMask) & ~CryptBlockSizeMask;

   EncryptedBuf        = (UINT8 *) AllocateZeroPool (EncryptedSize);

-  ASSERT (EncryptedBuf != NULL);

-

+  if (EncryptedBuf == NULL) {

+    Status = EFI_OUT_OF_RESOURCES;

+    goto ON_EXIT;

+  }

+  

   //

   // Copy all payload into EncryptedIkePayload

   //

@@ -2878,7 +2989,10 @@ Ikev2EncryptPacket (
   //

   EncryptPayloadSize = sizeof(IKEV2_ENCRYPTED) + IvSize + EncryptedSize + CheckSumSize;

   EncryptPayloadBuf  = AllocateZeroPool (EncryptPayloadSize);

-  ASSERT (EncryptPayloadBuf != NULL);

+  if (EncryptPayloadBuf == NULL) {

+    Status = EFI_OUT_OF_RESOURCES;

+    goto ON_EXIT;

+  }

 

   //

   // Fill in Header of  Encrypted Payload

@@ -2965,7 +3079,10 @@ Ikev2EncryptPacket (
   // Create Encrypted Payload and add into IkePacket->PayloadList

   //

   EncryptPayload = IkePayloadAlloc ();

-  ASSERT (EncryptPayload != NULL);

+  if (EncryptPayload == NULL) {

+    Status = EFI_OUT_OF_RESOURCES;

+    goto ON_EXIT;

+  }

 

   //

   // Fill the encrypted payload into the IKE_PAYLOAD structure.

@@ -3211,7 +3328,9 @@ Ikev2SendIkePacket (
   // Transform IkePacke to NetBuf

   //

   IkePacketNetbuf = IkeNetbufFromPacket ((UINT8 *) SessionCommon, IkePacket, IkeType);

-  ASSERT (IkePacketNetbuf != NULL);

+  if (IkePacketNetbuf == NULL) {

+    return EFI_OUT_OF_RESOURCES;

+  }

 

   ZeroMem (&EndPoint, sizeof (UDP_END_POINT));

   EndPoint.RemotePort = IKE_DEFAULT_PORT;

diff --git a/NetworkPkg/IpSecDxe/Ikev2/Sa.c b/NetworkPkg/IpSecDxe/Ikev2/Sa.c
index 9967e1a682..c83d45678b 100644
--- a/NetworkPkg/IpSecDxe/Ikev2/Sa.c
+++ b/NetworkPkg/IpSecDxe/Ikev2/Sa.c
@@ -2,7 +2,7 @@
   The operations for IKEv2 SA.

 

   (C) Copyright 2015 Hewlett-Packard Development Company, L.P.<BR>

-  Copyright (c) 2010 - 2011, Intel Corporation. All rights reserved.<BR>

+  Copyright (c) 2010 - 2016, Intel Corporation. All rights reserved.<BR>

 

   This program and the accompanying materials

   are licensed and made available under the terms and conditions of the BSD License

@@ -102,7 +102,9 @@ Ikev2InitPskGenerator (
   // 1. Allocate IKE packet

   //

   IkePacket = IkePacketAlloc ();

-  ASSERT (IkePacket != NULL);

+  if (IkePacket == NULL) {

+    goto CheckError;

+  }

 

   //

   // 1.a Fill the IkePacket->Hdr

@@ -176,7 +178,9 @@ Ikev2InitPskGenerator (
   if ((IkeSaSession->SessionCommon.IsInitiator) && (IkeSaSession->NCookie == NULL)) {

     IkeSaSession->NiBlkSize = IKE_NONCE_SIZE;

     IkeSaSession->NiBlock   = IkeGenerateNonce (IKE_NONCE_SIZE);

-    ASSERT (IkeSaSession->NiBlock != NULL);

+    if (IkeSaSession->NiBlock == NULL) {

+      goto CheckError;

+    }

   }

 

   if (IkeSaSession->SessionCommon.IsInitiator) {

@@ -298,7 +302,11 @@ Ikev2InitPskParser (
   //

   NonceSize   = NoncePayload->PayloadSize - sizeof (IKEV2_COMMON_PAYLOAD_HEADER);

   NonceBuffer = (UINT8 *) AllocatePool (NonceSize);

-  ASSERT (NonceBuffer != NULL);

+  if (NonceBuffer == NULL) {

+    Status = EFI_OUT_OF_RESOURCES;

+    goto CheckError;

+  }

+  

   CopyMem (

     NonceBuffer,

     NoncePayload->PayloadBuf + sizeof (IKEV2_COMMON_PAYLOAD_HEADER),

@@ -444,7 +452,9 @@ Ikev2AuthPskGenerator (
   // 1. Allocate IKE Packet

   //

   IkePacket= IkePacketAlloc ();

-  ASSERT (IkePacket != NULL);

+  if (IkePacket == NULL) {

+    return NULL;

+  }

 

   //

   // 1.a Fill the IkePacket Header.

@@ -745,7 +755,10 @@ Ikev2AuthPskParser (
     //

     if (ChildSaSession->IkeSaSession->Spd == NULL) {

       ChildSaSession->IkeSaSession->Spd = ChildSaSession->Spd;

-      Ikev2ChildSaSessionSpdSelectorCreate (ChildSaSession);

+      Status = Ikev2ChildSaSessionSpdSelectorCreate (ChildSaSession);

+      if (EFI_ERROR (Status)) {

+        return Status;

+      }

     }

   } else {

     //

@@ -930,7 +943,9 @@ Ikev2AuthCertGenerator (
   // 1. Allocate IKE Packet

   //

   IkePacket= IkePacketAlloc ();

-  ASSERT (IkePacket != NULL);

+  if (IkePacket == NULL) {

+    return NULL;

+  }

 

   //

   // 1.a Fill the IkePacket Header.

@@ -1280,7 +1295,10 @@ Ikev2AuthCertParser (
     //

     if (ChildSaSession->IkeSaSession->Spd == NULL) {

       ChildSaSession->IkeSaSession->Spd = ChildSaSession->Spd;

-      Ikev2ChildSaSessionSpdSelectorCreate (ChildSaSession);

+      Status = Ikev2ChildSaSessionSpdSelectorCreate (ChildSaSession);

+      if (EFI_ERROR (Status)) {

+        goto Exit;

+      }

     }

   } else {

     //

@@ -1360,17 +1378,27 @@ Ikev2GenerateSaDhPublicKey (
   IKEV2_SESSION_KEYS *IkeKeys;

 

   IkeSaSession->IkeKeys = AllocateZeroPool (sizeof (IKEV2_SESSION_KEYS));

-  ASSERT (IkeSaSession->IkeKeys != NULL);

+  if (IkeSaSession->IkeKeys == NULL) {

+    return EFI_OUT_OF_RESOURCES;

+  }

+  

   IkeKeys = IkeSaSession->IkeKeys;

   IkeKeys->DhBuffer = AllocateZeroPool (sizeof (IKEV2_DH_BUFFER));

-  ASSERT (IkeKeys->DhBuffer != NULL);

+  if (IkeKeys->DhBuffer == NULL) {

+    FreePool (IkeSaSession->IkeKeys);

+    return EFI_OUT_OF_RESOURCES;

+  }

 

   //

   // Init DH with the certain DH Group Description.

   //

   IkeKeys->DhBuffer->GxSize   = OakleyModpGroup[(UINT8)IkeSaSession->SessionCommon.PreferDhGroup].Size >> 3;

   IkeKeys->DhBuffer->GxBuffer = AllocateZeroPool (IkeKeys->DhBuffer->GxSize);

-  ASSERT (IkeKeys->DhBuffer->GxBuffer != NULL);

+  if (IkeKeys->DhBuffer->GxBuffer == NULL) {

+    FreePool (IkeKeys->DhBuffer);

+    FreePool (IkeSaSession->IkeKeys);

+    return EFI_OUT_OF_RESOURCES;

+  }

 

   //

   // Get X PublicKey

@@ -1385,6 +1413,13 @@ Ikev2GenerateSaDhPublicKey (
              );

   if (EFI_ERROR (Status)) {

     DEBUG ((DEBUG_ERROR, "Error CPLKeyManGetKeyParam X public key error Status = %r\n", Status));

+    

+    FreePool (IkeKeys->DhBuffer->GxBuffer);

+    

+    FreePool (IkeKeys->DhBuffer);

+    

+    FreePool (IkeSaSession->IkeKeys);

+    

     return Status;

   }

 

@@ -1422,7 +1457,9 @@ Ikev2GenerateSaDhComputeKey (
   PubKeySize          = KePayload->PayloadSize - sizeof (IKEV2_KEY_EXCHANGE);

   DhBuffer->GxySize   = DhBuffer->GxSize;

   DhBuffer->GxyBuffer = AllocateZeroPool (DhBuffer->GxySize);

-  ASSERT (DhBuffer->GxyBuffer != NULL);

+  if (DhBuffer->GxyBuffer == NULL) {

+    return EFI_OUT_OF_RESOURCES;

+  }

 

   //

   // Get GxyBuf

@@ -1436,6 +1473,9 @@ Ikev2GenerateSaDhComputeKey (
              );

   if (EFI_ERROR (Status)) {

     DEBUG ((DEBUG_ERROR, "Error CPLKeyManGetKeyParam Y session key error Status = %r\n", Status));

+

+    FreePool (DhBuffer->GxyBuffer);

+    

     return Status;

   }

 

@@ -1444,7 +1484,12 @@ Ikev2GenerateSaDhComputeKey (
   //

   DhBuffer->GySize   = PubKeySize;

   DhBuffer->GyBuffer = AllocateZeroPool (DhBuffer->GySize);

-  ASSERT (DhBuffer->GyBuffer != NULL);

+  if (DhBuffer->GyBuffer == NULL) {

+    FreePool (DhBuffer->GxyBuffer);

+    

+    return Status;

+  }

+  

   CopyMem (DhBuffer->GyBuffer, PubKey, DhBuffer->GySize);

 

   IPSEC_DUMP_BUF ("DH Public Key (g^y) Dump", DhBuffer->GyBuffer, DhBuffer->GySize);

@@ -1524,7 +1569,10 @@ Ikev2GenerateSaKeys (
   //

   KeyBufferSize = IkeSaSession->NiBlkSize + IkeSaSession->NrBlkSize;

   KeyBuffer     = AllocateZeroPool (KeyBufferSize);

-  ASSERT (KeyBuffer != NULL);

+  if (KeyBuffer == NULL) {

+    Status = EFI_OUT_OF_RESOURCES;

+    goto Exit;

+  }

 

   CopyMem (KeyBuffer, IkeSaSession->NiBlock, IkeSaSession->NiBlkSize);

   CopyMem (KeyBuffer + IkeSaSession->NiBlkSize, IkeSaSession->NrBlock, IkeSaSession->NrBlkSize);

diff --git a/NetworkPkg/IpSecDxe/Ikev2/Utility.c b/NetworkPkg/IpSecDxe/Ikev2/Utility.c
index aa3e440a0e..5b26ba1d02 100644
--- a/NetworkPkg/IpSecDxe/Ikev2/Utility.c
+++ b/NetworkPkg/IpSecDxe/Ikev2/Utility.c
@@ -2,7 +2,7 @@
   The Common operations used by IKE Exchange Process.

 

   (C) Copyright 2015 Hewlett-Packard Development Company, L.P.<BR>

-  Copyright (c) 2010 - 2015, Intel Corporation. All rights reserved.<BR>

+  Copyright (c) 2010 - 2016, Intel Corporation. All rights reserved.<BR>

 

   This program and the accompanying materials

   are licensed and made available under the terms and conditions of the BSD License

@@ -57,7 +57,9 @@ Ikev2SaSessionAlloc (
   IKEV2_SA_SESSION      *IkeSaSession;

 

   IkeSaSession = AllocateZeroPool (sizeof (IKEV2_SA_SESSION));

-  ASSERT (IkeSaSession != NULL);

+  if (IkeSaSession == NULL) {

+    return NULL;

+  }

 

   //

   // Initialize the fields of IkeSaSession and its SessionCommon.

@@ -908,9 +910,9 @@ Ikev2ChildSaSilentDelete (
   

   SelectorSize  = sizeof (EFI_IPSEC_CONFIG_SELECTOR);

   Selector      = AllocateZeroPool (SelectorSize);

-  ASSERT (Selector != NULL);

-

-  

+  if (Selector == NULL) {

+    return EFI_OUT_OF_RESOURCES;

+  }

 

   while (1) {

     Status = EfiIpSecConfigGetNextSelector (

@@ -923,7 +925,11 @@ Ikev2ChildSaSilentDelete (
       FreePool (Selector);

 

       Selector = AllocateZeroPool (SelectorSize);

-      ASSERT (Selector != NULL);

+      if (Selector == NULL) {

+        Status = EFI_OUT_OF_RESOURCES;

+        break;

+      }

+      

       Status   = EfiIpSecConfigGetNextSelector (

                    &Private->IpSecConfig,

                    IPsecConfigDataTypeSad,

@@ -943,7 +949,11 @@ Ikev2ChildSaSilentDelete (
       //

       IsRemoteFound   = TRUE;

       RemoteSelector  = AllocateZeroPool (SelectorSize);

-      ASSERT (RemoteSelector != NULL);

+      if (RemoteSelector == NULL) {

+        Status = EFI_OUT_OF_RESOURCES;

+        break;

+      }

+      

       CopyMem (RemoteSelector, Selector, SelectorSize);

     }

 

@@ -954,7 +964,11 @@ Ikev2ChildSaSilentDelete (
       //

       IsLocalFound  = TRUE;

       LocalSelector = AllocateZeroPool (SelectorSize);

-      ASSERT (LocalSelector != NULL);

+      if (LocalSelector == NULL) {

+        Status = EFI_OUT_OF_RESOURCES;

+        break;

+      }

+      

       CopyMem (LocalSelector, Selector, SelectorSize);

     }

   }

@@ -1270,7 +1284,11 @@ Ikev2InitializeSaData (
     ChildSaSession              = IKEV2_CHILD_SA_SESSION_FROM_COMMON (SessionCommon);

     ProposalData->ProtocolId    = IPSEC_PROTO_IPSEC_ESP;

     ProposalData->Spi           = AllocateZeroPool (sizeof (ChildSaSession->LocalPeerSpi));

-    ASSERT (ProposalData->Spi != NULL);

+    if (ProposalData->Spi == NULL) {

+      FreePool (SaData);

+      return NULL;

+    }

+    

     CopyMem (

       ProposalData->Spi,

       &ChildSaSession->LocalPeerSpi,

@@ -1338,7 +1356,12 @@ Ikev2InitializeSaData (
     ProposalData->ProtocolId    = IPSEC_PROTO_IPSEC_ESP;

     ProposalData->NumTransforms = 3;

     ProposalData->Spi           = AllocateZeroPool (sizeof (ChildSaSession->LocalPeerSpi));

-    ASSERT (ProposalData->Spi != NULL);

+    if (ProposalData->Spi == NULL) {

+      FreePool (((IKEV2_PROPOSAL_DATA *) (SaData + 1))->Spi);

+      FreePool (SaData);

+      return NULL;

+    }

+    

     CopyMem (

       ProposalData->Spi,

       &ChildSaSession->LocalPeerSpi,

@@ -1731,17 +1754,27 @@ Ikev2ResendNotify (
   than the one in ChildSaSession->Spd, especially for the tunnel mode.

     

   @param[in, out]  ChildSaSession  Pointer to IKEV2_CHILD_SA_SESSION related to.

+

+  @retval EFI_SUCCESS            The operation complete successfully.

+  @retval EFI_OUT_OF_RESOURCES   If the required resource can't be allocated.

   

 **/

-VOID

+EFI_STATUS

 Ikev2ChildSaSessionSpdSelectorCreate (

   IN OUT IKEV2_CHILD_SA_SESSION *ChildSaSession

   ) 

 {

+  EFI_STATUS          Status;

+

+  Status = EFI_SUCCESS;

+

   if (ChildSaSession->Spd != NULL && ChildSaSession->Spd->Selector != NULL) {

     if (ChildSaSession->SpdSelector == NULL) {

       ChildSaSession->SpdSelector = AllocateZeroPool (sizeof (EFI_IPSEC_SPD_SELECTOR));

-      ASSERT (ChildSaSession->SpdSelector != NULL);

+      if (ChildSaSession->SpdSelector == NULL) {

+        Status = EFI_OUT_OF_RESOURCES;

+        return Status;

+      }

     }

     CopyMem (

       ChildSaSession->SpdSelector, 

@@ -1753,18 +1786,34 @@ Ikev2ChildSaSessionSpdSelectorCreate (
                                                    sizeof (EFI_IP_ADDRESS_INFO), 

                                                    ChildSaSession->Spd->Selector->RemoteAddress

                                                    );

+    if (ChildSaSession->SpdSelector->RemoteAddress == NULL) {

+      Status = EFI_OUT_OF_RESOURCES;

+

+      FreePool (ChildSaSession->SpdSelector);

+      

+      return Status;

+    }

+    

     ChildSaSession->SpdSelector->LocalAddress = AllocateCopyPool (

                                                   ChildSaSession->Spd->Selector->LocalAddressCount * 

                                                   sizeof (EFI_IP_ADDRESS_INFO), 

                                                   ChildSaSession->Spd->Selector->LocalAddress

                                                   );

+    if (ChildSaSession->SpdSelector->LocalAddress == NULL) {

+      Status = EFI_OUT_OF_RESOURCES;

+

+      FreePool (ChildSaSession->SpdSelector->RemoteAddress);

 

-    ASSERT (ChildSaSession->SpdSelector->LocalAddress != NULL);

-    ASSERT (ChildSaSession->SpdSelector->RemoteAddress != NULL);

+      FreePool (ChildSaSession->SpdSelector);

+      

+      return Status;

+    }

 

     ChildSaSession->SpdSelector->RemoteAddressCount = ChildSaSession->Spd->Selector->RemoteAddressCount;

     ChildSaSession->SpdSelector->LocalAddressCount = ChildSaSession->Spd->Selector->LocalAddressCount; 

   }

+

+  return Status;

 }

 

 /**

@@ -1789,7 +1838,9 @@ Ikev2ChildSaSessionCreate (
   // Create a new ChildSaSession.Insert it into processing list and initiate the common parameters.

   //

   ChildSaSession = Ikev2ChildSaSessionAlloc (UdpService, IkeSaSession);

-  ASSERT (ChildSaSession != NULL);

+  if (ChildSaSession == NULL) {

+    return NULL;

+  }

 

   //

   // Set the specific parameters.

@@ -1810,18 +1861,29 @@ Ikev2ChildSaSessionCreate (
   // The ChildSaSession->SpdSelector might be changed after the traffic selector

   // negoniation and it will be copied into the SAData after ChildSA established.

   //

-  Ikev2ChildSaSessionSpdSelectorCreate (ChildSaSession);

+  if (EFI_ERROR (Ikev2ChildSaSessionSpdSelectorCreate (ChildSaSession))) {

+    Ikev2ChildSaSessionFree (ChildSaSession);

+    return NULL;

+  }

 

   //

   // Copy first NiBlock and NrBlock to ChildSa Session

   //

   ChildSaSession->NiBlock   = AllocateZeroPool (IkeSaSession->NiBlkSize);

-  ASSERT (ChildSaSession->NiBlock != NULL);

+  if (ChildSaSession->NiBlock == NULL) {

+    Ikev2ChildSaSessionFree (ChildSaSession);

+    return NULL;

+  }

+  

   ChildSaSession->NiBlkSize = IkeSaSession->NiBlkSize;

   CopyMem (ChildSaSession->NiBlock, IkeSaSession->NiBlock, IkeSaSession->NiBlkSize);

 

   ChildSaSession->NrBlock   = AllocateZeroPool (IkeSaSession->NrBlkSize);

-  ASSERT (ChildSaSession->NrBlock != NULL);

+  if (ChildSaSession->NrBlock == NULL) {

+    Ikev2ChildSaSessionFree (ChildSaSession);

+    return NULL;

+  }

+  

   ChildSaSession->NrBlkSize = IkeSaSession->NrBlkSize;

   CopyMem (ChildSaSession->NrBlock, IkeSaSession->NrBlock, IkeSaSession->NrBlkSize);

 

@@ -2194,7 +2256,10 @@ Ikev2SaParseSaPayload (
             // Find the matched one. 

             //

             IkeSaSession->SessionCommon.SaParams = AllocateZeroPool (sizeof (IKEV2_SA_PARAMS));

-            ASSERT (IkeSaSession->SessionCommon.SaParams != NULL);

+            if (IkeSaSession->SessionCommon.SaParams == NULL) {

+              return FALSE;

+            }

+            

             IkeSaSession->SessionCommon.SaParams->EncAlgId   = PreferEncryptAlgorithm;

             IkeSaSession->SessionCommon.SaParams->EnckeyLen  = PreferEncryptKeylength;

             IkeSaSession->SessionCommon.SaParams->DhGroup    = PreferDhGroup;

@@ -2209,7 +2274,10 @@ Ikev2SaParseSaPayload (
                                    sizeof (IKEV2_PROPOSAL_DATA) +

                                    sizeof (IKEV2_TRANSFORM_DATA) * 4;

             IkeSaSession->SaData = AllocateZeroPool (SaDataSize);

-            ASSERT (IkeSaSession->SaData != NULL);

+            if (IkeSaSession->SaData == NULL) {

+              FreePool (IkeSaSession->SessionCommon.SaParams);

+              return FALSE;

+            }

 

             IkeSaSession->SaData->NumProposals  = 1;

 

@@ -2225,6 +2293,7 @@ Ikev2SaParseSaPayload (
               );

 

             ((IKEV2_PROPOSAL_DATA *) (IkeSaSession->SaData + 1))->ProposalIndex = 1;

+            

             return TRUE;

           } else {

             PreferEncryptAlgorithm   = 0;

@@ -2300,7 +2369,10 @@ Ikev2SaParseSaPayload (
 

     if (IsMatch) {

         IkeSaSession->SessionCommon.SaParams = AllocateZeroPool (sizeof (IKEV2_SA_PARAMS));

-        ASSERT (IkeSaSession->SessionCommon.SaParams != NULL);

+        if (IkeSaSession->SessionCommon.SaParams == NULL) {

+          return FALSE;

+        }

+        

         IkeSaSession->SessionCommon.SaParams->EncAlgId   = PreferEncryptAlgorithm;

         IkeSaSession->SessionCommon.SaParams->EnckeyLen  = PreferEncryptKeylength;

         IkeSaSession->SessionCommon.SaParams->DhGroup    = PreferDhGroup;

@@ -2311,6 +2383,7 @@ Ikev2SaParseSaPayload (
         return TRUE;

     }

   }

+  

   return FALSE;

 }

 

@@ -2391,7 +2464,10 @@ Ikev2ChildSaParseSaPayload (
           // Find the matched one. 

           //

           ChildSaSession->SessionCommon.SaParams = AllocateZeroPool (sizeof (IKEV2_SA_PARAMS));

-          ASSERT (ChildSaSession->SessionCommon.SaParams != NULL);

+          if (ChildSaSession->SessionCommon.SaParams == NULL) {

+            return FALSE;

+          }

+          

           ChildSaSession->SessionCommon.SaParams->EncAlgId   = PreferEncryptAlgorithm;

           ChildSaSession->SessionCommon.SaParams->EnckeyLen  = PreferEncryptKeylength;

           ChildSaSession->SessionCommon.SaParams->IntegAlgId = PreferIntegrityAlgorithm;

@@ -2405,7 +2481,10 @@ Ikev2ChildSaParseSaPayload (
                                  sizeof (IKEV2_TRANSFORM_DATA) * 4;

 

           ChildSaSession->SaData = AllocateZeroPool (SaDataSize);

-          ASSERT (ChildSaSession->SaData != NULL);

+          if (ChildSaSession->SaData == NULL) {

+            FreePool (ChildSaSession->SessionCommon.SaParams);

+            return FALSE;

+          }

 

           ChildSaSession->SaData->NumProposals  = 1;

 

@@ -2426,7 +2505,14 @@ Ikev2ChildSaParseSaPayload (
                                                                           sizeof (ChildSaSession->LocalPeerSpi), 

                                                                           &ChildSaSession->LocalPeerSpi

                                                                           );

-          ASSERT (((IKEV2_PROPOSAL_DATA *) (ChildSaSession->SaData + 1))->Spi != NULL);

+          if (((IKEV2_PROPOSAL_DATA *) (ChildSaSession->SaData + 1))->Spi == NULL) {

+            FreePool (ChildSaSession->SessionCommon.SaParams);

+

+            FreePool (ChildSaSession->SaData );

+            

+            return FALSE;

+          }

+          

           return TRUE;

 

         } else {

@@ -2496,7 +2582,10 @@ Ikev2ChildSaParseSaPayload (
     ProposalData  = (IKEV2_PROPOSAL_DATA *)((IKEV2_SA_DATA *)SaPayload->PayloadBuf + 1);

     if (IsMatch) {

         ChildSaSession->SessionCommon.SaParams = AllocateZeroPool (sizeof (IKEV2_SA_PARAMS));

-        ASSERT (ChildSaSession->SessionCommon.SaParams != NULL);

+        if (ChildSaSession->SessionCommon.SaParams == NULL) {

+          return FALSE;

+        }

+        

         ChildSaSession->SessionCommon.SaParams->EncAlgId   = PreferEncryptAlgorithm;

         ChildSaSession->SessionCommon.SaParams->EnckeyLen  = PreferEncryptKeylength;

         ChildSaSession->SessionCommon.SaParams->IntegAlgId = PreferIntegrityAlgorithm;

@@ -2605,7 +2694,11 @@ Ikev2SaGenerateKey (
   }

 

   LocalFragments[1].Data     = AllocateZeroPool (FragmentsSize);

-  ASSERT (LocalFragments[1].Data != NULL);

+  if (LocalFragments[1].Data == NULL) {

+    Status = EFI_OUT_OF_RESOURCES;

+    goto Exit;

+  }

+  

   LocalFragments[1].DataSize = FragmentsSize;

 

   //

@@ -2631,7 +2724,11 @@ Ikev2SaGenerateKey (
   // Allocate buffer for the first fragment

   //

   LocalFragments[0].Data     = AllocateZeroPool (AuthKeyLength);

-  ASSERT (LocalFragments[0].Data != NULL);

+  if (LocalFragments[0].Data == NULL) {

+    Status = EFI_OUT_OF_RESOURCES;

+    goto Exit;

+  }

+  

   LocalFragments[0].DataSize = AuthKeyLength;

 

   Round = (OutputKeyLength - 1) / AuthKeyLength + 1;

diff --git a/NetworkPkg/IpSecDxe/Ikev2/Utility.h b/NetworkPkg/IpSecDxe/Ikev2/Utility.h
index c018f425fc..319b6cb32c 100644
--- a/NetworkPkg/IpSecDxe/Ikev2/Utility.h
+++ b/NetworkPkg/IpSecDxe/Ikev2/Utility.h
@@ -2,7 +2,7 @@
   The interfaces of IKE/Child session operations and payload related operations 

   used by IKE Exchange Process.

 

-  Copyright (c) 2010, Intel Corporation. All rights reserved.<BR>

+  Copyright (c) 2010 - 2016, Intel Corporation. All rights reserved.<BR>

 

   This program and the accompanying materials

   are licensed and made available under the terms and conditions of the BSD License

@@ -1119,9 +1119,12 @@ Ikev2SaGenerateKey (
   than the one in ChildSaSession->Spd, especially for the tunnel mode.

     

   @param[in, out]  ChildSaSession  Pointer to IKEV2_CHILD_SA_SESSION related to.

+

+  @retval EFI_SUCCESS            The operation complete successfully.

+  @retval EFI_OUT_OF_RESOURCES   If the required resource can't be allocated.

   

 **/

-VOID

+EFI_STATUS

 Ikev2ChildSaSessionSpdSelectorCreate (

   IN OUT IKEV2_CHILD_SA_SESSION *ChildSaSession

   );

diff --git a/NetworkPkg/IpSecDxe/IpSecConfigImpl.c b/NetworkPkg/IpSecDxe/IpSecConfigImpl.c
index e1b24e4355..cfee978d94 100644
--- a/NetworkPkg/IpSecDxe/IpSecConfigImpl.c
+++ b/NetworkPkg/IpSecDxe/IpSecConfigImpl.c
@@ -2175,7 +2175,10 @@ IpSecGetVariable (
   VariableNameLength  = StrLen (VariableName);

   VariableNameISize   = (VariableNameLength + 5) * sizeof (CHAR16);

   VariableNameI       = AllocateZeroPool (VariableNameISize);

-  ASSERT (VariableNameI != NULL);

+  if (VariableNameI == NULL) {

+    Status = EFI_OUT_OF_RESOURCES;

+    goto ON_EXIT;

+  }

   

   //

   // Construct the varible name of ipsecconfig meta data.

diff --git a/NetworkPkg/IpSecDxe/IpSecCryptIo.c b/NetworkPkg/IpSecDxe/IpSecCryptIo.c
index 8396c59003..dca44231be 100644
--- a/NetworkPkg/IpSecDxe/IpSecCryptIo.c
+++ b/NetworkPkg/IpSecDxe/IpSecCryptIo.c
@@ -1,7 +1,7 @@
 /** @file

   Common interfaces to call Security library.

 

-  Copyright (c) 2009 - 2014, Intel Corporation. All rights reserved.<BR>

+  Copyright (c) 2009 - 2016, Intel Corporation. All rights reserved.<BR>

 

   This program and the accompanying materials

   are licensed and made available under the terms and conditions of the BSD License

@@ -968,7 +968,10 @@ IpSecCryptoIoGetPublicKeyFromCert (
   RsaGetKey (RsaContext, RsaKeyN, NULL, PublicKeyLen);

  

   *PublicKey = AllocateZeroPool (*PublicKeyLen);

-  ASSERT (*PublicKey != NULL);

+  if (*PublicKey == NULL) {

+    Status = EFI_OUT_OF_RESOURCES;

+    goto EXIT;

+  }

 

   if (!RsaGetKey (RsaContext, RsaKeyN, *PublicKey, PublicKeyLen)) {

     Status = EFI_INVALID_PARAMETER;

diff --git a/NetworkPkg/IpSecDxe/IpSecImpl.c b/NetworkPkg/IpSecDxe/IpSecImpl.c
index 854a9a58d8..625f154ff4 100644
--- a/NetworkPkg/IpSecDxe/IpSecImpl.c
+++ b/NetworkPkg/IpSecDxe/IpSecImpl.c
@@ -2,7 +2,7 @@
   The implementation of IPsec.

 

   (C) Copyright 2015 Hewlett-Packard Development Company, L.P.<BR>

-  Copyright (c) 2009 - 2011, Intel Corporation. All rights reserved.<BR>

+  Copyright (c) 2009 - 2016, Intel Corporation. All rights reserved.<BR>

 

   This program and the accompanying materials

   are licensed and made available under the terms and conditions of the BSD License

@@ -1190,9 +1190,6 @@ IpSecTunnelInboundPacket (
                                      on return.

   @param[in]      FragmentCount      The number of fragments.

 

-  @retval EFI_SUCCESS              The operation was successful.

-  @retval EFI_OUT_OF_RESOURCES     The required system resources can't be allocated.

-

 **/

 UINT8 *

 IpSecTunnelOutboundPacket (

@@ -1220,7 +1217,10 @@ IpSecTunnelOutboundPacket (
 

   if (IpVersion == IP_VERSION_4) {

     InnerHead = AllocateZeroPool (sizeof (IP4_HEAD) + *OptionsLength);

-    ASSERT (InnerHead != NULL);

+    if (InnerHead == NULL) {

+      return NULL;

+    }

+    

     CopyMem (

       InnerHead,

       IpHead,

@@ -1233,7 +1233,10 @@ IpSecTunnelOutboundPacket (
       );

   } else {

     InnerHead = AllocateZeroPool (sizeof (EFI_IP6_HEADER) + *OptionsLength);

-    ASSERT (InnerHead != NULL);

+    if (InnerHead == NULL) {

+      return NULL;

+    }

+    

     CopyMem (

       InnerHead,

       IpHead,

@@ -1264,7 +1267,11 @@ IpSecTunnelOutboundPacket (
              IpSecOnRecyclePacket,

              NULL

              );

-  ASSERT (Packet != NULL);

+  if (Packet == NULL) {

+    FreePool (InnerHead);

+    return NULL;

+  }

+  

   //

   // 3. Check the Last Header, if it is TCP, UDP or ICMP recalcualate its pesudo

   //    CheckSum.