Download/Enumerate: support ShouldRecurse/ShouldInclude.

tmds · Apr 3, 2024 · 077de6c · 077de6c
1 parent 5c77d15
commit 077de6c
Show file tree

Hide file tree

Showing 7 changed files with 83 additions and 2 deletions.
diff --git a/README.md b/README.md
@@ -250,6 +250,8 @@ class EnumerationOptions
     bool FollowFileLinks { get; set; } = true;
     bool FollowDirectoryLinks { get; set; } = true;
     UnixFileTypeFilter FileTypeFilter { get; set; } = RegularFile | Directory | SymbolicLink | CharacterDevice | BlockDevice | Socket | Fifo;
+    SftpFileEntryPredicate? ShouldRecurse { get; set; }
+    SftpFileEntryPredicate? ShouldInclude { get; set; }
 }
 class DownloadEntriesOptions
 {
@@ -258,6 +260,8 @@ class DownloadEntriesOptions
     bool FollowFileLinks { get; set; } = true;
     bool FollowDirectoryLinks { get; set; } = true;
     UnixFileTypeFilter FileTypeFilter { get; set; } = RegularFile | Directory | SymbolicLink;
+    SftpFileEntryPredicate? ShouldRecurse { get; set; }
+    SftpFileEntryPredicate? ShouldInclude { get; set; }   
 }
 class UploadEntriesOptions
 {
@@ -267,6 +271,7 @@ class UploadEntriesOptions
     bool FollowDirectoryLinks { get; set; } = true;
 }
 delegate T SftpFileEntryTransform<T>(ref SftpFileEntry entry);
+delegate bool SftpFileEntryPredicate(ref SftpFileEntry entry);
 ref struct SftpFileEntry
 {
     long Length { get; }

diff --git a/src/Tmds.Ssh/DownloadEntriesOptions.cs b/src/Tmds.Ssh/DownloadEntriesOptions.cs
@@ -13,4 +13,6 @@ public sealed class DownloadEntriesOptions
         UnixFileTypeFilter.RegularFile |
         UnixFileTypeFilter.Directory |
         UnixFileTypeFilter.SymbolicLink;
+    public SftpFileEntryPredicate? ShouldRecurse { get; set; }
+    public SftpFileEntryPredicate? ShouldInclude { get; set; }
 }
diff --git a/src/Tmds.Ssh/EnumerationOptions.cs b/src/Tmds.Ssh/EnumerationOptions.cs
@@ -16,4 +16,6 @@ public sealed class EnumerationOptions
         UnixFileTypeFilter.BlockDevice |
         UnixFileTypeFilter.Socket |
         UnixFileTypeFilter.Fifo;
+    public SftpFileEntryPredicate? ShouldRecurse { get; set; }
+    public SftpFileEntryPredicate? ShouldInclude { get; set; }
 }
diff --git a/src/Tmds.Ssh/SftpClient.cs b/src/Tmds.Ssh/SftpClient.cs
@@ -611,7 +611,9 @@ public async ValueTask DownloadDirectoryEntriesAsync(string remoteDirPath, strin
                         RecurseSubdirectories = options.RecurseSubdirectories,
                         FollowDirectoryLinks = options.FollowDirectoryLinks,
                         FollowFileLinks = options.FollowFileLinks,
-                        FileTypeFilter = options.FileTypeFilter });
+                        FileTypeFilter = options.FileTypeFilter,
+                        ShouldInclude = options.ShouldInclude,
+                        ShouldRecurse = options.ShouldRecurse });
 
             var onGoing = new Queue<ValueTask>();
             try

diff --git a/src/Tmds.Ssh/SftpFileEntry.cs b/src/Tmds.Ssh/SftpFileEntry.cs
@@ -9,6 +9,7 @@
 namespace Tmds.Ssh;
 
 public delegate T SftpFileEntryTransform<T>(ref SftpFileEntry entry);
