From 8c62bc7993daaead74d7522cc247030894a16ee2 Mon Sep 17 00:00:00 2001
From: "Spruit, Neil R" <neil.r.spruit@intel.com>
Date: Tue, 24 Sep 2024 01:49:20 +0200
Subject: [PATCH] [L0] Check if buffer is freed before allowing a zehandle to
 be read

- Avoid use after free of buffers by tracking if a buffer has already
  been called with free() and verify that the parent of a sub buffer has
not already been released before attempting to get ze handles from the
parent.
- In either case, the ze handle returned will be a nullptr along with
  the UR error UR_RESULT_ERROR_INVALID_NULL_HANDLE

Signed-off-by: Spruit, Neil R <neil.r.spruit@intel.com>
---
 source/adapters/level_zero/memory.cpp | 41 ++++++++++++++++++++-------
 source/adapters/level_zero/memory.hpp |  3 ++
 2 files changed, 34 insertions(+), 10 deletions(-)

diff --git a/source/adapters/level_zero/memory.cpp b/source/adapters/level_zero/memory.cpp
index 69edf83a78..de65d062aa 100644
--- a/source/adapters/level_zero/memory.cpp
+++ b/source/adapters/level_zero/memory.cpp
@@ -1682,6 +1682,7 @@ ur_result_t urMemRelease(
     Buffer->free();
   }
   delete Mem;
+  Mem = nullptr;
 
   return UR_RESULT_SUCCESS;
 }
@@ -2024,18 +2025,37 @@ ur_result_t _ur_buffer::getZeHandle(char *&ZeHandle, access_mode_t AccessMode,
 
   auto &Allocation = Allocations[Device];
 
+  if (this->isFreed) {
+    logger::debug("getZeHandle(ur_device_handle{}) buffer already released, no "
+                  "valid handles.",
+                  (void *)Device);
+    ZeHandle = nullptr;
+    return UR_RESULT_ERROR_INVALID_NULL_HANDLE;
+  }
+
   // Sub-buffers don't maintain own allocations but rely on parent buffer.
   if (SubBuffer) {
-    UR_CALL(SubBuffer->Parent->getZeHandle(ZeHandle, AccessMode, Device));
-    ZeHandle += SubBuffer->Origin;
-    // Still store the allocation info in the PI sub-buffer for
-    // getZeHandlePtr to work. At least zeKernelSetArgumentValue needs to
-    // be given a pointer to the allocation handle rather than its value.
-    //
-    Allocation.ZeHandle = ZeHandle;
-    Allocation.ReleaseAction = allocation_t::keep;
-    LastDeviceWithValidAllocation = Device;
-    return UR_RESULT_SUCCESS;
+    // Verify that the Parent Buffer is still valid or if it has been freed.
+    if (SubBuffer->Parent && !SubBuffer->Parent->isFreed) {
+      UR_CALL(SubBuffer->Parent->getZeHandle(ZeHandle, AccessMode, Device));
+      ZeHandle += SubBuffer->Origin;
+      // Still store the allocation info in the PI sub-buffer for
+      // getZeHandlePtr to work. At least zeKernelSetArgumentValue needs to
+      // be given a pointer to the allocation handle rather than its value.
+      //
+      Allocation.ZeHandle = ZeHandle;
+      Allocation.ReleaseAction = allocation_t::keep;
+      LastDeviceWithValidAllocation = Device;
+      return UR_RESULT_SUCCESS;
+    } else {
+      // Return a null pointer handle and an error if the parent buffer is
+      // already gone.
+      logger::debug("getZeHandle(ur_device_handle{}) SubBuffer's parent "
+                    "already released, no valid handles.",
+                    (void *)Device);
+      ZeHandle = nullptr;
+      return UR_RESULT_ERROR_INVALID_NULL_HANDLE;
+    }
   }
 
   // First handle case where the buffer is represented by only
@@ -2266,6 +2286,7 @@ ur_result_t _ur_buffer::free() {
       die("_ur_buffer::free(): Unhandled release action");
     }
     ZeHandle = nullptr; // don't leave hanging pointers
+    this->isFreed = true;
   }
   return UR_RESULT_SUCCESS;
 }
diff --git a/source/adapters/level_zero/memory.hpp b/source/adapters/level_zero/memory.hpp
index 71d102e9dd..b39dd58eb8 100644
--- a/source/adapters/level_zero/memory.hpp
+++ b/source/adapters/level_zero/memory.hpp
@@ -129,6 +129,9 @@ struct _ur_buffer final : ur_mem_handle_t_ {
   // Frees all allocations made for the buffer.
   ur_result_t free();
 
+  // Tracks if this buffer is freed already or should be considered valid.
+  bool isFreed{false};
+
   // Information about a single allocation representing this buffer.
   struct allocation_t {
     // Level Zero memory handle is really just a naked pointer.