From 4c191258204213499eccc4cdec24908ad423d6b1 Mon Sep 17 00:00:00 2001
From: Savva Mitrofanov <savvamtr@gmail.com>
Date: Wed, 1 Mar 2023 16:23:36 +0600
Subject: [PATCH] Utilities: Corrects missing memory allocation checks in
 TestNtfsDxe

---
 Utilities/TestNtfsDxe/TestNtfsDxe.c | 15 +++++++++++++++
 1 file changed, 15 insertions(+)

diff --git a/Utilities/TestNtfsDxe/TestNtfsDxe.c b/Utilities/TestNtfsDxe/TestNtfsDxe.c
index 3d9f3aac..f70ac34c 100644
--- a/Utilities/TestNtfsDxe/TestNtfsDxe.c
+++ b/Utilities/TestNtfsDxe/TestNtfsDxe.c
@@ -202,6 +202,11 @@ LLVMFuzzerTestOneInput (
     Status = FileGetInfo (NewHandle, &gEfiFileInfoGuid, &Len, Info);
     if (Status == EFI_BUFFER_TOO_SMALL) {
       Info = AllocateZeroPool (Len);
+      if (Info == NULL) {
+        FreeAll (FileName, Instance);
+        return 0;
+      }
+
       FileGetInfo (NewHandle, &gEfiFileInfoGuid, &Len, Info);
       FreePool (Info);
     }
@@ -210,6 +215,11 @@ LLVMFuzzerTestOneInput (
     Status = FileGetInfo (NewHandle, &gEfiFileSystemInfoGuid, &Len, Info);
     if (Status == EFI_BUFFER_TOO_SMALL) {
       Info = AllocateZeroPool (Len);
+      if (Info == NULL) {
+        FreeAll (FileName, Instance);
+        return 0;
+      }
+
       FileGetInfo (NewHandle, &gEfiFileSystemInfoGuid, &Len, Info);
       FreePool (Info);
     }
@@ -218,6 +228,11 @@ LLVMFuzzerTestOneInput (
     Status = FileGetInfo (NewHandle, &gEfiFileSystemVolumeLabelInfoIdGuid, &Len, Info);
     if (Status == EFI_BUFFER_TOO_SMALL) {
       Info = AllocateZeroPool (Len);
+      if (Info == NULL) {
+        FreeAll (FileName, Instance);
+        return 0;
+      }
+
       FileGetInfo (NewHandle, &gEfiFileSystemVolumeLabelInfoIdGuid, &Len, Info);
       FreePool (Info);
     }