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>

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.

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 ();
-
+  if (NoncePayload == NULL) {
-  ASSERT (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;

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);

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;
 
-    ASSERT (ChildSaSession->SpdSelector->LocalAddress != NULL);
+      FreePool (ChildSaSession->SpdSelector->RemoteAddress);
-    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;

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
   );

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.

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;

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.