feat: FileMutationReverser undo primitive + shared AtomicFileWriter

2026-05-31 10:36:47 +04:00
parent b778c06f6c
commit 3c90ac4aef
4 changed files with 173 additions and 23 deletions
@@ -0,0 +1,29 @@
+package com.correx.infrastructure.tools.filesystem
+
+import java.nio.file.Files
+import java.nio.file.Path
+import java.nio.file.Paths
+import java.nio.file.StandardCopyOption
+
+/**
+ * Crash-safe write primitive shared by the write tool and the undo reverser:
+ * stage into a temp file in the target's OWN directory (same filesystem, so the
+ * rename is atomic) then atomically swap. A partial write is never observable at
+ * [target]; on failure the temp is cleaned and the original is left untouched.
+ * Falls back to a non-atomic replace only if the platform rejects ATOMIC_MOVE.
+ */
+internal object AtomicFileWriter {
+    fun write(target: Path, content: ByteArray) {
+        target.parent?.let { Files.createDirectories(it) }
+        val dir = target.parent ?: Paths.get(".")
+        val tmp = Files.createTempFile(dir, ".${target.fileName}", ".tmp")
+        runCatching {
+            Files.write(tmp, content)
+            runCatching { Files.move(tmp, target, StandardCopyOption.ATOMIC_MOVE) }
+                .getOrElse { Files.move(tmp, target, StandardCopyOption.REPLACE_EXISTING) }
+        }.onFailure {
+            Files.deleteIfExists(tmp)
+            throw it
+        }
+    }
+}
@@ -0,0 +1,52 @@
+package com.correx.infrastructure.tools.filesystem
+
+import com.correx.core.artifactstore.ArtifactStore
+import com.correx.core.events.events.FileWrittenEvent
+import com.correx.core.events.types.ArtifactId
+import java.nio.file.Files
+import java.nio.file.Path
+
+/**
+ * Reverses a recorded [FileWrittenEvent] using only the event + CAS (invariants #1/#8):
+ * restores the pre-image blob, or deletes the file if it didn't previously exist.
+ * The restore goes through the same atomic write seam ([AtomicFileWriter]) and is
+ * jailed through [PathJail] (fail-closed) and idempotent. Pure function of recorded
+ * state — never re-observes the environment to decide what to do.
+ */
+class FileMutationReverser(
+    private val artifactStore: ArtifactStore,
+    private val allowedRoots: Set<Path>,
+) {
+    @Suppress("ReturnCount")
+    suspend fun revert(event: FileWrittenEvent): RevertResult {
+        val target = Path.of(event.path)
+        if (!PathJail.isContained(target, allowedRoots)) {
+            return RevertResult.Rejected("target outside allowed roots: ${event.path}")
+        }
+        return if (event.preExisted) {
+            val hash = event.preImageHash
+                ?: return RevertResult.Rejected("preExisted but no pre-image hash: ${event.path}")
+            val bytes = artifactStore.get(ArtifactId(hash))
+                ?: return RevertResult.Rejected("pre-image blob not found in CAS: $hash")
+            runCatching { AtomicFileWriter.write(target, bytes) }
+                .fold({ RevertResult.Restored }, { RevertResult.Failed(it.message ?: "io error") })
+        } else {
+            runCatching { Files.deleteIfExists(target) }
+                .fold({ RevertResult.Deleted }, { RevertResult.Failed(it.message ?: "io error") })
+        }
+    }
+}
+
+sealed interface RevertResult {
+    /** Pre-image bytes were written back to the target. */
+    data object Restored : RevertResult
+
+    /** The (previously non-existent) target was removed. */
+    data object Deleted : RevertResult
+
+    /** The event could not be reversed safely (jail/missing blob); no change made. */
+    data class Rejected(val reason: String) : RevertResult
+
+    /** An IO error occurred during the restore. */
+    data class Failed(val reason: String) : RevertResult
+}
@@ -31,7 +31,6 @@ import java.nio.file.Files
 import java.nio.file.InvalidPathException
 import java.nio.file.Path
 import java.nio.file.Paths
