diff --git a/aurora-main/lib/webgpu/gpu_cache.cpp b/aurora-main/lib/webgpu/gpu_cache.cpp index 35e9642..e494bc7 100644 --- a/aurora-main/lib/webgpu/gpu_cache.cpp +++ b/aurora-main/lib/webgpu/gpu_cache.cpp @@ -6,6 +6,7 @@ #include #include #include +#include #include #include #include @@ -313,8 +314,33 @@ static size_t load_from_database(const XXH128_hash_t& keyHash, void* value, size if (ret == SQLITE_ROW) { // Hit const auto foundPtr = sqlite3_column_blob(load_stmt, 0); - foundSize = sqlite3_column_int64(load_stmt, 1); - const bool compressed = sqlite3_column_int(load_stmt, 2) != 0; + const auto declaredSize = sqlite3_column_int64(load_stmt, 1); + const auto storedSize = sqlite3_column_bytes(load_stmt, 0); + const auto compression = sqlite3_column_int(load_stmt, 2); + const bool compressed = compression == 1; + // Dawn asks for the size before allocating its destination. Validate here, + // not only during the copy: corrupt metadata must become a cache miss. + bool valid = declaredSize > 0 && + static_cast(declaredSize) <= std::numeric_limits::max() && + foundPtr != nullptr && storedSize > 0 && (compression == 0 || compression == 1); + if (valid && compressed) { +#if defined(AURORA_CACHE_USE_ZSTD) + // Our writer uses ZSTD_compress, which records the original content size. + const auto frameSize = ZSTD_getFrameContentSize(foundPtr, static_cast(storedSize)); + valid = frameSize != ZSTD_CONTENTSIZE_ERROR && frameSize != ZSTD_CONTENTSIZE_UNKNOWN && + frameSize == static_cast(declaredSize); +#else + valid = false; +#endif + } else if (valid) { + valid = declaredSize == storedSize; + } + if (!valid) { + Log.error("Ignoring cache entry with inconsistent size or compression metadata"); + check(sqlite3_reset(load_stmt)); + return 0; + } + foundSize = static_cast(declaredSize); if (value == nullptr) { g_hits.fetch_add(1, std::memory_order_relaxed); } else {