Merge pull request #3785 from MrMha/fix/shader-cache-read-allocation

fix(shader): handle cache file allocation failures

CMakeLists.txt

@@ -48,6 +48,12 @@ elseif(NOT APPLE)
 endif()
 
 if(BUILD_TESTING)
+    add_test_executable(shader_cache_file_read_tests
+        ${CMAKE_CURRENT_SOURCE_DIR}/core/shader/recompiler/tests/CacheFileReadTests.cpp)
+    target_include_directories(shader_cache_file_read_tests PRIVATE
+        ${CMAKE_CURRENT_SOURCE_DIR}/core/shader/recompiler)
+    add_test(NAME shader_cache_file_read COMMAND shader_cache_file_read_tests)
+    set_tests_properties(shader_cache_file_read PROPERTIES TIMEOUT 20)
     find_package(Python3 QUIET COMPONENTS Interpreter)
     if(Python3_Interpreter_FOUND)
         add_test(NAME pr_overlap_encoding COMMAND ${Python3_EXECUTABLE}

core/shader/recompiler/ShaderCacheDirectory.hpp

@@ -4,11 +4,15 @@
 #include <atomic>
 #include <chrono>
 #include <cstddef>
+#include <cstdint>
 #include <cstdlib>
 #include <cstring>
 #include <filesystem>
 #include <fstream>
 #include <functional>
+#include <limits>
+#include <new>
+#include <stdexcept>
 #include <span>
 #include <string>
 #include <system_error>
@@ -89,17 +93,27 @@ inline bool WriteFileAtomically(const std::filesystem::path& path, std::span<con
 
 inline bool ReadWholeFile(const std::filesystem::path& path, std::vector<std::byte>& bytes) {
     bytes.clear();
-    std::ifstream file(path, std::ios::binary | std::ios::ate);
-    if (!file) return false;
-    const auto size = static_cast<std::streamoff>(file.tellg());
-    if (size < 0) return false;
-    file.seekg(0, std::ios::beg);
-    bytes.resize(static_cast<std::size_t>(size));
-    if (!file.read(reinterpret_cast<char*>(bytes.data()), static_cast<std::streamsize>(bytes.size()))) {
+    try {
+        std::ifstream file(path, std::ios::binary | std::ios::ate);
+        if (!file) return false;
+        const auto size = static_cast<std::streamoff>(file.tellg());
+        if (size < 0 || static_cast<std::uintmax_t>(size) > bytes.max_size() ||
+            static_cast<std::uintmax_t>(size) > static_cast<std::uintmax_t>(std::numeric_limits<std::streamsize>::max()))
+            return false;
+        if (!file.seekg(0, std::ios::beg)) return false;
+        bytes.resize(static_cast<std::size_t>(size));
+        if (!file.read(reinterpret_cast<char*>(bytes.data()), static_cast<std::streamsize>(bytes.size()))) {
+            bytes.clear();
+            return false;
+        }
+        return true;
+    } catch (const std::bad_alloc&) {
+        bytes.clear();
+        return false;
+    } catch (const std::length_error&) {
         bytes.clear();
         return false;
     }
-    return true;
 }
 
 inline std::uint64_t HashBytes(std::span<const std::byte> bytes, std::uint64_t seed = 0) {

core/shader/recompiler/tests/CacheFileReadTests.cpp

@@ -0,0 +1,44 @@
+#include "ShaderCacheDirectory.hpp"
+#include <iostream>
+#include <new>
+#include <stdexcept>
+
+static bool FailLargeAllocation = false;
+
+void* operator new(std::size_t size) {
+    if (FailLargeAllocation && size > 1024 * 1024) throw std::bad_alloc();
+    if (auto* memory = std::malloc(size == 0 ? 1 : size)) return memory;
+    throw std::bad_alloc();
+}
+void operator delete(void* memory) noexcept { std::free(memory); }
+void operator delete(void* memory, std::size_t) noexcept { std::free(memory); }
+void* operator new[](std::size_t size) { return ::operator new(size); }
+void operator delete[](void* memory) noexcept { ::operator delete(memory); }
+void operator delete[](void* memory, std::size_t) noexcept { ::operator delete(memory); }
+
+static void Require(bool ok) { if (!ok) throw std::runtime_error("Cache file read failed"); }
+
+int main() {
+    const auto path = std::filesystem::temp_directory_path() / ("anyps5-cache-read-" +
+        std::to_string(std::chrono::steady_clock::now().time_since_epoch().count()));
+    int status = 0;
+    try {
+        std::vector<std::byte> bytes{std::byte{1}};
+        Require(!ShaderRecompiler::ReadWholeFile(path, bytes) && bytes.empty());
+        { std::ofstream file(path, std::ios::binary); }
+        Require(ShaderRecompiler::ReadWholeFile(path, bytes) && bytes.empty());
+        const std::vector<char> payload(2 * 1024 * 1024, 'S');
+        { std::ofstream file(path, std::ios::binary | std::ios::trunc); file.write(payload.data(), payload.size()); Require(static_cast<bool>(file)); }
+        Require(ShaderRecompiler::ReadWholeFile(path, bytes) && bytes.size() == payload.size());
+        Require(bytes.front() == std::byte{'S'} && bytes.back() == std::byte{'S'});
+        std::vector<std::byte> failed{std::byte{1}};
+        FailLargeAllocation = true;
+        const auto read = ShaderRecompiler::ReadWholeFile(path, failed);
+        FailLargeAllocation = false;
+        Require(!read && failed.empty());
+        std::cout << "5 cache file read cases passed\n";
+    } catch (const std::exception& error) { FailLargeAllocation = false; std::cerr << error.what() << '\n'; status = 1; }
+    std::error_code error;
+    std::filesystem::remove(path, error);
+    return status;
+}