From 2712b3276bd31c1535a97b696934654005c1cc3a Mon Sep 17 00:00:00 2001
From: PixelyIon <pixelyion@protonmail.com>
Date: Thu, 2 Jun 2022 22:14:22 +0530
Subject: [PATCH] Fix incorrect `VkBufferImageCopy` offset calculations

The `VkBufferImageCopy` offset calculations were wrong inside `CopyIntoStagingBuffer` as it multiplied the mip level's linear size by `levelCount` rather than `layerCount`. This led to substantial UB in games which called this function as it led to an overflow and resulted in writing to other areas of the buffer which caused major issues such as vertex/index buffer corruption and corresponding graphical glitches alongside likely being the cause of some crashes.
---
 .../main/cpp/skyline/gpu/texture/texture.cpp  | 70 +++++++------------
 .../main/cpp/skyline/gpu/texture/texture.h    |  5 ++
 2 files changed, 30 insertions(+), 45 deletions(-)

diff --git a/app/src/main/cpp/skyline/gpu/texture/texture.cpp b/app/src/main/cpp/skyline/gpu/texture/texture.cpp
index cf79623f..217edc56 100644
--- a/app/src/main/cpp/skyline/gpu/texture/texture.cpp
+++ b/app/src/main/cpp/skyline/gpu/texture/texture.cpp
@@ -300,25 +300,9 @@ namespace skyline::gpu {
         return stagingBuffer;
     }
 
-    void Texture::CopyFromStagingBuffer(const vk::raii::CommandBuffer &commandBuffer, const std::shared_ptr<memory::StagingBuffer> &stagingBuffer) {
-        auto image{GetBacking()};
-        if (layout == vk::ImageLayout::eUndefined)
-            commandBuffer.pipelineBarrier(vk::PipelineStageFlagBits::eHost, vk::PipelineStageFlagBits::eTransfer, {}, {}, {}, vk::ImageMemoryBarrier{
-                .image = image,
-                .srcAccessMask = vk::AccessFlagBits::eMemoryRead,
-                .dstAccessMask = vk::AccessFlagBits::eTransferWrite,
-                .oldLayout = std::exchange(layout, vk::ImageLayout::eGeneral),
-                .newLayout = vk::ImageLayout::eGeneral,
-                .srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED,
-                .dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED,
-                .subresourceRange = {
-                    .aspectMask = format->vkAspect,
-                    .levelCount = levelCount,
-                    .layerCount = layerCount,
-                },
-            });
+    boost::container::small_vector<vk::BufferImageCopy, 10> Texture::GetBufferImageCopies() {
+        boost::container::small_vector<vk::BufferImageCopy, 10> bufferImageCopies;
 
-        std::vector<vk::BufferImageCopy> bufferImageCopies;
         auto pushBufferImageCopyWithAspect{[&](vk::ImageAspectFlagBits aspect) {
             vk::DeviceSize bufferOffset{};
             u32 mipLevel{};
@@ -345,6 +329,28 @@ namespace skyline::gpu {
         if (format->vkAspect & vk::ImageAspectFlagBits::eStencil)
             pushBufferImageCopyWithAspect(vk::ImageAspectFlagBits::eStencil);
 
+        return bufferImageCopies;
+    }
+
+    void Texture::CopyFromStagingBuffer(const vk::raii::CommandBuffer &commandBuffer, const std::shared_ptr<memory::StagingBuffer> &stagingBuffer) {
+        auto image{GetBacking()};
+        if (layout == vk::ImageLayout::eUndefined)
+            commandBuffer.pipelineBarrier(vk::PipelineStageFlagBits::eHost, vk::PipelineStageFlagBits::eTransfer, {}, {}, {}, vk::ImageMemoryBarrier{
+                .image = image,
+                .srcAccessMask = vk::AccessFlagBits::eMemoryRead,
+                .dstAccessMask = vk::AccessFlagBits::eTransferWrite,
+                .oldLayout = std::exchange(layout, vk::ImageLayout::eGeneral),
+                .newLayout = vk::ImageLayout::eGeneral,
+                .srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED,
+                .dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED,
+                .subresourceRange = {
+                    .aspectMask = format->vkAspect,
+                    .levelCount = levelCount,
+                    .layerCount = layerCount,
+                },
+            });
+
+        auto bufferImageCopies{GetBufferImageCopies()};
         commandBuffer.copyBufferToImage(stagingBuffer->vkBuffer, image, layout, vk::ArrayProxy(static_cast<u32>(bufferImageCopies.size()), bufferImageCopies.data()));
     }
 
@@ -365,33 +371,7 @@ namespace skyline::gpu {
             },
         });
 
-        boost::container::small_vector<vk::BufferImageCopy, 10> bufferImageCopies;
-        auto pushBufferImageCopyWithAspect{[&](vk::ImageAspectFlagBits aspect) {
-            vk::DeviceSize bufferOffset{};
-            u32 mipLevel{};
-            for (auto &level : mipLayouts) {
-                bufferImageCopies.emplace_back(
-                    vk::BufferImageCopy{
-                        .bufferOffset = bufferOffset,
-                        .imageSubresource = {
-                            .aspectMask = aspect,
-                            .mipLevel = mipLevel++,
-                            .layerCount = layerCount,
-                        },
-                        .imageExtent = level.dimensions,
-                    }
-                );
-                bufferOffset += level.targetLinearSize * levelCount;
-            }
-        }};
-
-        if (format->vkAspect & vk::ImageAspectFlagBits::eColor)
-            pushBufferImageCopyWithAspect(vk::ImageAspectFlagBits::eColor);
-        if (format->vkAspect & vk::ImageAspectFlagBits::eDepth)
-            pushBufferImageCopyWithAspect(vk::ImageAspectFlagBits::eDepth);
-        if (format->vkAspect & vk::ImageAspectFlagBits::eStencil)
-            pushBufferImageCopyWithAspect(vk::ImageAspectFlagBits::eStencil);
-
+        auto bufferImageCopies{GetBufferImageCopies()};
         commandBuffer.copyImageToBuffer(image, layout, stagingBuffer->vkBuffer, vk::ArrayProxy(static_cast<u32>(bufferImageCopies.size()), bufferImageCopies.data()));
 
         commandBuffer.pipelineBarrier(vk::PipelineStageFlagBits::eTransfer, vk::PipelineStageFlagBits::eHost, {}, {}, vk::BufferMemoryBarrier{
diff --git a/app/src/main/cpp/skyline/gpu/texture/texture.h b/app/src/main/cpp/skyline/gpu/texture/texture.h
index 3c02a24b..c8cdc95d 100644
--- a/app/src/main/cpp/skyline/gpu/texture/texture.h
+++ b/app/src/main/cpp/skyline/gpu/texture/texture.h
@@ -414,6 +414,11 @@ namespace skyline::gpu {
             ~TextureBufferCopy();
         };
 
+        /**
+         * @return A vector of all the buffer image copies that need to be done for every aspect of every level of every layer of the texture
+         */
+        boost::container::small_vector<vk::BufferImageCopy, 10> GetBufferImageCopies();
+
       public:
         std::weak_ptr<FenceCycle> cycle; //!< A fence cycle for when any host operation mutating the texture has completed, it must be waited on prior to any mutations to the backing
         std::optional<GuestTexture> guest;