fix(agc): stamp the watched blocks of a store into partly watched memory

core/libs/prx/libSceAgcDriver/Execution/src/GuestMemory.cpp

@@ -918,6 +918,32 @@ struct WriteTracker {
         }
     }
 
+    bool blockPartWatched(std::uint64_t first, std::uint64_t stop) const {
+#ifdef _WIN32
+        return covers(first, 1) || covers(stop - 1, 1);
+#else
+        return GuestWriteWatch::GuestWriteWatchOverlaps_nid_postfix(static_cast<std::uintptr_t>(first), static_cast<std::size_t>(stop - first));
+#endif
+    }
+
+    void stampWatchedPart(std::uint64_t address, std::uint64_t end, std::uint32_t stampGeneration) {
+#ifdef _WIN32
+        address = std::max<std::uint64_t>(address, base);
+        end = std::min<std::uint64_t>(end, base + size);
+        if (address >= end) return;
+#else
+        end = std::min<std::uint64_t>(end, LeafCount * LeafBlocks * WriteBlockBytes);
+        if (address >= end || !GuestWriteWatch::GuestWriteWatchOverlaps_nid_postfix(static_cast<std::uintptr_t>(address), static_cast<std::size_t>(end - address))) return;
+#endif
+        for (auto block = blockOf(address); block <= blockOf(end - 1); ++block) {
+            const auto first = std::max(address, blockBegin(block));
+            const auto stop = std::min<std::uint64_t>(end, blockBegin(block) + WriteBlockBytes);
+            if (!blockPartWatched(first, stop)) continue;
+            stamp(block, stampGeneration, StampKind::Driver);
+            noteDriverStore(block, address, end, stampGeneration);
+        }
+    }
+
     void noteDriverStore(std::uint64_t block, std::uint64_t address, std::uint64_t end, std::uint32_t stampGeneration) {
         const auto begin = blockBegin(block);
         const auto from = static_cast<std::uint32_t>(std::max(address, begin) - begin);
@@ -1447,6 +1473,7 @@ std::uint64_t storeOwn(std::uint64_t address, std::size_t bytes, const std::func
         ++tracker.generation;
         if (!tracker.watched || !tracker.covers(stored.first, static_cast<std::size_t>(stored.second - stored.first))) {
             tracker.noteUnstampedStore(stored.first, stored.second);
+            if (tracker.watched) tracker.stampWatchedPart(stored.first, stored.second, tracker.generation);
             return 0;
         }
         for (auto block = tracker.blockOf(stored.first); block <= tracker.blockOf(stored.second - 1); ++block) {
@@ -1492,6 +1519,7 @@ std::uint64_t MarkWritten(std::uint64_t address, std::size_t bytes) {
     ++tracker.generation;
     if (!tracker.watched || !tracker.covers(address, bytes)) {
         tracker.noteUnstampedStore(address, address + bytes);
+        if (tracker.watched) tracker.stampWatchedPart(address, address + bytes, tracker.generation);
         return 0;
     }
     {

core/libs/prx/libSceAgcDriver/tests/WriteTracking.cpp

@@ -9,6 +9,7 @@
 #include <windows.h>
 #else
 #include <sys/mman.h>
+#include <unistd.h>
 #endif
 #include <array>
 #include <cstdint>
@@ -280,6 +281,75 @@ void CheckNewestStamp() {
     Require(NewestStamp(base + 3 * Block - 8, Block + 8) > partial, "a CPU store into the watched part of a partly watched range did not advance its stamp");
 }
 
+void CheckPartlyWatchedStore() {
+    void* memory = AllocateWatched(3 * Block);
+    const auto base = reinterpret_cast<std::uint64_t>(memory);
+    const auto stored = base + Block;
+    const auto unwatched = base + 2 * Block;
+    std::memset(memory, 0x11, 3 * Block);
+    BumpCollectEpoch();
+    Unwatch(unwatched, Block);
+    const auto before = CollectWrites(base, 2 * Block);
+    Require(before != 0 && UnchangedSince(base, 2 * Block, before), "the watched blocks were not collected");
+    Require(MarkWritten(unwatched - 64, 128) == 0, "a driver store into partly watched memory is trusted");
+    Require(!UnchangedSince(stored, Block, before), "a driver store into partly watched memory left its watched block unchanged");
+    Require(StoredOver(unwatched - 64, 64, before), "a driver store into partly watched memory is not seen over its watched bytes");
+    Require(!StoredOver(stored, Block - 64, before), "a driver store into partly watched memory is seen over watched bytes it did not store");
+    Require(UnchangedSince(base, Block, before), "a driver store into partly watched memory stamped a block it did not store into");
+    const auto beforeOwn = TrackerGeneration();
+    Require(StoreOwnBytes(unwatched - 64, 128, [&] { std::memset(reinterpret_cast<void*>(unwatched - 64), 0x22, 128); }) == 0, "an own store into partly watched memory is trusted");
+    Require(!UnchangedSince(stored, Block, beforeOwn), "an own store into partly watched memory left its watched block unchanged");
+    Require(UnchangedSince(base, Block, beforeOwn), "an own store into partly watched memory stamped a block it did not store into");
+#ifndef _WIN32
+    void* split = AllocateWatched(Block);
+    const auto splitBase = reinterpret_cast<std::uint64_t>(split);
+    const auto watchedHalf = splitBase + Block / 2;
+    std::memset(split, 0x11, Block);
+    Unwatch(splitBase, Block / 2);
+    const auto beforeSplit = CollectWrites(watchedHalf, Block / 2);
+    Require(beforeSplit != 0 && UnchangedSince(watchedHalf, Block / 2, beforeSplit), "the watched half of a split block was not collected");
+    Require(MarkWritten(splitBase + 64, 128) == 0 && UnchangedSince(watchedHalf, Block / 2, beforeSplit), "a driver store into the unwatched half of a split block stamped its watched half");
+    Require(MarkWritten(watchedHalf - 64, 128) == 0, "a driver store across the halves of a split block is trusted");
+    Require(!UnchangedSince(watchedHalf, Block / 2, beforeSplit), "a driver store across the halves of a split block left its watched half unchanged");
+#endif
+}
+
+#ifndef _WIN32
+void CheckWatchOverlaps() {
+    const auto page = static_cast<std::size_t>(sysconf(_SC_PAGESIZE));
+    void* memory = AllocateWatched(Block);
+    const auto base = reinterpret_cast<std::uintptr_t>(memory);
+    Require(GuestWriteWatch::GuestWriteWatchUnregister_nid_postfix(memory, page), "cannot unwatch the block's first page");
+    Require(GuestWriteWatch::GuestWriteWatchUnregister_nid_postfix(reinterpret_cast<void*>(base + 3 * page), Block - 3 * page), "cannot unwatch the block's last pages");
+    Require(!GuestWriteWatch::GuestWriteWatchCovers_nid_postfix(base, Block), "a block with unwatched pages is covered");
+    Require(GuestWriteWatch::GuestWriteWatchOverlaps_nid_postfix(base, Block), "a block with watched pages inside does not overlap the watch");
+    Require(GuestWriteWatch::GuestWriteWatchOverlaps_nid_postfix(base + page - 1, 2), "a range ending in a watched page does not overlap the watch");
+    Require(GuestWriteWatch::GuestWriteWatchOverlaps_nid_postfix(base + 3 * page - 1, 2), "a range starting in a watched page does not overlap the watch");
+    Require(!GuestWriteWatch::GuestWriteWatchOverlaps_nid_postfix(base, page), "an unwatched page before the watched ones overlaps the watch");
+    Require(!GuestWriteWatch::GuestWriteWatchOverlaps_nid_postfix(base + 3 * page, Block - 3 * page), "unwatched pages after the watched ones overlap the watch");
+    Require(!GuestWriteWatch::GuestWriteWatchOverlaps_nid_postfix(base + page, 0), "an empty range overlaps the watch");
+}
+
+void CheckWatchedMiddleStore() {
+    const auto page = static_cast<std::size_t>(sysconf(_SC_PAGESIZE));
+    void* memory = AllocateWatched(Block);
+    const auto base = reinterpret_cast<std::uint64_t>(memory);
+    const auto middle = base + page;
+    std::memset(memory, 0x11, Block);
+    Unwatch(base, page);
+    Unwatch(base + 3 * page, Block - 3 * page);
+    const auto before = CollectWrites(middle, 2 * page);
+    Require(before != 0 && UnchangedSince(middle, 2 * page, before), "the watched middle of a block was not collected");
+    Require(MarkWritten(base + 3 * page + 64, 128) == 0 && UnchangedSince(middle, 2 * page, before), "a driver store into an unwatched end of a block stamped its watched middle");
+    Require(MarkWritten(base + page / 2, 3 * page) == 0, "a driver store around the watched middle of a block is trusted");
+    Require(!UnchangedSince(middle, 2 * page, before), "a driver store around the watched middle of a block left it unchanged");
+    Require(StoredOver(middle, 2 * page, before), "a driver store around the watched middle of a block is not seen over it");
+    const auto beforeOwn = TrackerGeneration();
+    Require(StoreOwnBytes(base + page / 2, 3 * page, [&] { std::memset(reinterpret_cast<void*>(base + page / 2), 0x22, 3 * page); }) == 0, "an own store around the watched middle of a block is trusted");
+    Require(!UnchangedSince(middle, 2 * page, beforeOwn), "an own store around the watched middle of a block left it unchanged");
+}
+#endif
+
 #ifdef _WIN32
 void CheckDeferredSharedImports() {
     constexpr std::size_t bytes = 4 * Block;
@@ -730,6 +800,24 @@ void ProtectPage(std::uint64_t address, bool readable) {
     GuestAllocations::GuestAllocationsInvalidate_nid_postfix(static_cast<std::uintptr_t>(address), 4096);
 }
 
+void CheckStoreAcrossImports() {
+    void* memory = AllocateWatched(3 * Block);
+    const auto base = reinterpret_cast<std::uint64_t>(memory);
+    const auto boundary = base + 2 * Block;
+    std::memset(memory, 0x11, 3 * Block);
+    Require(ImportWatched(base, 2 * Block, [] { return true; }) && Watched(base, 2 * Block), "an import that stays watched left the watch");
+    Require(ImportWatched(boundary, Block, [] { return true; }, false) && !Watched(boundary, Block), "an import that leaves the watch stayed watched");
+    const auto before = CollectWrites(base, 2 * Block);
+    Require(before != 0 && UnchangedSince(base, 2 * Block, before), "the watched import was not collected");
+    Require(MarkWritten(boundary - 64, 128) == 0, "a GPU write across a watched import and an unwatched one is trusted");
+    Require(!UnchangedSince(base + Block, Block, before), "a GPU write across a watched import and an unwatched one left the watched import's last block unchanged");
+    Require(StoredOver(boundary - 64, 64, before) && !StoredOver(base + Block, Block - 64, before), "a GPU write across a watched import and an unwatched one is not seen over exactly its watched bytes");
+    Require(UnchangedSince(base, Block, before), "a GPU write across a watched import and an unwatched one stamped a block it did not write");
+    const auto beforeOwn = TrackerGeneration();
+    Require(StoreOwnBytes(boundary - 64, 128, [&] { std::memset(reinterpret_cast<void*>(boundary - 64), 0x22, 128); }) == 0, "an own store across a watched import and an unwatched one is trusted");
+    Require(!UnchangedSince(base + Block, Block, beforeOwn) && UnchangedSince(base, Block, beforeOwn), "an own store across a watched import and an unwatched one did not stamp exactly the watched block it wrote");
+}
+
 void CheckUnwatchedImport() {
     void* memory = AllocateWatched(3 * Block);
     const auto base = reinterpret_cast<std::uint64_t>(memory);
@@ -935,6 +1023,11 @@ int main() {
         CheckOwnStore();
         CheckUnwatch();
         CheckNewestStamp();
+        CheckPartlyWatchedStore();
+#ifndef _WIN32
+        CheckWatchOverlaps();
+        CheckWatchedMiddleStore();
+#endif
 #ifdef _WIN32
         CheckDeferredSharedImports();
         CheckPrivateMappingReuse();
@@ -946,6 +1039,7 @@ int main() {
         CheckImportWindow();
         CheckImportWindowZeroPages();
         CheckImportWindowOtherThread();
+        CheckStoreAcrossImports();
         CheckUnwatchedImport();
         CheckCommittedRanges();
     } catch (const std::exception& error) {

core/libs/prx/libc/include/GuestWriteWatch.hpp

@@ -12,6 +12,7 @@ bool GuestWriteWatchAvailable_nid_postfix();
 void GuestWriteWatchRegister_nid_postfix(const void* pointer, std::size_t bytes);
 bool GuestWriteWatchUnregister_nid_postfix(const void* pointer, std::size_t bytes);
 bool GuestWriteWatchCovers_nid_postfix(std::uintptr_t address, std::size_t bytes);
+bool GuestWriteWatchOverlaps_nid_postfix(std::uintptr_t address, std::size_t bytes);
 bool GuestWriteWatchCollect_nid_postfix(std::uintptr_t address, std::size_t bytes, void (*written)(void* context, std::uintptr_t begin, std::uintptr_t end), void* context);
 void GuestWriteWatchBeginHostWrite_nid_postfix(const void* pointer, std::size_t bytes);
 void GuestWriteWatchEndHostWrite_nid_postfix(const void* pointer, std::size_t bytes);

core/libs/prx/libc/src/GuestWriteWatch.cpp

@@ -90,6 +90,13 @@ class Watch {
         return covers(begin, end);
     }
 
+    bool Overlaps(std::uintptr_t begin, std::uintptr_t end) {
+        if (!Available()) return false;
+        std::shared_lock lock(_lock);
+        const auto next = _ranges.upper_bound(begin);
+        return (next != _ranges.begin() && std::prev(next)->second > begin) || (next != _ranges.end() && next->first < end);
+    }
+
     bool Collect(std::uintptr_t begin, std::uintptr_t end, void (*written)(void*, std::uintptr_t, std::uintptr_t), void* context) {
         if (!Available()) return false;
         begin &= ~(PageBytes - 1);
@@ -379,6 +386,13 @@ class Watch {
         return covers(begin, end);
     }
 
+    bool Overlaps(std::uintptr_t begin, std::uintptr_t end) {
+        if (!_available) return false;
+        const Lock lock(_lock);
+        const auto next = _ranges.upper_bound(begin);
+        return (next != _ranges.begin() && std::prev(next)->second.end > begin) || (next != _ranges.end() && next->first < end);
+    }
+
     enum class Take { All, Armed, Fresh };
 
     bool Collect(std::uintptr_t begin, std::uintptr_t end, void (*written)(void*, std::uintptr_t, std::uintptr_t), void* context, Take take = Take::All) {
@@ -660,6 +674,16 @@ bool GuestWriteWatchCovers_nid_postfix(std::uintptr_t address, std::size_t bytes
 #endif
 }
 
+bool GuestWriteWatchOverlaps_nid_postfix(std::uintptr_t address, std::size_t bytes) {
+#ifdef APS5_GUEST_WRITE_WATCH
+    return bytes != 0 && address + bytes > address && Watch::Get().Overlaps(address, address + bytes);
+#else
+    static_cast<void>(address);
+    static_cast<void>(bytes);
+    return false;
+#endif
+}
+
 bool GuestWriteWatchCollect_nid_postfix(std::uintptr_t address, std::size_t bytes, void (*written)(void* context, std::uintptr_t begin, std::uintptr_t end), void* context) {
 #ifdef APS5_GUEST_WRITE_WATCH
     if (bytes == 0 || address + bytes < address) return false;

core/libs/tests/WriteWatch.cpp

@@ -231,6 +231,21 @@ void CheckProtection() {
     Require(!GuestWriteWatchCovers_nid_postfix(region.Address(2), region.page), "an unregistered range is covered");
 }
 
+void CheckOverlaps() {
+    Region region(6);
+    Require(!GuestWriteWatchOverlaps_nid_postfix(region.Address(0), region.bytes), "an unregistered range overlaps the watch");
+    GuestWriteWatchRegister_nid_postfix(region.base, region.bytes);
+    Require(GuestWriteWatchUnregister_nid_postfix(region.base, region.page), "cannot unregister the first page");
+    Require(GuestWriteWatchUnregister_nid_postfix(region.base + 3 * region.page, 3 * region.page), "cannot unregister the last pages");
+    Require(!GuestWriteWatchCovers_nid_postfix(region.Address(0), region.bytes), "a range around unregistered pages is covered");
+    Require(GuestWriteWatchOverlaps_nid_postfix(region.Address(0), region.bytes), "a range around registered pages does not overlap the watch");
+    Require(GuestWriteWatchOverlaps_nid_postfix(region.Address(1) - 1, 2), "a range ending in a registered page does not overlap the watch");
+    Require(GuestWriteWatchOverlaps_nid_postfix(region.Address(3) - 1, 2), "a range starting in a registered page does not overlap the watch");
+    Require(!GuestWriteWatchOverlaps_nid_postfix(region.Address(0), region.page), "an unregistered page before the registered ones overlaps the watch");
+    Require(!GuestWriteWatchOverlaps_nid_postfix(region.Address(3), 3 * region.page), "unregistered pages after the registered ones overlap the watch");
+    Require(!GuestWriteWatchOverlaps_nid_postfix(region.Address(1), 0), "an empty range overlaps the watch");
+}
+
 void CheckThreads() {
     constexpr std::size_t Threads = 4;
     constexpr std::size_t PagesPerThread = 64;
@@ -269,7 +284,7 @@ void CheckThreads() {
 
 int main(int argc, char** argv) {
     try {
-        Require(argc == 2, "expected collect, fresh, armed, register, mapped, host, empty, protection or threads");
+        Require(argc == 2, "expected collect, fresh, armed, register, mapped, host, empty, protection, overlaps or threads");
         const std::string_view mode(argv[1]);
         if (mode == "collect") CheckCollect();
         else if (mode == "fresh") CheckFresh();
@@ -279,6 +294,7 @@ int main(int argc, char** argv) {
         else if (mode == "host") CheckHostWrite();
         else if (mode == "empty") CheckEmptyRanges();
         else if (mode == "protection") CheckProtection();
+        else if (mode == "overlaps") CheckOverlaps();
         else if (mode == "threads") CheckThreads();
         else throw std::runtime_error("unknown write watch test");
         std::printf("write watch %s passed\n", argv[1]);

core/libs/tests/write_watch/CMakeLists.txt

@@ -15,7 +15,7 @@ set_target_properties(write_watch_tests PROPERTIES
     CXX_EXTENSIONS OFF
     RUNTIME_OUTPUT_DIRECTORY "$<1:${ANYPS5_TEST_OUTPUT_DIRECTORY}>")
 
-foreach(mode collect fresh armed register mapped host empty protection threads)
+foreach(mode collect fresh armed register mapped host empty protection overlaps threads)
     add_test(NAME write_watch_${mode} COMMAND write_watch_tests ${mode})
     set_tests_properties(write_watch_${mode} PROPERTIES TIMEOUT 30)
 endforeach()