Additional test cases for probing MTU/transact size for various dialects

pike/__init__.py

@@ -13,5 +13,5 @@
         'test',
         'transport',
 ]
-__version_info__ = (0, 2, 20)
+__version_info__ = (0, 2, 21)
 __version__ = "{0}.{1}.{2}".format(*__version_info__)

pike/test/__init__.py

@@ -161,7 +161,7 @@ class _AssertErrorContext(object):
         pass
 
     @contextlib.contextmanager
-    def assert_error(self, status):
+    def assert_error(self, *status):
         e = None
         o = PikeTest._AssertErrorContext()
 
@@ -172,7 +172,7 @@ def assert_error(self, status):
 
         if e is None:
             raise self.failureException('No error raised when "%s" expected' % status)
-        elif e.response.status != status:
+        elif e.response.status not in status:
             raise self.failureException('"%s" raised when "%s" expected' % (e.response.status, status))
 
         o.response = e.response
@@ -236,6 +236,28 @@ def assertBufferEqual(self, buf1, buf2):
             raise AssertionError("Block mismatch at byte {0}: "
                                  "{1} != {2}".format(low, buf1[low], buf2[low]))
 
+
+class TreeConnectWithDialect(object):
+    """
+    Mixin class provides `tree_connect_with_dialect_and_caps` contextmanager
+    """
+    @contextlib.contextmanager
+    def tree_connect_with_dialect_and_caps(self, dialect=None, caps=0):
+        client = model.Client(capabilities=caps)
+        if dialect is not None:
+            client.dialects = [dialect, smb2.DIALECT_SMB2_002]
+        chan, tree = self.tree_connect(client)
+        try:
+            yield chan, tree
+        finally:
+            if chan.connection.connected:
+                try:
+                    chan.logoff()
+                    chan.connection.close()
+                except EOFError:
+                    pass
+
+
 class _Decorator(object):
     def __init__(self, value):
         self.value = value

pike/test/querydirectory.py

@@ -39,6 +39,8 @@
 import pike.test
 import pike.ntstatus
 
+from contextlib import contextmanager
+
 
 class QueryDirectoryTest(pike.test.PikeTest):
     # Enumerate directory at FILE_DIRECTORY_INFORMATION level.
@@ -67,7 +69,9 @@ def test_specific_name(self):
 
         hello = chan.create(tree,
                             name,
-                            access=pike.smb2.GENERIC_WRITE | pike.smb2.GENERIC_READ | pike.smb2.DELETE,
+                            access=(pike.smb2.GENERIC_WRITE |
+                                    pike.smb2.GENERIC_READ |
+                                    pike.smb2.DELETE),
                             disposition=pike.smb2.FILE_SUPERSEDE,
                             options=pike.smb2.FILE_DELETE_ON_CLOSE).result()
 
@@ -124,3 +128,127 @@ def test_restart_scan(self):
         self.assertIn('.', names)
 
         chan.close(root)
