Introduce filterByRoot extensions for FS providers for easier agent…

.github/CODEOWNERS

@@ -1,3 +1,7 @@
 rag/rag-base/src/commonMain/kotlin/ai/koog/rag/base/files/FileSystemProvider.kt @skarpovdev @sproshev
+
 rag/rag-base/src/jvmMain/kotlin/ai/koog/rag/base/files/JVMFileSystemProvider.kt @skarpovdev @sproshev
-rag/rag-base/src/jvmTest/kotlin/ai/koog/rag/base/files/JVMFileSystemProviderTest.kt @skarpovdev @sproshev
+rag/rag-base/src/jvmTest/kotlin/ai/koog/rag/base/files/JVMFileSystemProviderTest.kt @skarpovdev @sproshev
+
+rag/rag-base/src/commonMain/kotlin/ai/koog/rag/base/files/FilteredFileSystemProvider.kt @skarpovdev @sproshev
+rag/rag-base/src/jvmTest/kotlin/ai/koog/rag/base/files/JVMFilteredFileSystemProviderTest.kt @skarpovdev @sproshev

rag/rag-base/src/commonMain/kotlin/ai/koog/rag/base/files/Extensions.kt

@@ -0,0 +1,21 @@
+package ai.koog.rag.base.files
+
+internal fun <Path> Path.contains(
+    other: Path,
+    fs: FileSystemProvider.ReadOnly<Path>
+): Boolean {
+    val currentComponents = this.components(fs)
+    val otherComponents = other.components(fs)
+    return currentComponents.zip(otherComponents)
+        .all { it.first == it.second } && otherComponents.size >= currentComponents.size
+}
+
+private fun <Path> Path.components(fs: FileSystemProvider.ReadOnly<Path>): List<String> {
+    return buildList {
+        var path: Path? = this@components
+        while (path != null) {
+            add(fs.name(path))
+            path = fs.parent(path)
+        }
+    }.asReversed()
+}

rag/rag-base/src/commonMain/kotlin/ai/koog/rag/base/files/FilteredFileSystemProvider.kt

