From d98d152d2359151f5b6c2ce0ddda35bf96437f07 Mon Sep 17 00:00:00 2001
From: kd-11 <karokidii@gmail.com>
Date: Sun, 27 Mar 2022 20:25:34 +0300
Subject: [PATCH] rsx: Fix leaking surface cache refs from texture cache

- Lock surfaces in use by texture cache to prevent complete deletion
- Remove discarded surfaces from the reprotect cache to avoid uaf
---
 rpcs3/Emu/RSX/Common/texture_cache.h       |  6 ++++++
 rpcs3/Emu/RSX/Common/texture_cache_utils.h | 24 ++++++++++++++++++++++
 rpcs3/Emu/RSX/GL/GLTextureCache.cpp        |  6 ------
 rpcs3/Emu/RSX/GL/GLTextureCache.h          |  5 +++++
 rpcs3/Emu/RSX/VK/VKTextureCache.h          | 11 +++++-----
 5 files changed, 40 insertions(+), 12 deletions(-)

diff --git a/rpcs3/Emu/RSX/Common/texture_cache.h b/rpcs3/Emu/RSX/Common/texture_cache.h
index 0e7f1ad736..cea4e0e1e9 100644
--- a/rpcs3/Emu/RSX/Common/texture_cache.h
+++ b/rpcs3/Emu/RSX/Common/texture_cache.h
@@ -874,6 +874,12 @@ namespace rsx
 							tex.set_dirty(true);
 							result.invalidate_samplers = true;
 						}
+
+						if (tex.is_dirty() && tex.get_context() == rsx::texture_upload_context::framebuffer_storage)
+						{
+							// Make sure the region is not going to get immediately reprotected
+							m_flush_always_cache.erase(tex.get_section_range());
+						}
 					}
 				}
 
diff --git a/rpcs3/Emu/RSX/Common/texture_cache_utils.h b/rpcs3/Emu/RSX/Common/texture_cache_utils.h
index 8a3d1cd512..644f32b6a0 100644
--- a/rpcs3/Emu/RSX/Common/texture_cache_utils.h
+++ b/rpcs3/Emu/RSX/Common/texture_cache_utils.h
@@ -1353,9 +1353,28 @@ namespace rsx
 
 				set_dirty(false);
 			}
+
+			if (context == rsx::texture_upload_context::framebuffer_storage)
+			{
+				// Lock, unlock
+				if (prot == utils::protection::no && old_prot != utils::protection::no)
+				{
+					// Locked memory. We have to take ownership of the object in the surface cache as well
+					auto surface = derived()->get_render_target();
+					surface->add_ref();
+				}
+				else if (old_prot == utils::protection::no && prot != utils::protection::no)
+				{
+					// Release the surface, the cache can remove it if needed
+					ensure(prot == utils::protection::rw);
+					auto surface = derived()->get_render_target();
+					surface->release();
+				}
+			}
 		}
 
 	public:
+
 		inline void protect(utils::protection prot)
 		{
 			utils::protection old_prot = get_protection();
@@ -1442,6 +1461,11 @@ namespace rsx
 			}
 
 			flush_exclusions.clear();
+
+			if (context == rsx::texture_upload_context::framebuffer_storage)
+			{
+				derived()->get_render_target()->sync_tag();
+			}
 		}
 
 		void on_speculative_flush()
diff --git a/rpcs3/Emu/RSX/GL/GLTextureCache.cpp b/rpcs3/Emu/RSX/GL/GLTextureCache.cpp
index bf9f1e90ec..4550d1c37a 100644
--- a/rpcs3/Emu/RSX/GL/GLTextureCache.cpp
+++ b/rpcs3/Emu/RSX/GL/GLTextureCache.cpp
@@ -81,12 +81,6 @@ namespace gl
 				rsx_log.error("Unexpected swizzled texture format 0x%x", static_cast<u32>(format));
 			}
 		}
-
-		if (context == rsx::texture_upload_context::framebuffer_storage)
-		{
-			// Update memory tag
-			static_cast<gl::render_target*>(vram_texture)->sync_tag();
-		}
 	}
 
 	gl::texture_view* texture_cache::create_temporary_subresource_impl(gl::command_context& cmd, gl::texture* src, GLenum sized_internal_fmt, GLenum dst_type,
diff --git a/rpcs3/Emu/RSX/GL/GLTextureCache.h b/rpcs3/Emu/RSX/GL/GLTextureCache.h
index 5f8c2d37ae..7de17a47dd 100644
--- a/rpcs3/Emu/RSX/GL/GLTextureCache.h
+++ b/rpcs3/Emu/RSX/GL/GLTextureCache.h
@@ -406,6 +406,11 @@ namespace gl
 			return managed_texture.get();
 		}
 
+		gl::render_target* get_render_target()
+		{
+			return gl::as_rtt(vram_texture);
+		}
+
 		gl::texture_view* get_raw_view()
 		{
 			return vram_texture->get_view(0xAAE4, rsx::default_remap_vector);
diff --git a/rpcs3/Emu/RSX/VK/VKTextureCache.h b/rpcs3/Emu/RSX/VK/VKTextureCache.h
index 3026579334..37a79ca320 100644
--- a/rpcs3/Emu/RSX/VK/VKTextureCache.h
+++ b/rpcs3/Emu/RSX/VK/VKTextureCache.h
@@ -151,6 +151,11 @@ namespace vk
 			return managed_texture;
 		}
 
+		vk::render_target* get_render_target()
+		{
+			return vk::as_rtt(vram_texture);
+		}
+
 		VkFormat get_format() const
 		{
 			if (context == rsx::texture_upload_context::dma)
@@ -293,12 +298,6 @@ namespace vk
 					rsx_log.error("Unexpected swizzled texture format 0x%x", gcm_format);
 				}
 			}
-
-			if (context == rsx::texture_upload_context::framebuffer_storage)
-			{
-				// Update memory tag
-				static_cast<vk::render_target*>(vram_texture)->sync_tag();
-			}
 		}
 
 		void* map_synchronized(u32, u32)