Fix several issues in BaseCryptLib:

[mirror_edk2.git] / CryptoPkg / Library / BaseCryptLib / Pk / CryptDh.c

diff --git a/CryptoPkg/Library/BaseCryptLib/Pk/CryptDh.c b/CryptoPkg/Library/BaseCryptLib/Pk/CryptDh.c
index 30792d27ca08f6b30eceff3d185448f8d468825f..942b3d103e260e49b942efe2395e8717d6570b2e 100644 (file)
--- a/CryptoPkg/Library/BaseCryptLib/Pk/CryptDh.c
+++ b/CryptoPkg/Library/BaseCryptLib/Pk/CryptDh.c
@@ -91,7 +91,7 @@ DhGenerateParameter (
   //\r
   // Check input parameters.\r
   //\r
-  if (DhContext == NULL || Prime == NULL) {\r
+  if (DhContext == NULL || Prime == NULL || PrimeLength > INT_MAX) {\r
     return FALSE;\r
   }\r
 \r
@@ -139,12 +139,13 @@ DhSetParameter (
   IN      CONST UINT8  *Prime\r
   )\r
 {\r
-  DH  *Dh;\r
+  DH      *Dh;\r
+  BIGNUM  *Bn;\r
 \r
   //\r
   // Check input parameters.\r
   //\r
-  if (DhContext == NULL || Prime == NULL) {\r
+  if (DhContext == NULL || Prime == NULL || PrimeLength > INT_MAX) {\r
     return FALSE;\r
   }\r
   \r
@@ -152,14 +153,46 @@ DhSetParameter (
     return FALSE;\r
   }\r
 \r
+  Bn = NULL;\r
+\r
   Dh = (DH *) DhContext;\r
-  Dh->p = BN_new();\r
-  Dh->g = BN_new();\r
+  Dh->g = NULL;\r
+  Dh->p = BN_new ();\r
+  if (Dh->p == NULL) {\r
+    goto Error;\r
+  }\r
+  \r
+  Dh->g = BN_new ();\r
+  if (Dh->g == NULL) {\r
+    goto Error;\r
+  }\r
 \r
-  BN_bin2bn (Prime, (UINT32) (PrimeLength / 8), Dh->p);\r
-  BN_set_word (Dh->g, (UINT32) Generator);\r
+  Bn = BN_bin2bn (Prime, (UINT32) (PrimeLength / 8), Dh->p);\r
+  if (Bn == NULL) {\r
+    goto Error;\r
+  }\r
+\r
+  if (BN_set_word (Dh->g, (UINT32) Generator) == 0) {\r
+    goto Error;\r
+  }\r
 \r
   return TRUE;\r
+\r
+Error:\r
+\r
+  if (Dh->p != NULL) {\r
+    BN_free (Dh->p);\r
+  }\r
+\r
+  if (Dh->g != NULL) {\r
+    BN_free (Dh->g);\r
+  }\r
+\r
+  if (Bn != NULL) {\r
+    BN_free (Bn);\r
+  }\r
+  \r
+  return FALSE;\r
 }\r
 \r
 /**\r
@@ -194,6 +227,7 @@ DhGenerateKey (
 {\r
   BOOLEAN RetVal;\r
   DH      *Dh;\r
+  INTN    Size;\r
 \r
   //\r
   // Check input parameters.\r
@@ -207,12 +241,17 @@ DhGenerateKey (
   }\r
   \r
   Dh = (DH *) DhContext;\r
-  *PublicKeySize = 0;\r
 \r
   RetVal = (BOOLEAN) DH_generate_key (DhContext);\r
   if (RetVal) {\r
+    Size = BN_num_bytes (Dh->pub_key);\r
+    if ((Size > 0) && (*PublicKeySize < (UINTN) Size)) {\r
+      *PublicKeySize = Size;\r
+      return FALSE;\r
+    }\r
+    \r
     BN_bn2bin (Dh->pub_key, PublicKey);\r
-    *PublicKeySize  = BN_num_bytes (Dh->pub_key);\r
+    *PublicKeySize = Size;\r
   }\r
 \r
   return RetVal;\r
@@ -227,7 +266,8 @@ DhGenerateKey (
   If DhContext is NULL, then return FALSE.\r
   If PeerPublicKey is NULL, then return FALSE.\r
   If KeySize is NULL, then return FALSE.\r
-  If KeySize is large enough but Key is NULL, then return FALSE.\r
+  If Key is NULL, then return FALSE.\r
+  If KeySize is not large enough, then return FALSE.\r
 \r
   @param[in, out]  DhContext          Pointer to the DH context.\r
   @param[in]       PeerPublicKey      Pointer to the peer's public key.\r
@@ -252,23 +292,37 @@ DhComputeKey (
   )\r
 {\r
   BIGNUM  *Bn;\r
+  INTN    Size;\r
 \r
   //\r
   // Check input parameters.\r
   //\r
-  if (DhContext == NULL || PeerPublicKey == NULL || KeySize == NULL) {\r
+  if (DhContext == NULL || PeerPublicKey == NULL || KeySize == NULL || Key == NULL) {\r
     return FALSE;\r
   }\r
 \r
-  if (Key == NULL && *KeySize != 0) {\r
+  if (PeerPublicKeySize > INT_MAX) {\r
     return FALSE;\r
   }\r
   \r
   Bn = BN_bin2bn (PeerPublicKey, (UINT32) PeerPublicKeySize, NULL);\r
+  if (Bn == NULL) {\r
+    return FALSE;\r
+  }\r
 \r
-  *KeySize = (BOOLEAN) DH_compute_key (Key, Bn, DhContext);\r
+  Size = DH_compute_key (Key, Bn, DhContext);\r
+  if (Size < 0) {\r
+    BN_free (Bn);\r
+    return FALSE;\r
+  }\r
 \r
-  BN_free (Bn);\r
+  if (*KeySize < (UINTN) Size) {\r
+    *KeySize = Size;\r
+    BN_free (Bn);\r
+    return FALSE;\r
+  }\r
 \r
+  *KeySize = Size;\r
+  BN_free (Bn);\r
   return TRUE;\r
 }\r