Vulkan: make mappableBuffersForEagerTransition ref buffers

There are some scenarios where buffers can be used in pending
commands that are not retained by a command buffer.

They must be retained in the set of mappable buffers for eager
transition to prevent a use-after-free violation.

Fixed: chromium:1421170
Change-Id: I452d80b2513a7726a003d44e2a7850292d798bb1
Reviewed-on: https://dawn-review.googlesource.com/c/dawn/+/122580
Auto-Submit: Austin Eng <enga@chromium.org>
Kokoro: Kokoro <noreply+kokoro@google.com>
Commit-Queue: Austin Eng <enga@chromium.org>
Reviewed-by: Loko Kung <lokokung@google.com>
Reviewed-by: Corentin Wallez <cwallez@chromium.org>
Reviewed-by: Peng Huang <penghuang@chromium.org>

src/dawn/common/RefBase.h

@@ -110,6 +110,8 @@ class RefBase {
     const T operator->() const { return mValue; }
     T operator->() { return mValue; }
 
+    bool operator<(const RefBase<T, Traits>& other) const { return mValue < other.mValue; }
+
     // Smart pointer methods.
     const T& Get() const { return mValue; }
     T& Get() { return mValue; }

src/dawn/native/vulkan/BufferVk.cpp

@@ -416,9 +416,8 @@ bool Buffer::EnsureDataInitializedAsDestination(CommandRecordingContext* recordi
 // static
 void Buffer::TransitionMappableBuffersEagerly(const VulkanFunctions& fn,
                                               CommandRecordingContext* recordingContext,
-                                              std::set<Buffer*> buffers) {
+                                              const std::set<Ref<Buffer>>& buffers) {
     ASSERT(!buffers.empty());
-    ASSERT(recordingContext->mappableBuffersForEagerTransition.empty());
 
     VkPipelineStageFlags srcStages = 0;
     VkPipelineStageFlags dstStages = 0;
@@ -426,7 +425,8 @@ void Buffer::TransitionMappableBuffersEagerly(const VulkanFunctions& fn,
     std::vector<VkBufferMemoryBarrier> barriers;
     barriers.reserve(buffers.size());
 
-    for (Buffer* buffer : buffers) {
+    size_t originalBufferCount = buffers.size();
+    for (const Ref<Buffer>& buffer : buffers) {
         wgpu::BufferUsage mapUsage = buffer->GetUsage() & kMapUsages;
         ASSERT(mapUsage == wgpu::BufferUsage::MapRead || mapUsage == wgpu::BufferUsage::MapWrite);
         VkBufferMemoryBarrier barrier;
@@ -435,9 +435,9 @@ void Buffer::TransitionMappableBuffersEagerly(const VulkanFunctions& fn,
                                                     &srcStages, &dstStages)) {
             barriers.push_back(barrier);
         }
-        // TrackUsageAndGetResourceBarrier() should not modify recordingContext for map usages.
-        ASSERT(recordingContext->mappableBuffersForEagerTransition.empty());
     }
+    // TrackUsageAndGetResourceBarrier() should not modify recordingContext for map usages.
+    ASSERT(buffers.size() == originalBufferCount);
 
     if (barriers.empty()) {
         return;

src/dawn/native/vulkan/BufferVk.h

@@ -58,7 +58,7 @@ class Buffer final : public BufferBase {
 
     static void TransitionMappableBuffersEagerly(const VulkanFunctions& fn,
                                                  CommandRecordingContext* recordingContext,
-                                                 std::set<Buffer*> buffers);
+                                                 const std::set<Ref<Buffer>>& buffers);
 
   private:
     ~Buffer() override;

src/dawn/native/vulkan/CommandRecordingContext.h

@@ -50,7 +50,7 @@ struct CommandRecordingContext {
 
     // Mappable buffers which will be eagerly transitioned to usage MapRead or MapWrite after
     // VkSubmit.
-    std::set<Buffer*> mappableBuffersForEagerTransition;
+    std::set<Ref<Buffer>> mappableBuffersForEagerTransition;
 
     // For Device state tracking only.
     VkCommandPool commandPool = VK_NULL_HANDLE;

src/dawn/native/vulkan/DeviceVk.cpp

@@ -310,7 +310,7 @@ MaybeError Device::SubmitPendingCommands() {
     if (!mRecordingContext.mappableBuffersForEagerTransition.empty()) {
         // Transition mappable buffers back to map usages with the submit.
         Buffer::TransitionMappableBuffersEagerly(
-            fn, &mRecordingContext, std::move(mRecordingContext.mappableBuffersForEagerTransition));
+            fn, &mRecordingContext, mRecordingContext.mappableBuffersForEagerTransition);
     }
 
     ScopedSignalSemaphore scopedSignalSemaphore(this, VK_NULL_HANDLE);

src/dawn/tests/end2end/BufferTests.cpp

@@ -534,6 +534,24 @@ TEST_P(BufferMappingTests, MapWrite_ZeroSizedTwice) {
     MapAsyncAndWait(buffer, wgpu::MapMode::Write, 0, wgpu::kWholeMapSize);
 }
 
+// Regression test for crbug.com/1421170 where dropping a buffer which needs
+// padding bytes initialization resulted in a use-after-free.
+TEST_P(BufferMappingTests, RegressChromium1421170) {
+    // Create a mappable buffer of size 7. It will be internally
+    // aligned such that the padding bytes need to be zero initialized.
+    wgpu::BufferDescriptor descriptor;
+    descriptor.size = 7;
+    descriptor.usage = wgpu::BufferUsage::MapWrite;
+    descriptor.mappedAtCreation = false;
+    wgpu::Buffer buffer = device.CreateBuffer(&descriptor);
+
+    // Drop the buffer. The pending commands to zero initialize the
+    // padding bytes should stay valid.
+    buffer = nullptr;
+    // Flush pending commands.
+    device.Tick();
+}
+
 DAWN_INSTANTIATE_TEST(BufferMappingTests,
                       D3D12Backend(),
                       MetalBackend(),