@@ -0,0 +1,140 @@
+package ai.koog.rag.base.files
+
+import kotlinx.io.IOException
+import kotlinx.io.Sink
+import kotlinx.io.Source
+
+/**
+ * Filters the current read-only file system implementation based on the specified root path such that
+ * only paths that are contained within the given root or paths that contain the root are visible and accessible.
+ *
+ * @param root The root path to use as the basis for filtering.
+ * @return A new file system provider instance that is filtered based on the given root path.
+ */
+public fun <Path> FileSystemProvider.ReadOnly<Path>.filterByRoot(root: Path): FileSystemProvider.ReadOnly<Path> {
+    val filter = PathFilter { path, fs ->
+        root.contains(path, fs) || path.contains(root, fs)
+    }
+    return FilteredReadOnly(this, filter)
+}
+
+/**
+ * Filters the current read-write file system implementation based on the specified root path such that
+ * only paths that are contained within the given root or paths that contain the root are visible and accessible.
+ *
+ * @param root The root path to use as the basis for filtering.
+ * @return A new file system provider instance that is filtered based on the given root path.
+ */
+public fun <Path> FileSystemProvider.ReadWrite<Path>.filterByRoot(root: Path): FileSystemProvider.ReadWrite<Path> {
+    val filter = PathFilter { path, fs ->
+        root.contains(path, fs) || path.contains(root, fs)
+    }
+    return FilteredReadWrite(this, filter)
+}
+
+internal fun interface PathFilter<Path> {
+    fun show(path: Path, fs: FileSystemProvider.ReadOnly<Path>): Boolean
+    fun hide(path: Path, fs: FileSystemProvider.ReadOnly<Path>): Boolean = !show(path, fs)
+}
+
+internal open class FilteredReadOnly<P>(
+    private val fs: FileSystemProvider.ReadOnly<P>,
+    private val filter: PathFilter<P>
+) : FileSystemProvider.ReadOnly<P> {
+    private fun requireAllowed(path: P) {
+        require(filter.show(path, fs)) { "Path $path is hidden by filter" }
+    }
+
+    override suspend fun read(path: P): ByteArray {
+        requireAllowed(path)
+        return fs.read(path)
+    }
+
+    override suspend fun source(path: P): Source {
+        requireAllowed(path)
+        return fs.source(path)
+    }
+
+    override suspend fun size(path: P): Long {
+        requireAllowed(path)
+        return fs.size(path)
+    }
+
+    override suspend fun list(directory: P): List<P> {
+        requireAllowed(directory)
+        return fs.list(directory).filter { child ->
+            filter.show(child, fs)
+        }
+    }
+
+    override suspend fun metadata(path: P): FileMetadata? {
+        return if (filter.show(path, fs)) {
+            fs.metadata(path)
+        } else {
+            null
+        }
+    }
+
+    override suspend fun exists(path: P): Boolean {
+        return if (filter.show(path, fs)) {
+            fs.exists(path)
+        } else {
+            false
+        }
+    }
+
+    override fun parent(path: P): P? = fs.parent(path)
+
+    override fun relativize(root: P, path: P): String? = fs.relativize(root, path)
+
+    override fun toAbsolutePathString(path: P): String = fs.toAbsolutePathString(path)
+
+    override fun fromAbsoluteString(path: String): P = fs.fromAbsoluteString(path)
+
+    override fun fromRelativeString(base: P, path: String): P = fs.fromRelativeString(base, path)
+
+    override fun name(path: P): String = fs.name(path)
+
+    override fun extension(path: P): String = fs.extension(path)
+}
+
+internal class FilteredReadWrite<P>(
+    private val fs: FileSystemProvider.ReadWrite<P>,
+    private val filter: PathFilter<P>
+) : FileSystemProvider.ReadWrite<P>, FilteredReadOnly<P>(fs, filter) {
+    private fun ensureAllowed(path: P) {
+        if (filter.hide(path, fs)) {
+            throw IOException("Path $path is not allowed by filter")
+        }
+    }
+
+    override suspend fun create(parent: P, name: String, type: FileMetadata.FileType) {
+        ensureAllowed(parent)
+        ensureAllowed(fs.fromRelativeString(parent, name))
+        fs.create(parent, name, type)
+    }
+
+    override suspend fun write(path: P, content: ByteArray) {
+        ensureAllowed(path)
+        fs.write(path, content)
+    }
+
+    override suspend fun sink(path: P, append: Boolean): Sink {
+        ensureAllowed(path)
+        return fs.sink(path, append)
+    }
+
+    override suspend fun move(source: P, target: P) {
+        ensureAllowed(source)
+        ensureAllowed(target)
+        fs.move(source, target)
+    }
+
+    override suspend fun delete(parent: P, name: String) {
+        ensureAllowed(parent)
+        ensureAllowed(fs.fromRelativeString(parent, name))
+        fs.delete(parent, name)
+    }
+
+
+}

rag/rag-base/src/jvmTest/kotlin/ai/koog/rag/base/files/JVMFilteredFileSystemProviderTest.kt

