From 26b42e28371d23a1d63502edade06eb6764216c4 Mon Sep 17 00:00:00 2001
From: Jordan Borean <jborean93@gmail.com>
Date: Fri, 6 Dec 2019 06:32:36 +1000
Subject: [PATCH] Fix for query_directory on large dirs (#23)

---
 CHANGELOG.md               |  1 +
 smbclient/_io.py           | 19 ++++++++-----------
 smbprotocol/exceptions.py  |  1 +
 tests/test_smbclient_os.py | 20 ++++++++++++++++++++
 4 files changed, 30 insertions(+), 11 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index a7ebea2..bfbb339 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -3,6 +3,7 @@
 ## 1.0.1 - TBD
 
 * Fix issue when reading a large file that exceeds 65KB and raises `STATUS_END_OF_FILE`.
+* Fix issue where `listdir`, `scandir`, `walk` would only enumerate a subset of entries in a directories with lots of sub files/folders
 
 
 ## 1.0.0 - 2019-11-30
diff --git a/smbclient/_io.py b/smbclient/_io.py
index b583ffa..1a1ffc2 100644
--- a/smbclient/_io.py
+++ b/smbclient/_io.py
@@ -485,22 +485,19 @@ class SMBDirectoryIO(SMBRawIO):
     _INVALID_MODE = 'w+'
 
     def query_directory(self, pattern, info_class):
-        idx = 0
+        query_flags = QueryDirectoryFlags.SMB2_RESTART_SCANS
         while True:
-            entries = self.fd.query_directory(pattern, info_class, flags=QueryDirectoryFlags.SMB2_INDEX_SPECIFIED,
-                                              file_index=idx)
+            try:
+                entries = self.fd.query_directory(pattern, info_class, flags=query_flags)
+            except SMBResponseException as exc:
+                if exc.status == NtStatus.STATUS_NO_MORE_FILES:
+                    break
+                raise
 
-            end = False
+            query_flags = 0  # Only the first request should have set SMB2_RESTART_SCANS
             for entry in entries:
-                idx = entry['next_entry_offset'].get_value()
-                if idx == 0:
-                    end = True
-
                 yield entry
 
-            if end:
-                break
-
     def readable(self):
         return False
 
diff --git a/smbprotocol/exceptions.py b/smbprotocol/exceptions.py
index 436bfcd..33768be 100644
--- a/smbprotocol/exceptions.py
+++ b/smbprotocol/exceptions.py
@@ -49,6 +49,7 @@ class NtStatus(object):
     STATUS_NOTIFY_CLEANUP = 0x0000010B
     STATUS_NOTIFY_ENUM_DIR = 0x0000010C
     STATUS_BUFFER_OVERFLOW = 0x80000005
+    STATUS_NO_MORE_FILES = 0x80000006
     STATUS_END_OF_FILE = 0xC0000011
     STATUS_INVALID_EA_NAME = 0x80000013
     STATUS_EA_LIST_INCONSISTENT = 0x80000014
diff --git a/tests/test_smbclient_os.py b/tests/test_smbclient_os.py
index b94edc8..0e40520 100644
--- a/tests/test_smbclient_os.py
+++ b/tests/test_smbclient_os.py
@@ -1147,6 +1147,26 @@ def test_rmdir_symlink_with_src(smb_share):
     assert smbclient.listdir(smb_share) == ['dir']
 
 
+def test_scandir_large(smb_share):
+    dir_path = ntpath.join(smb_share, 'directory')
+
+    # Create lots of directories with the maximum name possible to ensure they won't be returned in 1 request.
+    smbclient.mkdir(dir_path)
+    for i in range(150):
+        dirname = str(i).zfill(255)
+        smbclient.mkdir(ntpath.join(smb_share, 'directory', dirname))
+
+    actual = []
+    for entry in smbclient.scandir(dir_path):
+        actual.append(entry.path)
+
+    # Just a test optimisation, remove all the dirs so we don't have to re-enumerate them again in rmtree.
+    for path in actual:
+        smbclient.rmdir(path)
+
+    assert len(actual) == 150
+
+
 def test_scandir(smb_share):
     dir_path = ntpath.join(smb_share, 'directory')
     smbclient.makedirs(dir_path, exist_ok=True)
-- 
GitLab