+
+
+class QueryDirectoryTestMaxMtu(pike.test.PikeTest,
+                               pike.test.TreeConnectWithDialect):
+    root_dir_name = "mtu_transport_query_dir"
+    filename_prefix = "A" * 211
+    filename_pattern = "{0}.test.{{0:05}}".format(filename_prefix)
+    payload_size = 65536
+    n_entries = 0
+    filenames = []
+    dataset_created = False
+
+    def create_dataset(self):
+        if QueryDirectoryTestMaxMtu.dataset_created:
+            return
+        # 66 bytes is the struct size of FILE_DIRECTORY_INFORMATION
+        n_entries = (self.payload_size /
+                     ((len(self.filename_pattern) - 1) * 2 + 66)) + 1
+        QueryDirectoryTestMaxMtu.n_entries = n_entries
+        self.info("Creating {0} files to fill {1} bytes".format(
+                        self.n_entries,
+                        self.payload_size))
+
+        chan, tree = self.tree_connect()
+        with self.get_test_root(chan, tree) as _:
+            pass            # ensure the test root dir is created
+        # Creates files within the root directory.
+        remaining_files = self.n_entries
+        while remaining_files > 0:
+            batch_futures = []
+            batch_size = 30
+            if remaining_files < batch_size:
+                batch_size = remaining_files
+            for ix in xrange(batch_size):
+                filename = self.filename_pattern.format(remaining_files)
+                path = "{0}\\{1}".format(self.root_dir_name, filename)
+                batch_futures.append((filename, chan.create(tree, path)))
+                remaining_files -= 1
+            for filename, fu in batch_futures:
+                fh = fu.result()
+                chan.close(fh)
+                self.filenames.append(filename)
+        chan.logoff()
+        QueryDirectoryTestMaxMtu.dataset_created = True
+
+    def setUp(self):
+        self.create_dataset()
+
+    @contextmanager
+    def get_test_root(self, chan, tree):
+        root_handle = chan.create(
+            tree,
+            self.root_dir_name,
+            access=pike.smb2.GENERIC_READ,
+            options=pike.smb2.FILE_DIRECTORY_FILE,
+            share=(pike.smb2.FILE_SHARE_READ |
+                   pike.smb2.FILE_SHARE_WRITE |
+                   pike.smb2.FILE_SHARE_DELETE)).result()
+        try:
+            yield root_handle
+        finally:
+            chan.close(root_handle)
+
+    def gen_file_directory_info_max_mtu(self, dialect=None, caps=0):
+        """
+        Enumerate directory at FILE_DIRECTORY_INFORMATION level using the
+        maximum transfer size.
+        """
+
+        with self.tree_connect_with_dialect_and_caps(dialect, caps) as (chan,
+                                                                        tree):
+            # Retrieves the maximum transaction size to determine how big the
+            # command's payload can be (MTU).
+            max_trans_size = chan.connection.negotiate_response.max_transact_size
+
+            # Enumerates the root directory and validates that the expected
+            # files are present in it.
+            with self.get_test_root(chan, tree) as root_handle:
+                dir_query1 = chan.query_directory(
+                    root_handle,
+                    file_information_class=pike.smb2.FILE_DIRECTORY_INFORMATION,
+                    flags=0,
+                    file_index=0,
+                    file_name='{0}*'.format(self.filename_prefix),
+                    output_buffer_length=self.payload_size)
+                dir_query2 = chan.query_directory(
+                    root_handle,
+                    file_information_class=pike.smb2.FILE_DIRECTORY_INFORMATION,
+                    flags=0,
+                    file_name='{0}*'.format(self.filename_prefix),
+                    output_buffer_length=self.payload_size)
+
+                transaction_size1 = dir_query1[-1].end - dir_query1[0].start
+                self.info("Transaction size for query: {0}/{1}".format(
+                                transaction_size1, max_trans_size))
+                transaction_size2 = dir_query2[-1].end - dir_query2[0].start
+                self.info("Transaction size for query: {0}/{1}".format(
+                                transaction_size2, max_trans_size))
+                self.info("Dir entries returned: {0}/{1}".format(
+                                len(dir_query1) + len(dir_query2),
+                                self.n_entries))
+                self.assertLessEqual(transaction_size1,
+                                     max_trans_size)
+                self.assertGreaterEqual(transaction_size1 + transaction_size2,
+                                        self.payload_size)
+
+    def test_file_directory_info(self):
+        self.gen_file_directory_info_max_mtu()
+
+    @pike.test.RequireDialect(pike.smb2.DIALECT_SMB2_002)
+    def test_file_directory_info_2_002(self):
+        self.gen_file_directory_info_max_mtu(dialect=pike.smb2.DIALECT_SMB2_002)
+
+    @pike.test.RequireDialect(pike.smb2.DIALECT_SMB2_1)
+    def test_file_directory_info_2_1(self):
+        self.gen_file_directory_info_max_mtu(dialect=pike.smb2.DIALECT_SMB2_1)
+
+    @pike.test.RequireDialect(pike.smb2.DIALECT_SMB3_0)
+    def test_file_directory_info_3_0(self):
+        self.gen_file_directory_info_max_mtu(dialect=pike.smb2.DIALECT_SMB3_0)
+
+    @pike.test.RequireCapabilities(pike.smb2.SMB2_GLOBAL_CAP_LARGE_MTU)
+    def test_file_directory_info_large_mtu(self):
+        self.gen_file_directory_info_max_mtu(caps=pike.smb2.SMB2_GLOBAL_CAP_LARGE_MTU)