From 98fef69edd1fb01a34bc31d2f2d2d39c600942f5 Mon Sep 17 00:00:00 2001
From: Satyanvesh Dittakavi <Satyanvesh.Dittakavi@amd.com>
Date: Tue, 11 Jan 2022 15:49:48 +0000
Subject: [PATCH] SWDEV-313867 - SWDEV-314101 - Address Negative scenarios with
 hipGraphAddMemcpyNode and hipGraphAddMemsetNode

Change-Id: I51527ce6953aee9a3ef7d821754819b6c8087939


[ROCm/clr commit: c58ba64a0bd279cbb83e4d01baf2860dc4772767]
---
 projects/clr/hipamd/src/hip_graph.cpp  | 10 ++++++++--
 projects/clr/hipamd/src/hip_memory.cpp |  9 +++++++--
 2 files changed, 15 insertions(+), 4 deletions(-)

diff --git a/projects/clr/hipamd/src/hip_graph.cpp b/projects/clr/hipamd/src/hip_graph.cpp
index 8ef1f5490f..faf8b5c2ba 100644
--- a/projects/clr/hipamd/src/hip_graph.cpp
+++ b/projects/clr/hipamd/src/hip_graph.cpp
@@ -129,10 +129,16 @@ hipError_t ihipGraphAddMemcpyNode1D(hipGraphNode_t* pGraphNode, hipGraph_t graph
 hipError_t ihipGraphAddMemsetNode(hipGraphNode_t* pGraphNode, hipGraph_t graph,
                                   const hipGraphNode_t* pDependencies, size_t numDependencies,
                                   const hipMemsetParams* pMemsetParams) {
-  if (pGraphNode == nullptr || graph == nullptr ||
-      (numDependencies > 0 && pDependencies == nullptr) || pMemsetParams == nullptr) {
+  if (pGraphNode == nullptr || graph == nullptr || pMemsetParams == nullptr ||
+      (numDependencies > 0 && pDependencies == nullptr) || pMemsetParams->height == 0) {
     return hipErrorInvalidValue;
   }
+  // The element size must be 1, 2, or 4 bytes
+  if (pMemsetParams->elementSize != sizeof(int8_t) && pMemsetParams->elementSize != sizeof(int16_t)
+      && pMemsetParams->elementSize != sizeof(int32_t)) {
+    return hipErrorInvalidValue;
+  }
+
   hipError_t status;
   status = ihipGraphMemsetParams_validate(pMemsetParams);
   if (status != hipSuccess) {
diff --git a/projects/clr/hipamd/src/hip_memory.cpp b/projects/clr/hipamd/src/hip_memory.cpp
index 1848b5bc4c..2b3dec2b0c 100644
--- a/projects/clr/hipamd/src/hip_memory.cpp
+++ b/projects/clr/hipamd/src/hip_memory.cpp
@@ -2168,13 +2168,18 @@ hipError_t hipMemcpyAtoH(void* dstHost,
 }
 
 hipError_t ihipMemcpy3D_validate(const hipMemcpy3DParms* p) {
-  // The struct passed to hipMemcpy3D() must specify one of srcArray or srcPtr and one of dstArray
-  // or dstPtr. Passing more than one non-zero source or destination will cause hipMemcpy3D() to
+  // Passing more than one non-zero source or destination will cause hipMemcpy3D() to
   // return an error.
   if (p == nullptr || ((p->srcArray != nullptr) && (p->srcPtr.ptr != nullptr)) ||
       ((p->dstArray != nullptr) && (p->dstPtr.ptr != nullptr))) {
     return hipErrorInvalidValue;
   }
+  // The struct passed to hipMemcpy3D() must specify one of srcArray or srcPtr and one of dstArray
+  // or dstPtr.
+  if (((p->srcArray == nullptr) && (p->srcPtr.ptr == nullptr)) ||
+      ((p->dstArray == nullptr) && (p->dstPtr.ptr == nullptr))) {
+    return hipErrorInvalidValue;
+  }
 
   // If the source and destination are both arrays, hipMemcpy3D() will return an error if they do
   // not have the same element size.