[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

@@ -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;
 }

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.