From ac4041d90628125d12b614dcc06fc27bf505e22f Mon Sep 17 00:00:00 2001
From: Michal Mrozek <michal.mrozek@intel.com>
Date: Wed, 22 Jan 2020 15:59:11 +0100
Subject: [PATCH] Add more robust validation of inputs.

- prevent USM device pointers on transfer calls
- prevent pointers that do not hold enough storage to service transfer.

Change-Id: I678808c034f708e9d0ae477d632788aae7f70452
Signed-off-by: Michal Mrozek <michal.mrozek@intel.com>
---
 runtime/command_queue/enqueue_write_buffer.h  |  7 +++
 .../unified_memory_manager_tests.cpp          | 52 +++++++++++++++++++
 2 files changed, 59 insertions(+)

diff --git a/runtime/command_queue/enqueue_write_buffer.h b/runtime/command_queue/enqueue_write_buffer.h
index 13d5089c6a..12276cc8c8 100644
--- a/runtime/command_queue/enqueue_write_buffer.h
+++ b/runtime/command_queue/enqueue_write_buffer.h
@@ -70,6 +70,13 @@ cl_int CommandQueueHw<GfxFamily>::enqueueWriteBuffer(
     if (!mapAllocation && this->getContext().getSVMAllocsManager()) {
         auto svmEntry = this->getContext().getSVMAllocsManager()->getSVMAlloc(ptr);
         if (svmEntry) {
+            if (svmEntry->memoryType == DEVICE_UNIFIED_MEMORY) {
+                return CL_INVALID_OPERATION;
+            }
+            if ((svmEntry->gpuAllocation->getGpuAddress() + svmEntry->size) < (castToUint64(ptr) + size)) {
+                return CL_INVALID_OPERATION;
+            }
+
             mapAllocation = svmEntry->cpuAllocation ? svmEntry->cpuAllocation : svmEntry->gpuAllocation;
         }
     }
diff --git a/unit_tests/memory_manager/unified_memory_manager_tests.cpp b/unit_tests/memory_manager/unified_memory_manager_tests.cpp
index 70a8e8dc11..3297b28824 100644
--- a/unit_tests/memory_manager/unified_memory_manager_tests.cpp
+++ b/unit_tests/memory_manager/unified_memory_manager_tests.cpp
@@ -614,6 +614,58 @@ TEST(UnfiedSharedMemoryTransferCalls, givenHostUSMllocationWhenPointerIsUsedAsWr
     ASSERT_EQ(CL_SUCCESS, status);
     clReleaseCommandQueue(commandQueue);
 }
+TEST(UnfiedSharedMemoryTransferCalls, givenDeviceUsmAllocationWhenItIsPassedToWriteBufferAsSourceThenErrorIsReturned) {
+    MockContext mockContext;
+    cl_context clContext = &mockContext;
+
+    auto status = CL_SUCCESS;
+    cl_device_id clDevice = mockContext.getDevice(0u);
+
+    auto deviceMemory = clDeviceMemAllocINTEL(clContext, clDevice, nullptr, 4096u, 0u, &status);
+
+    ASSERT_EQ(CL_SUCCESS, status);
+    auto buffer = clCreateBuffer(clContext, CL_MEM_READ_WRITE, 4096u, nullptr, &status);
+    ASSERT_EQ(CL_SUCCESS, status);
+
+    auto commandQueue = clCreateCommandQueue(clContext, clDevice, 0u, &status);
+    ASSERT_EQ(CL_SUCCESS, status);
+
+    status = clEnqueueWriteBuffer(commandQueue, buffer, false, 0u, 4096u, deviceMemory, 0u, nullptr, nullptr);
+    EXPECT_EQ(CL_INVALID_OPERATION, status);
+
+    status = clReleaseMemObject(buffer);
+    ASSERT_EQ(CL_SUCCESS, status);
+    status = clMemFreeINTEL(clContext, deviceMemory);
+    ASSERT_EQ(CL_SUCCESS, status);
+    clReleaseCommandQueue(commandQueue);
+}
+
+TEST(UnfiedSharedMemoryTransferCalls, givenHostAllocationThatIsSmallerThenWriteBufferTranfserSizeWhenTransferCallIsEmittedThenErrorIsReturned) {
+    MockContext mockContext;
+    cl_context clContext = &mockContext;
+
+    auto status = CL_SUCCESS;
+
+    auto hostMemory = clHostMemAllocINTEL(clContext, nullptr, 4u, 0u, &status);
+
+    ASSERT_EQ(CL_SUCCESS, status);
+    auto buffer = clCreateBuffer(clContext, CL_MEM_READ_WRITE, 4096u, nullptr, &status);
+    ASSERT_EQ(CL_SUCCESS, status);
+
+    cl_device_id clDevice = mockContext.getDevice(0u);
+
+    auto commandQueue = clCreateCommandQueue(clContext, clDevice, 0u, &status);
+    ASSERT_EQ(CL_SUCCESS, status);
+
+    status = clEnqueueWriteBuffer(commandQueue, buffer, false, 0u, 4096u, hostMemory, 0u, nullptr, nullptr);
+    EXPECT_EQ(CL_INVALID_OPERATION, status);
+
+    status = clReleaseMemObject(buffer);
+    ASSERT_EQ(CL_SUCCESS, status);
+    status = clMemFreeINTEL(clContext, hostMemory);
+    ASSERT_EQ(CL_SUCCESS, status);
+    clReleaseCommandQueue(commandQueue);
+}
 
 TEST(UnfiedSharedMemoryTransferCalls, givenSharedUSMllocationWithoutLocalMemoryWhenPointerIsUsedAsWriteBufferSourceThenUSMAllocationIsReused) {
     DebugManagerStateRestore restore;