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>
2016-06-17 11:59:47 +08:00
parent 415aa2f1cb
commit 6b16c9e7ea
8 changed files with 379 additions and 97 deletions
--- 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;