Support premature termination of listing (#928)

* Support premature termination of listing

* Added license header + small refactor

---------

Co-authored-by: Jeroen van Erp <jeroen@hierynomus.com>

src/main/java/com/hierynomus/sshj/sftp/RemoteResourceFilterConverter.java

@@ -0,0 +1,30 @@
+/*
+ * Copyright (C)2009 - SSHJ Contributors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.hierynomus.sshj.sftp;
+
+import com.hierynomus.sshj.sftp.RemoteResourceSelector.Result;
+import net.schmizz.sshj.sftp.RemoteResourceFilter;
+
+public class RemoteResourceFilterConverter {
+
+    public static RemoteResourceSelector selectorFrom(RemoteResourceFilter filter) {
+        if (filter == null) {
+            return RemoteResourceSelector.ALL;
+        }
+
+        return resource -> filter.accept(resource) ? Result.ACCEPT : Result.CONTINUE;
+    }
+}

src/main/java/com/hierynomus/sshj/sftp/RemoteResourceSelector.java

@@ -0,0 +1,49 @@
+/*
+ * Copyright (C)2009 - SSHJ Contributors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.hierynomus.sshj.sftp;
+
+import net.schmizz.sshj.sftp.RemoteResourceInfo;
+
+public interface RemoteResourceSelector {
+    public static RemoteResourceSelector ALL = new RemoteResourceSelector() {
+        @Override
+        public Result select(RemoteResourceInfo resource) {
+            return Result.ACCEPT;
+        }
+    };
+
+    enum Result {
+        /**
+         * Accept the remote resource and add it to the result.
+         */
+        ACCEPT,
+
+        /**
+         * Do not add the remote resource to the result and continue with the next.
+         */
+        CONTINUE,
+
+        /**
+         * Do not add the remote resource to the result and stop further execution.
+         */
+        BREAK;
+    }
+
+    /**
+     * Decide whether the remote resource should be included in the result and whether execution should continue.
+     */
+    Result select(RemoteResourceInfo resource);
+}

src/main/java/net/schmizz/sshj/sftp/RemoteDirectory.java

@@ -15,6 +15,7 @@
  */
 package net.schmizz.sshj.sftp;
 
+import com.hierynomus.sshj.sftp.RemoteResourceSelector;
 import net.schmizz.sshj.sftp.Response.StatusCode;
 
 import java.io.IOException;
@@ -22,6 +23,8 @@ import java.util.LinkedList;
 import java.util.List;
 import java.util.concurrent.TimeUnit;
 
