From 68161ff0dd64dc6cb6fc2e1f3bf461b5f4db3979 Mon Sep 17 00:00:00 2001
From: Satyanvesh Dittakavi <Satyanvesh.Dittakavi@amd.com>
Date: Tue, 6 Apr 2021 06:37:14 -0700
Subject: [PATCH] SWDEV-279486 - Fix hipMemcpyPeer/hipMemcpy APIs

Addresses the scenarios when the size passed is more than
the allocated size and when the deviceid is invalid
Change-Id: I6c9b62639096f655ffb61976905b1ce8c5f51ee7

Change-Id: I8a0d660924a8e2300c517aba6f9088626b8f6ef5
---
 rocclr/hip_memory.cpp             |  6 ++++++
 rocclr/hip_peer.cpp               | 12 ++++++++++++
 tests/src/ipc/hipSimpleIpcMem.cpp |  2 +-
 3 files changed, 19 insertions(+), 1 deletion(-)

diff --git a/rocclr/hip_memory.cpp b/rocclr/hip_memory.cpp
index e1f937ae35..4a8f4f9434 100755
--- a/rocclr/hip_memory.cpp
+++ b/rocclr/hip_memory.cpp
@@ -212,6 +212,12 @@ hipError_t ihipMemcpy(void* dst, const void* src, size_t sizeBytes, hipMemcpyKin
   amd::Memory *srcMemory = getMemoryObject(src, sOffset);
   size_t dOffset = 0;
   amd::Memory *dstMemory = getMemoryObject(dst, dOffset);
+  // Return error if sizeBytes passed to memcpy is more than the actual size allocated
+  if ((dstMemory && sizeBytes > (dstMemory->getSize() - dOffset)) ||
+      (srcMemory && sizeBytes > (srcMemory->getSize() - sOffset))) {
+    return hipErrorInvalidValue;
+  }
+
   amd::Device* queueDevice = &queue.device();
   if ((srcMemory == nullptr) && (dstMemory == nullptr)) {
     if ((kind == hipMemcpyHostToHost) || (kind == hipMemcpyDefault)) {
diff --git a/rocclr/hip_peer.cpp b/rocclr/hip_peer.cpp
index bb9f118b10..74c30fb411 100755
--- a/rocclr/hip_peer.cpp
+++ b/rocclr/hip_peer.cpp
@@ -215,6 +215,12 @@ hipError_t hipMemcpyPeer(void* dst, int dstDevice, const void* src, int srcDevic
                          size_t sizeBytes) {
   HIP_INIT_API(hipMemcpyPeer, dst, dstDevice, src, srcDevice, sizeBytes);
 
+  if (srcDevice >= static_cast<int>(g_devices.size()) ||
+      dstDevice >= static_cast<int>(g_devices.size()) ||
+      srcDevice < 0 || dstDevice < 0) {
+    HIP_RETURN(hipErrorInvalidDevice);
+  }
+
   HIP_RETURN(hipMemcpy(dst, src, sizeBytes, hipMemcpyDeviceToDevice));
 }
 
@@ -222,6 +228,12 @@ hipError_t hipMemcpyPeerAsync(void* dst, int dstDevice, const void* src, int src
                               size_t sizeBytes, hipStream_t stream) {
   HIP_INIT_API(hipMemcpyPeerAsync, dst, dstDevice, src, srcDevice, sizeBytes, stream);
 
+  if (srcDevice >= static_cast<int>(g_devices.size()) ||
+      dstDevice >= static_cast<int>(g_devices.size()) ||
+      srcDevice < 0 || dstDevice < 0) {
+    HIP_RETURN(hipErrorInvalidDevice);
+  }
+
   HIP_RETURN(hipMemcpyAsync(dst, src, sizeBytes, hipMemcpyDeviceToDevice, stream));
 }
 
diff --git a/tests/src/ipc/hipSimpleIpcMem.cpp b/tests/src/ipc/hipSimpleIpcMem.cpp
index c9c3540dfe..e98201daaa 100755
--- a/tests/src/ipc/hipSimpleIpcMem.cpp
+++ b/tests/src/ipc/hipSimpleIpcMem.cpp
@@ -60,7 +60,7 @@ hipError_t single_process(int32_t offset) {
                                              ipc_offset_handle, 0));
 
   // Copy Values from Device to Host and Check for correctness
-  HIPCHECK_RETURN_ONFAIL(hipMemcpy(ipc_out_hptr, ipc_out_dptr, (NUM * sizeof(int)), hipMemcpyDeviceToHost));
+  HIPCHECK_RETURN_ONFAIL(hipMemcpy(ipc_out_hptr, ipc_out_dptr, ((NUM-offset) * sizeof(int)), hipMemcpyDeviceToHost));
   for (size_t idx = offset; idx < NUM; ++idx) {
     if (ipc_out_hptr[idx-offset] != ipc_dptr[idx]) {
       std::cout<<"Failing @ idx: "<<idx<<std::endl;