+public delegate bool SftpFileEntryPredicate(ref SftpFileEntry entry);
 
 public ref struct SftpFileEntry
 {

diff --git a/src/Tmds.Ssh/SftpFileSystemEnumerable.cs b/src/Tmds.Ssh/SftpFileSystemEnumerable.cs
@@ -45,6 +45,8 @@ sealed class SftpFileSystemEnumerator<T> : IAsyncEnumerator<T>
     private readonly bool _followFileLinks;
     private readonly bool _followDirectoryLinks;
     private readonly UnixFileTypeFilter _fileTypeFilter;
+    private readonly SftpFileEntryPredicate? _shouldRecurse;
+    private readonly SftpFileEntryPredicate? _shouldInclude;
 
     private Queue<string>? _pending;
 
@@ -66,6 +68,8 @@ public SftpFileSystemEnumerator(SftpClient client, string path, SftpFileEntryTra
         _followDirectoryLinks = options.FollowDirectoryLinks;
         _followFileLinks = options.FollowFileLinks;
         _fileTypeFilter = options.FileTypeFilter;
+        _shouldInclude = options.ShouldInclude;
+        _shouldRecurse = options.ShouldRecurse;
     }
 
     public T Current => _current!;
@@ -191,7 +195,9 @@ private bool ReadNextEntry(bool followLink, out string? linkPath, out Memory<byt
 
     private bool SetCurrent(ref SftpFileEntry entry)
     {
-        if (_recurseSubdirectories && entry.FileType == UnixFileType.Directory)
+        if (_recurseSubdirectories &&
+            entry.FileType == UnixFileType.Directory &&
+            (_shouldRecurse?.Invoke(ref entry) ?? true))
         {
             _pending ??= new();
             _pending.Enqueue(entry.ToPath());
@@ -202,6 +208,10 @@ private bool SetCurrent(ref SftpFileEntry entry)
             return false;
         }
 
+        if (!(_shouldInclude?.Invoke(ref entry) ?? true))
+        {
+            return false;
+        }
         _current = _transform(ref entry);
         return true;
     }

diff --git a/test/Tmds.Ssh.Tests/SftpClientTests.cs b/test/Tmds.Ssh.Tests/SftpClientTests.cs
@@ -326,7 +326,66 @@ public async Task EnumerateDirectoryRecursive(bool recurse)
                     entries.Select(entry => entry.Path).ToHashSet()
                 );
             }
+        }
+
+        [InlineData(true)]
+        [InlineData(false)]
+        [Theory]
+        public async Task EnumerateDirectoryShouldInclude(bool value)
+        {
+            using var client = await _sshServer.CreateClientAsync();
+            using var sftpClient = await client.CreateSftpClientAsync();
+            string directoryPath = $"/tmp/{Path.GetRandomFileName()}";
+
+            await sftpClient.CreateNewDirectoryAsync(directoryPath);
+
+            for (int i = 0; i < 2; i++)
+            {
+                using var file = await sftpClient.CreateNewFileAsync($"{directoryPath}/file{i}", FileAccess.Write);
+                await file.CloseAsync();
+            }
+
+            List<(string Path, FileEntryAttributes Attributes)> entries = await sftpClient.GetDirectoryEntriesAsync(directoryPath,
+                new EnumerationOptions() { ShouldInclude = (ref SftpFileEntry entry) => value }).ToListAsync();
 
+            if (value == true)
+            {
+                Assert.Equal(2, entries.Count);
+            }
+            else
+            {
+                Assert.Empty(entries);
+            }
+        }
+
+        [InlineData(true)]
+        [InlineData(false)]
+        [Theory]
+        public async Task EnumerateDirectoryShouldRecurse(bool value)
+        {
+            using var client = await _sshServer.CreateClientAsync();
+            using var sftpClient = await client.CreateSftpClientAsync();
+            string directoryPath = $"/tmp/{Path.GetRandomFileName()}";
+
+            await sftpClient.CreateNewDirectoryAsync($"{directoryPath}/dir", createParents: true);
+
+            for (int i = 0; i < 2; i++)
+            {
+                using var file = await sftpClient.CreateNewFileAsync($"{directoryPath}/dir/file{i}", FileAccess.Write);
+                await file.CloseAsync();
+            }
+
+            List<(string Path, FileEntryAttributes Attributes)> entries = await sftpClient.GetDirectoryEntriesAsync(directoryPath,
+                new EnumerationOptions() { ShouldRecurse = (ref SftpFileEntry entry) => value, RecurseSubdirectories = true }).ToListAsync();
+
+            if (value == true)
+            {
+                Assert.Equal(3, entries.Count);
+            }
+            else
+            {
+                Assert.Single(entries);
+            }
         }
 
         [Fact]