From 3a45a965e222c2472a8b020c25d1daa47b0afd10 Mon Sep 17 00:00:00 2001
From: sdashmiz <shadi.dashmiz@amd.com>
Date: Fri, 10 Mar 2023 15:13:38 -0500
Subject: [PATCH] SWDEV-386695 - fix param sanitization check

Signed-off-by: sdashmiz <shadi.dashmiz@amd.com>
Change-Id: I1317640453d56718fa41f455c512282fd0f01e2a
---
 hipamd/src/hip_graph.cpp | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/hipamd/src/hip_graph.cpp b/hipamd/src/hip_graph.cpp
index f9a699e130..406c259a0a 100644
--- a/hipamd/src/hip_graph.cpp
+++ b/hipamd/src/hip_graph.cpp
@@ -2162,6 +2162,16 @@ hipError_t hipGraphAddMemAllocNode(hipGraphNode_t* pGraphNode, hipGraph_t graph,
       (numDependencies > 0 && pDependencies == nullptr) || pNodeParams == nullptr) {
     HIP_RETURN(hipErrorInvalidValue);
   }
+  if (pNodeParams->bytesize == 0 || pNodeParams->poolProps.allocType != hipMemAllocationTypePinned
+      || pNodeParams->poolProps.location.type != hipMemLocationTypeDevice) {
+    pNodeParams->dptr = nullptr;
+    HIP_RETURN(hipErrorInvalidValue);
+  }
+  if (pNodeParams->poolProps.location.type == hipMemLocationTypeDevice) {
+    if (pNodeParams->poolProps.location.id < 0 || pNodeParams->poolProps.location.id >= g_devices.size()) {
+      HIP_RETURN(hipErrorInvalidValue);
+    }
+  }
   // Clear the pointer to allocated memory because it may contain stale/uninitialized data
   pNodeParams->dptr = nullptr;
   auto mem_alloc_node = new hipGraphMemAllocNode(pNodeParams);