From c4b9442fc4d2faeb5015e092ff541e68f1baf9e3 Mon Sep 17 00:00:00 2001
From: Ioannis Assiouras <Ioannis.Assiouras@amd.com>
Date: Fri, 19 Jan 2024 15:56:15 +0000
Subject: [PATCH] SWDEV-442126 - Fix use_after_free case in
 ExtractFatBinaryUsingCOMGR

Change-Id: Ie07cf82025b65e4f95d4ce3cf46a6b8081f29b6a
---
 hipamd/src/hip_fatbin.cpp | 10 +++-------
 1 file changed, 3 insertions(+), 7 deletions(-)

diff --git a/hipamd/src/hip_fatbin.cpp b/hipamd/src/hip_fatbin.cpp
index 9d5591d053..5d0d00ac97 100644
--- a/hipamd/src/hip_fatbin.cpp
+++ b/hipamd/src/hip_fatbin.cpp
@@ -115,7 +115,7 @@ void ListAllDeviceWithNoCOFromBundle(const std::unordered_map<std::string,
 }
 
 hipError_t FatBinaryInfo::ExtractFatBinaryUsingCOMGR(const std::vector<hip::Device*>& devices) {
-  amd_comgr_data_t data_object;
+  amd_comgr_data_t data_object {0};
   amd_comgr_status_t comgr_status = AMD_COMGR_STATUS_SUCCESS;
   hipError_t hip_status = hipSuccess;
   amd_comgr_code_object_info_t* query_list_array = nullptr;
@@ -266,12 +266,6 @@ hipError_t FatBinaryInfo::ExtractFatBinaryUsingCOMGR(const std::vector<hip::Devi
       fatbin_dev_info_[device->deviceId()]->program_
         = new amd::Program(*(device->asContext()));
     }
-
-    if ((comgr_status = amd_comgr_release_data(data_object)) != AMD_COMGR_STATUS_SUCCESS) {
-      LogPrintfError("Releasing COMGR data failed with status %d ", comgr_status);
-      return hipErrorInvalidValue;
-    }
-
   } while(0);
 
   // Clean up file and memory resouces if hip_status failed for some reason.
@@ -292,7 +286,9 @@ hipError_t FatBinaryInfo::ExtractFatBinaryUsingCOMGR(const std::vector<hip::Devi
       fdesc_ = 0;
       fsize_ = 0;
     }
+  }
 
+  if (data_object.handle) {
     if ((comgr_status = amd_comgr_release_data(data_object)) != AMD_COMGR_STATUS_SUCCESS) {
       LogPrintfError("Releasing COMGR data failed with status %d ", comgr_status);
       return hipErrorInvalidValue;