+import static com.hierynomus.sshj.sftp.RemoteResourceFilterConverter.selectorFrom;
+
 public class RemoteDirectory
         extends RemoteResource {
 
@@ -31,37 +34,55 @@ public class RemoteDirectory
 
     public List<RemoteResourceInfo> scan(RemoteResourceFilter filter)
             throws IOException {
-        List<RemoteResourceInfo> rri = new LinkedList<RemoteResourceInfo>();
-        // TODO: Remove GOTO! 
-        loop:
-        for (; ; ) {
-            final Response res = requester.request(newRequest(PacketType.READDIR))
-                    .retrieve(requester.getTimeoutMs(), TimeUnit.MILLISECONDS);
-            switch (res.getType()) {
+        return scan(selectorFrom(filter));
+    }
 
+    public List<RemoteResourceInfo> scan(RemoteResourceSelector selector)
+            throws IOException {
+        if (selector == null) {
+            selector = RemoteResourceSelector.ALL;
+        }
+
+        List<RemoteResourceInfo> remoteResourceInfos = new LinkedList<>();
+
+        while (true) {
+            final Response response = requester.request(newRequest(PacketType.READDIR))
+                    .retrieve(requester.getTimeoutMs(), TimeUnit.MILLISECONDS);
+
+            switch (response.getType()) {
                 case NAME:
-                    final int count = res.readUInt32AsInt();
+                    final int count = response.readUInt32AsInt();
                     for (int i = 0; i < count; i++) {
-                        final String name = res.readString(requester.sub.getRemoteCharset());
-                        res.readString(); // long name - IGNORED - shdve never been in the protocol
-                        final FileAttributes attrs = res.readFileAttributes();
+                        final String name = response.readString(requester.sub.getRemoteCharset());
+                        response.readString(); // long name - IGNORED - shdve never been in the protocol
+                        final FileAttributes attrs = response.readFileAttributes();
                         final PathComponents comps = requester.getPathHelper().getComponents(path, name);
                         final RemoteResourceInfo inf = new RemoteResourceInfo(comps, attrs);
-                        if (!(".".equals(name) || "..".equals(name)) && (filter == null || filter.accept(inf))) {
-                            rri.add(inf);
+
+                        if (".".equals(name) || "..".equals(name)) {
+                            continue;
+                        }
+
+                        final RemoteResourceSelector.Result selectionResult = selector.select(inf);
+                        switch (selectionResult) {
+                            case ACCEPT:
+                                remoteResourceInfos.add(inf);
+                                break;
+                            case CONTINUE:
+                                continue;
+                            case BREAK:
+                                return remoteResourceInfos;
                         }
                     }
                     break;
 
                 case STATUS:
-                    res.ensureStatusIs(StatusCode.EOF);
-                    break loop;
+                    response.ensureStatusIs(StatusCode.EOF);
+                    return remoteResourceInfos;
 
                 default:
-                    throw new SFTPException("Unexpected packet: " + res.getType());
+                    throw new SFTPException("Unexpected packet: " + response.getType());
             }
         }
-        return rri;
     }
-
 }

src/main/java/net/schmizz/sshj/sftp/SFTPClient.java

@@ -15,6 +15,7 @@
  */
 package net.schmizz.sshj.sftp;
 
+import com.hierynomus.sshj.sftp.RemoteResourceSelector;
 import net.schmizz.sshj.connection.channel.direct.SessionFactory;
 import net.schmizz.sshj.xfer.FilePermission;
 import net.schmizz.sshj.xfer.LocalDestFile;
@@ -25,6 +26,8 @@ import java.io.Closeable;
 import java.io.IOException;
 import java.util.*;
 
+import static com.hierynomus.sshj.sftp.RemoteResourceFilterConverter.selectorFrom;
+
 public class SFTPClient
         implements Closeable {
 
@@ -57,16 +60,18 @@ public class SFTPClient
 
     public List<RemoteResourceInfo> ls(String path)
             throws IOException {
-        return ls(path, null);
+        return ls(path, RemoteResourceSelector.ALL);
     }
 
     public List<RemoteResourceInfo> ls(String path, RemoteResourceFilter filter)
             throws IOException {
-        final RemoteDirectory dir = engine.openDir(path);
-        try {
-            return dir.scan(filter);
-        } finally {
-            dir.close();
+        return ls(path, selectorFrom(filter));
+    }
+
+    public List<RemoteResourceInfo> ls(String path, RemoteResourceSelector selector)
+            throws IOException {
+        try (RemoteDirectory dir = engine.openDir(path)) {
+            return dir.scan(selector == null ? RemoteResourceSelector.ALL : selector);
         }
     }
 

src/main/java/net/schmizz/sshj/sftp/StatefulSFTPClient.java

@@ -15,6 +15,7 @@
  */
 package net.schmizz.sshj.sftp;
 
+import com.hierynomus.sshj.sftp.RemoteResourceSelector;
 import net.schmizz.sshj.connection.channel.direct.SessionFactory;
 import net.schmizz.sshj.xfer.LocalDestFile;
 import net.schmizz.sshj.xfer.LocalSourceFile;
@@ -23,6 +24,8 @@ import java.io.IOException;
 import java.util.List;
 import java.util.Set;
 
+import static com.hierynomus.sshj.sftp.RemoteResourceFilterConverter.selectorFrom;
+
 public class StatefulSFTPClient
         extends SFTPClient {
 
@@ -57,7 +60,7 @@ public class StatefulSFTPClient
 
     public synchronized List<RemoteResourceInfo> ls()
             throws IOException {
-        return ls(cwd, null);
+        return ls(cwd, RemoteResourceSelector.ALL);
     }
 
     public synchronized List<RemoteResourceInfo> ls(RemoteResourceFilter filter)
@@ -70,20 +73,21 @@ public class StatefulSFTPClient
         return super.canonicalize(cwd);
     }
 
-    @Override
     public List<RemoteResourceInfo> ls(String path)
             throws IOException {
-        return ls(path, null);
+        return ls(path, RemoteResourceSelector.ALL);
+    }
+
+    public List<RemoteResourceInfo> ls(String path, RemoteResourceFilter filter)
+            throws IOException {
+        return ls(path, selectorFrom(filter));
     }
 
     @Override
-    public List<RemoteResourceInfo> ls(String path, RemoteResourceFilter filter)
+    public List<RemoteResourceInfo> ls(String path, RemoteResourceSelector selector)
             throws IOException {
-        final RemoteDirectory dir = getSFTPEngine().openDir(cwdify(path));
-        try {
-            return dir.scan(filter);
-        } finally {
-            dir.close();
+        try (RemoteDirectory dir = getSFTPEngine().openDir(cwdify(path))) {
+            return dir.scan(selector == null ? RemoteResourceSelector.ALL : selector);
         }
     }
 

src/test/groovy/com/hierynomus/sshj/sftp/SFTPClientSpec.groovy

@@ -19,6 +19,7 @@ import com.hierynomus.sshj.test.SshServerExtension
 import com.hierynomus.sshj.test.util.FileUtil
 import net.schmizz.sshj.SSHClient
 import net.schmizz.sshj.sftp.FileMode
+import net.schmizz.sshj.sftp.RemoteResourceInfo
 import net.schmizz.sshj.sftp.SFTPClient
 import org.junit.jupiter.api.extension.RegisterExtension
 import spock.lang.Specification
@@ -206,6 +207,60 @@ class SFTPClientSpec extends Specification {
         attrs.type == FileMode.Type.DIRECTORY
     }
 
+    def "should support premature termination of listing"() {
+        given:
+        SSHClient sshClient = fixture.setupConnectedDefaultClient()
+        sshClient.authPassword("test", "test")
+        SFTPClient sftpClient = sshClient.newSFTPClient()
+
+        final Path source = Files.createDirectory(temp.resolve("source")).toAbsolutePath()
+        final Path destination = Files.createDirectory(temp.resolve("destination")).toAbsolutePath()
+        final Path firstFile = Files.writeString(source.resolve("a_first.txt"), "first")
+        final Path secondFile = Files.writeString(source.resolve("b_second.txt"), "second")
+        final Path thirdFile = Files.writeString(source.resolve("c_third.txt"), "third")
+        final Path fourthFile = Files.writeString(source.resolve("d_fourth.txt"), "fourth")
+        sftpClient.put(firstFile.toString(), destination.resolve(firstFile.fileName).toString())
+        sftpClient.put(secondFile.toString(), destination.resolve(secondFile.fileName).toString())
+        sftpClient.put(thirdFile.toString(), destination.resolve(thirdFile.fileName).toString())
+        sftpClient.put(fourthFile.toString(), destination.resolve(fourthFile.fileName).toString())
+
+        def filesListed = 0
+        RemoteResourceInfo expectedFile = null
+        RemoteResourceSelector limitingSelector = new RemoteResourceSelector() {
+            @Override
+            RemoteResourceSelector.Result select(RemoteResourceInfo resource) {
+                filesListed += 1
+
+                switch(filesListed) {
+                    case 1:
+                        return RemoteResourceSelector.Result.CONTINUE
+                    case 2:
+                        expectedFile = resource
+                        return RemoteResourceSelector.Result.ACCEPT
+                    case 3:
+                        return RemoteResourceSelector.Result.BREAK
+                    default:
+                        throw new AssertionError((Object) "Should NOT select any more resources")
+                }
+            }
+        }
+
+        when:
+        def listingResult = sftpClient
+                .ls(destination.toString(), limitingSelector);
+
+        then:
+        // first should be skipped by CONTINUE
+        listingResult.contains(expectedFile) // second should be included by ACCEPT
+        // third should be skipped by BREAK
+        // fourth should be skipped by preceding BREAK
+        listingResult.size() == 1
+
+        cleanup:
+        sftpClient.close()
+        sshClient.disconnect()
+    }
+
     private void doUpload(File src, File dest) throws IOException {
         SSHClient sshClient = fixture.setupConnectedDefaultClient()
         sshClient.authPassword("test", "test")