@@ -0,0 +1,152 @@
+package ai.koog.rag.base.files
+
+import kotlinx.coroutines.runBlocking
+import kotlinx.io.IOException
+import org.junit.jupiter.api.Test
+import org.junit.jupiter.api.assertThrows
+import java.nio.file.Path
+import kotlin.io.path.name
+import kotlin.test.*
+
+class JVMFilteredFileSystemProviderTest : KoogTestBase() {
+
+    private lateinit var fsReadOnly: FileSystemProvider.ReadOnly<Path>
+    private lateinit var fsReadWrite: FileSystemProvider.ReadWrite<Path>
+
+    /* all the tests below have an assumption that JVMFileSystemProvider is covered by tests,
+     and hence only the filtering should be verified */
+
+    override fun setup() {
+        super.setup()
+        fsReadOnly = JVMFileSystemProvider.ReadOnly.filterByRoot(src1)
+        fsReadWrite = JVMFileSystemProvider.ReadWrite.filterByRoot(src1)
+    }
+
+    @Test
+    fun `test read`(): Unit = runBlocking {
+        assertThrows<IllegalArgumentException> {
+            fsReadOnly.read(file3)
+        }
+
+        fsReadOnly.read(resource1)
+    }
+
+    @Test
+    fun `test source`(): Unit = runBlocking {
+        assertThrows<IllegalArgumentException> {
+            fsReadOnly.source(file3).close()
+        }
+
+        fsReadOnly.source(resource1).close()
+    }
+
+    @Test
+    fun `test size`(): Unit = runBlocking {
+        assertThrows<IllegalArgumentException> {
+            fsReadOnly.size(file3)
+        }
+
+        fsReadOnly.size(resource1)
+    }
+
+    @Test
+    fun `test list`(): Unit = runBlocking {
+        assertThrows<IllegalArgumentException> {
+            fsReadOnly.list(src2)
+        }
+
+        val child = fsReadOnly.list(tempDir).single()
+        assertEquals(dir1, child)
+    }
+
+    @Test
+    fun `test metadata`(): Unit = runBlocking {
+        assertNull(fsReadOnly.metadata(dir3))
+        assertNull(fsReadOnly.metadata(file2))
+
+        assertNotNull(fsReadOnly.metadata(resources))
+        assertNotNull(fsReadOnly.metadata(file1))
+    }
+
+    @Test
+    fun `test exists`(): Unit = runBlocking {
+        assertFalse(fsReadOnly.exists(dir3))
+        assertFalse(fsReadOnly.exists(file2))
+
+        assertTrue(fsReadOnly.exists(resources))
+        assertTrue(fsReadOnly.exists(file1))
+    }
+
+    @Test
+    fun `test select does not throw`() {
+        assertNotNull(fsReadOnly.parent(src2))
+        assertNotNull(fsReadOnly.fromRelativeString(src2, assertNotNull(fsReadOnly.relativize(src2, file2))))
+        assertNotNull(fsReadOnly.fromAbsoluteString(assertNotNull(fsReadOnly.toAbsolutePathString(file2))))
+        assertNotNull(fsReadOnly.name(src2))
+        assertNotNull(fsReadOnly.extension(file2))
+    }
+
+    @Test
+    fun `test create`() = runBlocking {
+        assertThrows<IOException> {
+            // forbidden parent
+            fsReadWrite.create(dir2, "myNewDir", FileMetadata.FileType.Directory)
+        }
+
+        assertThrows<IOException> {
+            // allowed parent
+            fsReadWrite.create(dir1, "myNewFile", FileMetadata.FileType.File)
+        }
+
+        fsReadWrite.create(src1, "myNewDir", FileMetadata.FileType.Directory)
+        fsReadWrite.create(src1, "myNewFile", FileMetadata.FileType.File)
+    }
+
+    @Test
+    fun `test write`() = runBlocking {
+        assertThrows<IOException> {
+            fsReadWrite.write(file2, "myNewContent".toByteArray())
+        }
+
+        fsReadWrite.write(file1, "myNewContent".toByteArray())
+    }
+
+    @Test
+    fun `test sink`() = runBlocking {
+        assertThrows<IOException> {
+            fsReadWrite.sink(file2).close()
+        }
+
+        fsReadWrite.sink(file1).close()
+    }
+
+    @Test
+    fun `test move`() = runBlocking {
+        assertThrows<IOException> {
+            // source is not allowed
+            fsReadWrite.move(file2, src1.resolve(file2.name))
+        }
+
+        assertThrows<IOException> {
+            // destination is not allowed
+            fsReadWrite.move(resources, dir2.resolve(resources.name))
+        }
+
+        fsReadWrite.move(file1, resources.resolve(file1.name))
+    }
+
+    @Test
+    fun `test delete`() = runBlocking {
+        assertThrows<IOException> {
+            // forbidden parent
+            fsReadWrite.delete(src2, file2.name)
+        }
+
+        assertThrows<IOException> {
+            // allowed parent
+            fsReadWrite.delete(dir1, dir2.name)
+        }
+
+        fsReadWrite.delete(src1, file1.name)
+    }
+}

rag/rag-base/src/jvmTest/kotlin/ai/koog/rag/base/files/KoogTestBase.kt

@@ -56,7 +56,7 @@ open class KoogTestBase {
     private val testCode2 = loadTextFromResource("files/testCode2.kt")
 
     @BeforeEach
-    fun setup() {
+    open fun setup() {
         dir1 = tempDir.resolve("dir1").apply { this.createDirectories() }
         src1 = dir1.resolve("src").apply { this.createDirectories() }
         resources = src1.resolve("resources").apply { this.createDirectories() }