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