From 8d9d387e38b4ff081a776580c684dbd2c454a0d9 Mon Sep 17 00:00:00 2001
From: vit9696 <vit9696@users.noreply.github.com>
Date: Sun, 6 Mar 2022 15:24:20 +0300
Subject: [PATCH] OcCryptoLib: Drop BigNumMod memory allocation

BigNumCalculateMontParams previously did not check BigNumMod for memory
allocation failure anyway, which was thus harder to diagnose due to
the Montgomery Inverse of N still returning correctly.
---
 Library/OcCryptoLib/BigNumLibInternal.h |  9 ++++-----
 Library/OcCryptoLib/BigNumMontgomery.c  |  9 ++++++++-
 Library/OcCryptoLib/BigNumPrimitives.c  | 15 +++------------
 3 files changed, 15 insertions(+), 18 deletions(-)

diff --git a/Library/OcCryptoLib/BigNumLibInternal.h b/Library/OcCryptoLib/BigNumLibInternal.h
index fe694fd3ef6..e483a47ce3c 100644
--- a/Library/OcCryptoLib/BigNumLibInternal.h
+++ b/Library/OcCryptoLib/BigNumLibInternal.h
@@ -95,17 +95,16 @@ BigNumOrWord (
   @param[in]     A             The dividend.
   @param[in]     NumWordsA     The number of Words of A.
   @param[in]     B             The divisor.
-
-  @returns  Whether the operation was completes successfully.
-
+  @param[in]     Memory        Scratch buffer 2 * NumWordsA * OC_BN_WORD_SIZE.
 **/
-BOOLEAN
+VOID
 BigNumMod (
   IN OUT OC_BN_WORD        *Result,
   IN     OC_BN_NUM_WORDS   NumWordsRest,
   IN     CONST OC_BN_WORD  *A,
   IN     OC_BN_NUM_WORDS   NumWordsA,
-  IN     CONST OC_BN_WORD  *B
+  IN     CONST OC_BN_WORD  *B,
+  IN     VOID              *Memory
   );
 
 /**
diff --git a/Library/OcCryptoLib/BigNumMontgomery.c b/Library/OcCryptoLib/BigNumMontgomery.c
index 7a3a1e50455..ee30681023f 100644
--- a/Library/OcCryptoLib/BigNumMontgomery.c
+++ b/Library/OcCryptoLib/BigNumMontgomery.c
@@ -162,6 +162,7 @@ BigNumCalculateMontParams (
   UINTN           SizeRSqr;
   OC_BN_NUM_WORDS NumWordsRSqr;
   OC_BN_WORD      *RSqr;
+  VOID            *Memory;
 
   ASSERT (RSqrMod != NULL);
   ASSERT (NumWords > 0);
@@ -209,7 +210,13 @@ BigNumCalculateMontParams (
   // 2 * NumBits cannot overflow as per above.
   //
   BigNumOrWord (RSqr, NumWordsRSqr, 1, 2 * NumBits);
-  BigNumMod (RSqrMod, NumWords, RSqr, NumWordsRSqr, N);
+
+  Memory = AllocatePool (2 * NumWordsRSqr * OC_BN_WORD_SIZE);
+  if (Memory == NULL) {
+    return 0;
+  }
+  BigNumMod (RSqrMod, NumWords, RSqr, NumWordsRSqr, N, Memory);
+  FreePool (Memory);
 
   FreePool (RSqr);
 
diff --git a/Library/OcCryptoLib/BigNumPrimitives.c b/Library/OcCryptoLib/BigNumPrimitives.c
index cd993b6dc2b..d0082b0ea95 100644
--- a/Library/OcCryptoLib/BigNumPrimitives.c
+++ b/Library/OcCryptoLib/BigNumPrimitives.c
@@ -558,19 +558,18 @@ BigNumCmp (
   return 0;
 }
 
-BOOLEAN
+VOID
 BigNumMod (
   IN OUT OC_BN_WORD        *Result,
   IN     OC_BN_NUM_WORDS   NumWordsRest,
   IN     CONST OC_BN_WORD  *A,
   IN     OC_BN_NUM_WORDS   NumWordsA,
-  IN     CONST OC_BN_WORD  *B
+  IN     CONST OC_BN_WORD  *B,
+  IN     VOID              *Memory
   )
 {
   INTN            CmpResult;
 
-  VOID            *Memory;
-
   OC_BN_WORD      *ModTmp;
   OC_BN_NUM_BITS  SigBitsModTmp;
   OC_BN_NUM_WORDS SigWordsModTmp;
@@ -602,11 +601,6 @@ BigNumMod (
     "An overflow verification must be added"
     );
 
-  Memory = AllocatePool (2 * NumWordsA * OC_BN_WORD_SIZE);
-  if (Memory == NULL) {
-    return FALSE;
-  }
-
   ModTmp         = Memory;
   BigDiv         = &ModTmp[SigWordsModTmp];
   SigWordsBigDiv = SigWordsModTmp;
@@ -745,9 +739,6 @@ BigNumMod (
   //
   ASSERT (BigNumMostSignificantWord (ModTmp, SigWordsModTmp) + 1 <= NumWordsRest);
   CopyMem (Result, ModTmp, NumWordsRest * OC_BN_WORD_SIZE);
-
-  FreePool (Memory);
-  return TRUE;
 }
 
 VOID