diff --git a/pike/__init__.py b/pike/__init__.py index 4ac97bb6..5c3c93c8 100644 --- a/pike/__init__.py +++ b/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__) diff --git a/pike/test/__init__.py b/pike/test/__init__.py index de161723..4129a452 100644 --- a/pike/test/__init__.py +++ b/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 diff --git a/pike/test/querydirectory.py b/pike/test/querydirectory.py index 72909d3a..a95ac972 100644 --- a/pike/test/querydirectory.py +++ b/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) diff --git a/pike/test/readwrite.py b/pike/test/readwrite.py index f11f1827..fd2827ae 100644 --- a/pike/test/readwrite.py +++ b/pike/test/readwrite.py @@ -39,7 +39,10 @@ import pike.test import pike.ntstatus import array +import errno import random +import socket + class ReadWriteTest(pike.test.PikeTest): # Test that we can write to a file @@ -56,34 +59,12 @@ def test_write(self): disposition=pike.smb2.FILE_SUPERSEDE, options=pike.smb2.FILE_DELETE_ON_CLOSE, oplock_level=pike.smb2.SMB2_OPLOCK_LEVEL_EXCLUSIVE).result() - + bytes_written = chan.write(file, 0, buffer) self.assertEqual(bytes_written, len(buffer)) - - chan.close(file) - - # Test that a 0-byte write succeeds - def test_write_none(self): - chan, tree = self.tree_connect() - buffer = None - share_all = pike.smb2.FILE_SHARE_READ | pike.smb2.FILE_SHARE_WRITE | pike.smb2.FILE_SHARE_DELETE - - file = chan.create(tree, - 'write.txt', - access=pike.smb2.FILE_READ_DATA | pike.smb2.FILE_WRITE_DATA | pike.smb2.DELETE, - share=share_all, - disposition=pike.smb2.FILE_SUPERSEDE, - options=pike.smb2.FILE_DELETE_ON_CLOSE, - oplock_level=pike.smb2.SMB2_OPLOCK_LEVEL_EXCLUSIVE).result() - - bytes_written = chan.write(file, - 0, - buffer) - self.assertEqual(bytes_written, 0) - chan.close(file) # Test that a 0-byte write succeeds @@ -100,12 +81,12 @@ def test_write_none(self): disposition=pike.smb2.FILE_SUPERSEDE, options=pike.smb2.FILE_DELETE_ON_CLOSE, oplock_level=pike.smb2.SMB2_OPLOCK_LEVEL_EXCLUSIVE).result() - + bytes_written = chan.write(file, 0, buffer) self.assertEqual(bytes_written, 0) - + chan.close(file) # Test that a 0-byte write triggers access checks @@ -121,10 +102,10 @@ def test_write_none_access(self): share=share_all, disposition=pike.smb2.FILE_SUPERSEDE, oplock_level=pike.smb2.SMB2_OPLOCK_LEVEL_EXCLUSIVE).result() - + with self.assert_error(pike.ntstatus.STATUS_ACCESS_DENIED): chan.write(file, 0, None) - + chan.close(file) # Test that 0-byte write does not cause an oplock break @@ -196,3 +177,204 @@ def test_write_none_lease(self): chan.close(file) chan.close(file2) + + +class WriteReadMaxMtu(pike.test.PikeTest, + pike.test.TreeConnectWithDialect): + invalid_write_status = [ + pike.ntstatus.STATUS_INVALID_PARAMETER, # windows 2012+ + pike.ntstatus.STATUS_BUFFER_OVERFLOW] # windows 2008r2 / 7 + invalid_read_status = [ + pike.ntstatus.STATUS_INVALID_PARAMETER, # windows 2012+ + pike.ntstatus.STATUS_BUFFER_OVERFLOW, # windows 2008r2 / 7 + pike.ntstatus.STATUS_INVALID_NETWORK_RESPONSE] # onefs + + write_buf = None + + def setUp(self): + if WriteReadMaxMtu.write_buf is None: + with self.tree_connect_with_dialect_and_caps() as (chan, tree): + max_sz = max(chan.connection.negotiate_response.max_read_size, + chan.connection.negotiate_response.max_write_size) + WriteReadMaxMtu.write_buf = "%" * max_sz + + def pump_credits(self, chan, fh, size): + """ + do small write operations on the file, but request enough credits to + satisfy a write of `size` + """ + n_credits = size / 65536 + (1 if size % 65536 > 0 else 0) + if chan.connection.credits >= n_credits: + return + + self.info("pumping credits from {0} to {1}".format( + chan.connection.credits, + n_credits)) + while chan.connection.credits < n_credits: + with chan.let(credit_request=n_credits): + chan.write(fh, 0, "A") + + def gen_writeread_max_mtu(self, dialect=None, caps=0): + """ + Send the largest write and read the server will accept + """ + + filename = "gen_writeread_max_mtu" + write_resp = read_resp = None + + with self.tree_connect_with_dialect_and_caps(dialect, caps) as (chan, + tree): + max_read_size = chan.connection.negotiate_response.max_read_size + max_write_size = chan.connection.negotiate_response.max_write_size + self.info("Write {0} / Read {1}".format( + max_write_size, + max_read_size)) + self.assertGreaterEqual(len(self.write_buf), max_write_size) + fh = chan.create( + tree, + filename, + access=pike.smb2.GENERIC_ALL | pike.smb2.DELETE, + options=pike.smb2.FILE_DELETE_ON_CLOSE).result() + self.pump_credits(chan, fh, max_write_size) + write_resp = chan.write(fh, 0, self.write_buf[:max_write_size]) + self.assertEqual(write_resp, max_write_size) + self.pump_credits(chan, fh, max_read_size) + read_resp = chan.read(fh, max_read_size, 0) + self.assertBufferEqual(read_resp.tostring(), + self.write_buf[:max_read_size]) + return write_resp, read_resp + + def gen_writeread_over(self, writeover=0, readover=0, dialect=None, caps=0): + """ + Send more than the largest write and read the server will accept + """ + + filename = "gen_writeread_over" + write_resp = read_resp = None + + with self.tree_connect_with_dialect_and_caps(dialect, caps) as (chan, tree): + fh = chan.create( + tree, + filename, + access=pike.smb2.GENERIC_ALL | pike.smb2.DELETE, + options=pike.smb2.FILE_DELETE_ON_CLOSE).result() + if writeover: + max_write_size = chan.connection.negotiate_response.max_write_size + self.assertGreaterEqual(len(self.write_buf), max_write_size) + write_size = max_write_size + writeover + over_buf = "%" * writeover + self.info("Write {0} (over {1})".format(write_size, writeover)) + self.pump_credits(chan, fh, write_size) + write_resp = chan.write(fh, 0, self.write_buf + over_buf) + if readover: + max_read_size = chan.connection.negotiate_response.max_read_size + read_size = max_read_size + readover + self.info("Read {0} (over {1})".format(read_size, readover)) + self.pump_credits(chan, fh, read_size) + read_resp = chan.read(fh, read_size, 0) + return write_resp, read_resp + + @pike.test.RequireDialect(pike.smb2.DIALECT_SMB2_002) + def test_wr_2_002(self): + self.gen_writeread_max_mtu(dialect=pike.smb2.DIALECT_SMB2_002) + + @pike.test.RequireDialect(pike.smb2.DIALECT_SMB2_002) + def test_wr_2_002_writeover1(self): + try: + self.gen_writeread_over(writeover=1, + dialect=pike.smb2.DIALECT_SMB2_002) + self.fail("Writing more than max_write_size didn't raise error") + # special handling is needed here since windows resets the connection + # instead of checking the size and returning an error + except pike.model.ResponseError as err: + if err.response.status not in self.invalid_write_status: + raise + except socket.error as err: + if err.errno != errno.ECONNRESET: + raise + + @pike.test.RequireDialect(pike.smb2.DIALECT_SMB2_002) + def test_wr_2_002_readover1(self): + with self.assert_error(*self.invalid_read_status): + self.gen_writeread_over(readover=1, + dialect=pike.smb2.DIALECT_SMB2_002) + + @pike.test.RequireDialect(pike.smb2.DIALECT_SMB2_1) + def test_wr_2_1(self): + self.gen_writeread_max_mtu(dialect=pike.smb2.DIALECT_SMB2_1) + + @pike.test.RequireDialect(pike.smb2.DIALECT_SMB2_1) + def test_wr_2_1_writeover1(self): + with self.assert_error(*self.invalid_write_status): + self.gen_writeread_over(writeover=1, + dialect=pike.smb2.DIALECT_SMB2_1) + + @pike.test.RequireDialect(pike.smb2.DIALECT_SMB2_1) + def test_wr_2_1_readover1(self): + with self.assert_error(*self.invalid_read_status): + self.gen_writeread_over(readover=1, + dialect=pike.smb2.DIALECT_SMB2_1) + + @pike.test.RequireDialect(pike.smb2.DIALECT_SMB3_0) + def test_wr_3_0(self): + self.gen_writeread_max_mtu(dialect=pike.smb2.DIALECT_SMB3_0) + + @pike.test.RequireDialect(pike.smb2.DIALECT_SMB3_0) + def test_wr_3_0_writeover1(self): + with self.assert_error(*self.invalid_write_status): + self.gen_writeread_over(writeover=1, + dialect=pike.smb2.DIALECT_SMB3_0) + + @pike.test.RequireDialect(pike.smb2.DIALECT_SMB3_0) + def test_wr_3_0_readover1(self): + with self.assert_error(*self.invalid_read_status): + self.gen_writeread_over(readover=1, + dialect=pike.smb2.DIALECT_SMB3_0) + + @pike.test.RequireDialect(pike.smb2.DIALECT_SMB3_0_2) + def test_wr_3_0_2(self): + self.gen_writeread_max_mtu(dialect=pike.smb2.DIALECT_SMB3_0_2) + + @pike.test.RequireDialect(pike.smb2.DIALECT_SMB3_0_2) + def test_wr_3_0_2_writeover1(self): + with self.assert_error(*self.invalid_write_status): + self.gen_writeread_over(writeover=1, + dialect=pike.smb2.DIALECT_SMB3_0_2) + + @pike.test.RequireDialect(pike.smb2.DIALECT_SMB3_0_2) + def test_wr_3_0_2_readover1(self): + with self.assert_error(*self.invalid_read_status): + self.gen_writeread_over(readover=1, + dialect=pike.smb2.DIALECT_SMB3_0_2) + + @pike.test.RequireDialect(pike.smb2.DIALECT_SMB3_0_2) + def test_wr_3_1_1(self): + self.gen_writeread_max_mtu(dialect=pike.smb2.DIALECT_SMB3_1_1) + + @pike.test.RequireDialect(pike.smb2.DIALECT_SMB3_0_2) + def test_wr_3_1_1_writeover1(self): + with self.assert_error(*self.invalid_write_status): + self.gen_writeread_over(writeover=1, + dialect=pike.smb2.DIALECT_SMB3_1_1) + + @pike.test.RequireDialect(pike.smb2.DIALECT_SMB3_0_2) + def test_wr_3_1_1_readover1(self): + with self.assert_error(*self.invalid_read_status): + self.gen_writeread_over(readover=1, + dialect=pike.smb2.DIALECT_SMB3_1_1) + + @pike.test.RequireCapabilities(pike.smb2.SMB2_GLOBAL_CAP_LARGE_MTU) + def test_wr_large_mtu(self): + self.gen_writeread_max_mtu(caps=pike.smb2.SMB2_GLOBAL_CAP_LARGE_MTU) + + @pike.test.RequireCapabilities(pike.smb2.SMB2_GLOBAL_CAP_LARGE_MTU) + def test_wr_large_mtu_writeover1(self): + with self.assert_error(*self.invalid_write_status): + self.gen_writeread_over(writeover=1, + caps=pike.smb2.SMB2_GLOBAL_CAP_LARGE_MTU) + + @pike.test.RequireCapabilities(pike.smb2.SMB2_GLOBAL_CAP_LARGE_MTU) + def test_wr_large_mtu_readover1(self): + with self.assert_error(*self.invalid_read_status): + self.gen_writeread_over(readover=1, + caps=pike.smb2.SMB2_GLOBAL_CAP_LARGE_MTU)