Add context check and init in BaseCrypto2Hash().

[mirror_edk2.git] / SecurityPkg / Hash2DxeCrypto / Hash2DxeCrypto.c

diff --git a/SecurityPkg/Hash2DxeCrypto/Hash2DxeCrypto.c b/SecurityPkg/Hash2DxeCrypto/Hash2DxeCrypto.c
index dab02992362ed0c66caebfc769af69f89694317e..94057ab2e1fc7489ea344fbce509883b8db65352 100644 (file)
--- a/SecurityPkg/Hash2DxeCrypto/Hash2DxeCrypto.c
+++ b/SecurityPkg/Hash2DxeCrypto/Hash2DxeCrypto.c
@@ -358,6 +358,7 @@ BaseCrypto2Hash (
   UINTN                    CtxSize;\r
   BOOLEAN                  Ret;\r
   EFI_STATUS               Status;\r
+  HASH2_INSTANCE_DATA      *Instance;\r
 \r
   Status = EFI_SUCCESS;\r
 \r
@@ -373,6 +374,13 @@ BaseCrypto2Hash (
   if (HashInfo == NULL) {\r
     return EFI_UNSUPPORTED;\r
   }\r
+  \r
+  Instance = HASH2_INSTANCE_DATA_FROM_THIS(This);\r
+  if (Instance->HashContext != NULL) {\r
+    FreePool (Instance->HashContext);\r
+  }\r
+  Instance->HashInfoContext = NULL;\r
+  Instance->HashContext = NULL;\r
 \r
   //\r
   // Start hash sequence\r
@@ -392,6 +400,12 @@ BaseCrypto2Hash (
     goto Done;\r
   }\r
 \r
+  //\r
+  // Setup the context\r
+  //\r
+  Instance->HashContext = HashCtx;\r
+  Instance->HashInfoContext = HashInfo;\r
+\r
   Ret = HashInfo->Update (HashCtx, Message, MessageSize);\r
   if (!Ret) {\r
     Status = EFI_OUT_OF_RESOURCES;\r
@@ -404,7 +418,12 @@ BaseCrypto2Hash (
     goto Done;\r
   }\r
 Done:\r
+  //\r
+  // Cleanup the context\r
+  //\r
   FreePool (HashCtx);\r
+  Instance->HashInfoContext = NULL;\r
+  Instance->HashContext = NULL;\r
   return Status;\r
 }\r
 \r