-import java.nio.file.StandardCopyOption

 class FileWriteTool(
    allowedPaths: Set<Path> = emptySet(),
@@ -173,7 +172,7 @@ class FileWriteTool(
        "write" -> {
            val content = request.parameters["content"] as String
            withContext(Dispatchers.IO) {
-                atomicWrite(path, content)
+                AtomicFileWriter.write(path, content.toByteArray(Charsets.UTF_8))
            }
            storeArtifact(pathString, content, request.parameters["mode"] as? String ?: "0644")
            ToolResult.Success(
@@ -207,27 +206,6 @@ class FileWriteTool(
        )
    }

-    /**
-     * Crash-safe write: stage into a temp file in the target's OWN directory (same
-     * filesystem, so the rename is atomic) then atomically swap it into place. A
-     * partial write can never be observed at [target]; on failure the temp is
-     * cleaned and the original is left untouched. Falls back to a non-atomic
-     * replace only if the platform rejects ATOMIC_MOVE.
-     */
-    private fun atomicWrite(target: Path, content: String) {
-        target.parent?.let { Files.createDirectories(it) }
-        val dir = target.parent ?: Paths.get(".")
-        val tmp = Files.createTempFile(dir, ".${target.fileName}", ".tmp")
-        runCatching {
-            Files.writeString(tmp, content)
-            runCatching { Files.move(tmp, target, StandardCopyOption.ATOMIC_MOVE) }
-                .getOrElse { Files.move(tmp, target, StandardCopyOption.REPLACE_EXISTING) }
-        }.onFailure {
-            Files.deleteIfExists(tmp)
-            throw it
-        }
-    }
-
    @Suppress("ReturnCount")
    private suspend fun storeArtifact(pathString: String, content: String, mode: String) {
        val writer = materializingWriter ?: return
@@ -0,0 +1,91 @@
+package com.correx.infrastructure.tools.filesystem
+
+import com.correx.core.artifactstore.ArtifactStore
+import com.correx.core.events.events.FileWrittenEvent
+import com.correx.core.events.types.ArtifactId
+import com.correx.core.events.types.SessionId
+import com.correx.core.events.types.ToolInvocationId
+import kotlinx.coroutines.runBlocking
+import org.junit.jupiter.api.Assertions.assertEquals
+import org.junit.jupiter.api.Assertions.assertFalse
+import org.junit.jupiter.api.Assertions.assertTrue
+import org.junit.jupiter.api.Test
+import java.nio.file.Files
+
+class FileMutationReverserTest {
+
+    private class FakeStore : ArtifactStore {
+        val blobs = linkedMapOf<String, ByteArray>()
+        override suspend fun put(bytes: ByteArray): ArtifactId {
+            val id = ArtifactId(bytes.toString(Charsets.UTF_8).hashCode().toString())
+            blobs[id.value] = bytes
+            return id
+        }
+        override suspend fun get(id: ArtifactId): ByteArray? = blobs[id.value]
+        override suspend fun flushBefore(commit: suspend () -> Unit) = commit()
+    }
+
+    private fun event(path: String, preImageHash: String?, preExisted: Boolean) = FileWrittenEvent(
+        invocationId = ToolInvocationId("i"),
+        sessionId = SessionId("s"),
+        path = path,
+        preImageHash = preImageHash,
+        postImageHash = null,
+        preExisted = preExisted,
+        timestampMs = 0L,
+    )
+
+    @Test
+    fun `restores the pre-image and is idempotent`(): Unit = runBlocking {
+        val dir = Files.createTempDirectory("rev").toRealPath()
+        val target = dir.resolve("f.txt")
+        Files.writeString(target, "OLD")
+        val store = FakeStore()
+        val hash = store.put("OLD".toByteArray()).value
+        Files.writeString(target, "NEW")
+
+        val reverser = FileMutationReverser(store, setOf(dir))
+        assertEquals(RevertResult.Restored, reverser.revert(event(target.toString(), hash, preExisted = true)))
+        assertEquals("OLD", Files.readString(target))
+        // idempotent: reverting again is still a clean restore
+        assertEquals(RevertResult.Restored, reverser.revert(event(target.toString(), hash, preExisted = true)))
+        assertEquals("OLD", Files.readString(target))
+    }
+
+    @Test
+    fun `deletes a file that did not previously exist and is idempotent`(): Unit = runBlocking {
+        val dir = Files.createTempDirectory("rev-del").toRealPath()
+        val target = dir.resolve("new.txt")
+        Files.writeString(target, "X")
+        val reverser = FileMutationReverser(FakeStore(), setOf(dir))
+
+        assertEquals(RevertResult.Deleted, reverser.revert(event(target.toString(), null, preExisted = false)))
+        assertFalse(Files.exists(target))
+        assertEquals(RevertResult.Deleted, reverser.revert(event(target.toString(), null, preExisted = false)))
+    }
+
+    @Test
+    fun `rejects when the pre-image blob is missing and leaves the file untouched`(): Unit = runBlocking {
+        val dir = Files.createTempDirectory("rev-missing").toRealPath()
+        val target = dir.resolve("f.txt")
+        Files.writeString(target, "NEW")
+        val reverser = FileMutationReverser(FakeStore(), setOf(dir))
+
+        val res = reverser.revert(event(target.toString(), "no-such-hash", preExisted = true))
+        assertTrue(res is RevertResult.Rejected)
+        assertEquals("NEW", Files.readString(target))
+    }
+
+    @Test
+    fun `rejects a target outside the allowed roots without modifying it`(): Unit = runBlocking {
+        val dir = Files.createTempDirectory("rev-in").toRealPath()
+        val outside = Files.createTempDirectory("rev-out").toRealPath()
+        val target = outside.resolve("f.txt")
+        Files.writeString(target, "KEEP")
+        val reverser = FileMutationReverser(FakeStore(), setOf(dir))
+
+        val res = reverser.revert(event(target.toString(), null, preExisted = false))
+        assertTrue(res is RevertResult.Rejected)
+        assertTrue(Files.exists(target))
+    }
+}