"""SMB 读写传输的隔离行为与请求期限。""" import errno import io import logging import threading import unittest from concurrent.futures import ThreadPoolExecutor from datetime import datetime, timezone from types import SimpleNamespace from unittest.mock import Mock, patch from peach.smb_transport import ( SmbConnection, SmbTransport, _DeadlineConnection, _SmbReader, ) from smbprotocol.exceptions import NoMoreFiles, ObjectNameNotFound, AccessDenied, SMBAuthenticationError from smbprotocol.file_info import FileAttributes from smbprotocol.header import NtStatus, SMB2HeaderResponse from smbprotocol.open import ( CreateDisposition, CreateOptions, Open, ShareAccess, SMB2ReadResponse, ) def field(value): return SimpleNamespace(get_value=lambda: value) class SmbTransportTests(unittest.TestCase): def setUp(self): self.config = SmbConnection("files.invalid", "media", "reader", "secret", timeout=.1) self.conn = Mock() self.conn.transport.connected = True self.conn.max_read_size = 4 self.conn.sequence_window = {"high": 20, "low": 0} self.conn.send.return_value = object() self.conn.disconnect.side_effect = self.disconnect self.tree = Mock() self.tree.is_dfs_share = False self.tree.session.session_id = 1 self.tree.tree_connect_id = 2 self.patches = [ patch("peach.smb_transport._DeadlineConnection", return_value=self.conn), patch("peach.smb_transport.smbclient.register_session", return_value=self.tree.session), patch("peach.smb_transport.TreeConnect", return_value=self.tree), patch("peach.smb_transport.Open", side_effect=self.fd), ] self.mocks = [p.start() for p in self.patches] self.addCleanup(lambda: [p.stop() for p in reversed(self.patches)]) self.transport = SmbTransport(self.config) self.addCleanup(self.transport.close) self.fds = [] self.attrs = {} def disconnect(self, **_): self.conn.transport.connected = False def fd(self, tree, name): fd = Mock() fd.connection = self.conn fd.tree_connect = tree fd.file_name = name fd.connected = True fd.end_of_file = 10 fd.file_attributes = self.attrs.get(name, 0) fd.last_write_time = datetime(2026, 1, 2, 3, 4, 5, 123456, tzinfo=timezone.utc) fd.close.side_effect = lambda: setattr(fd, "connected", False) fd.read.side_effect = lambda offset, length, send: (object(), lambda _: b"0123456789"[offset:offset+length]) self.fds.append(fd) return fd def test_credentials_are_not_in_connection_repr(self): self.assertNotIn("reader", repr(self.config)) self.assertNotIn("secret", repr(self.config)) def test_write_uses_explicit_offsets_and_handles_partial_responses(self): self.conn.max_write_size = 4 with patch.object(self.transport, '_set_info') as set_information: writer = self.transport.open_write('folder/new.mp4') fd = self.fds[-1] fd.write.side_effect = lambda data, offset, send: (object(), lambda _: min(2, len(data))) self.assertEqual(writer.write(b'abcdef'), 6) self.assertEqual([call.args[1] for call in fd.write.call_args_list], [0, 2, 4]) self.assertEqual(writer.stat().st_size, 6) writer.flush() fd.flush.assert_called_once() writer.close() self.assertEqual(set_information.call_args.args[1]['end_of_file'].get_value(), 0) self.assertTrue(all(not fd.connected for fd in self.fds)) def test_rename_keeps_both_parent_chains_and_refuses_replacement(self): with patch.object(self.transport, '_set_info') as set_information: self.transport.rename('a/source.mp4', 'b/target.mp4') fd, information = set_information.call_args.args self.assertEqual(fd.file_name, 'a\\source.mp4') self.assertFalse(information['replace_if_exists'].get_value()) self.assertEqual(information['file_name'].get_value(), 'b\\target.mp4') self.assertTrue(all(not fd.connected for fd in self.fds)) def test_write_refuses_reparse_before_truncation(self): self.attrs['link.mp4'] = FileAttributes.FILE_ATTRIBUTE_REPARSE_POINT with patch.object(self.transport, '_set_info') as set_information: with self.assertRaises(PermissionError): self.transport.open_write('link.mp4') set_information.assert_not_called() def test_write_permission_failure_and_root_protection(self): for action in (lambda: self.transport.unlink(()), lambda: self.transport.mkdir(()), lambda: self.transport.rename('a', ())): with self.assertRaises(PermissionError): action() with patch.object(self.transport, '_set_info', side_effect=AccessDenied()): with self.assertRaises(PermissionError): self.transport.unlink('movie.mp4') self.assertTrue(all(not fd.connected for fd in self.fds)) def test_open_handle_exposes_its_own_size_and_timestamp(self): with self.transport.open('movie.mp4') as reader: info = reader.stat() self.assertEqual(info.st_size, 10) self.assertGreater(info.st_mtime_ns, 0) reader.seek(4) self.assertEqual(reader.stat(), info) self.assertTrue(all(not fd.connected for fd in self.fds)) def test_invalid_paths_never_connect(self): for path in ("../x", "/x", "\\x", "C:x", "x/../y", "x//y", "x:stream", "x/", ("a/b",), ("..",), ("x.",)): with self.subTest(path=path), self.assertRaises(ValueError): self.transport.open(path) self.conn.connect.assert_not_called() def test_stat_preserves_size_time_and_read_sharing(self): result = self.transport.stat(("folder", "film.mp4")) self.assertEqual(result.name, "film.mp4") self.assertEqual(result.st_size, 10) self.assertEqual(result.mtime_ns % 1_000_000_000, 123456000) self.assertEqual(result.st_mtime, result.mtime) self.assertFalse(result.is_dir) parent, file = self.fds self.assertEqual(file.create.call_args.args[3], ShareAccess.FILE_SHARE_READ) self.assertEqual(file.create.call_args.args[4], CreateDisposition.FILE_OPEN) self.assertTrue(file.create.call_args.args[5] & CreateOptions.FILE_OPEN_REPARSE_POINT) parent.close.assert_called_once() file.close.assert_called_once() def test_parent_reparse_is_rejected_before_child_open(self): self.attrs["folder"] = FileAttributes.FILE_ATTRIBUTE_REPARSE_POINT with self.assertRaises(PermissionError): self.transport.open("folder/film.mp4") self.assertEqual(len(self.fds), 1) self.fds[0].close.assert_called_once() def test_final_reparse_is_rejected(self): self.attrs["film.mp4"] = FileAttributes.FILE_ATTRIBUTE_REPARSE_POINT with self.assertRaises(PermissionError): self.transport.open("film.mp4") def test_directory_uses_enumeration_metadata_without_per_entry_open(self): def entries(*_): def entry(name, attributes=0): return {key: field(value) for key, value in { "file_name": name.encode("utf-16-le"), "file_attributes": attributes, "last_write_time": datetime(2026, 1, 1, tzinfo=timezone.utc), "end_of_file": 12, }.items()} return [entry("."), entry("movie.mp4"), entry("folder", FileAttributes.FILE_ATTRIBUTE_DIRECTORY), entry("link", FileAttributes.FILE_ATTRIBUTE_REPARSE_POINT)] original = self.mocks[3].side_effect def directory(tree, name): fd = original(tree, name) fd.query_directory.side_effect = [entries(), NoMoreFiles()] return fd self.mocks[3].side_effect = directory result = self.transport.scandir(()) self.assertEqual([x.name for x in result], ["movie.mp4", "folder", "link"]) self.assertEqual(result[0].size, 12) self.assertTrue(result[1].is_dir) self.assertTrue(result[2].is_symlink) self.assertEqual(len(self.fds), 1) self.assertEqual(self.fds[0].query_directory.call_count, 2) def test_connection_cache_is_private_and_dfs_is_rejected(self): self.tree.is_dfs_share = True with self.assertRaises(PermissionError): self.transport.stat(()) self.assertTrue(self.transport._closed) self.assertEqual(self.transport._cache, {}) self.conn.disconnect.assert_called_with(close=False) def test_missing_and_permission_failures_have_distinct_errno(self): for exception, expected in ((ObjectNameNotFound(), errno.ENOENT), (AccessDenied(), errno.EACCES), (SMBAuthenticationError("credential error"), errno.EACCES)): with self.subTest(expected=expected): original = self.fd def bad_fd(tree, name, original=original, exception=exception): fd = original(tree, name) fd.create.side_effect = exception return fd self.mocks[3].side_effect = bad_fd with self.assertRaises(OSError) as caught: self.transport.stat("missing") self.assertEqual(caught.exception.errno, expected) def test_seek_short_reads_and_eof(self): with self.transport.open("film.mp4") as reader: reader._fd.read.side_effect = lambda offset, length, send: (object(), lambda _: b"0123456789"[offset:offset+min(2, length)]) self.assertEqual(reader.read(5), b"01234") self.assertEqual(reader.tell(), 5) self.assertEqual(reader.seek(-2, io.SEEK_CUR), 3) self.assertEqual(reader.read(2), b"34") self.assertEqual(reader.seek(-2, io.SEEK_END), 8) self.assertEqual(reader.read(), b"89") self.assertEqual(reader.read(), b"") self.assertEqual(reader.seek(100), 100) self.assertEqual(reader.read(), b"") for offset, whence in ((-1, 0), (0, 8)): with self.assertRaises(ValueError): reader.seek(offset, whence) with self.assertRaises(ValueError): reader.read() def test_independent_handles_keep_offsets_and_parent_locks(self): first = self.transport.open("folder/film.mp4") second = self.transport.open("folder/film.mp4") self.assertIsNot(first._fd, second._fd) self.assertEqual(first.read(3), b"012") self.assertEqual(second.read(2), b"01") first.close() self.assertEqual(second.read(2), b"23") second.close() self.assertTrue(all(not fd.connected for fd in self.fds)) def test_parallel_handles_read_separate_offsets(self): first = self.transport.open("film.mp4") second = self.transport.open("film.mp4") second.seek(5) barrier = threading.Barrier(2, timeout=2) def read_fd(offset, length, send): def receive(_): barrier.wait() return b"0123456789"[offset:offset+length] return object(), receive first._fd.read.side_effect = second._fd.read.side_effect = read_fd with ThreadPoolExecutor() as executor: a = executor.submit(first.read, 3) b = executor.submit(second.read, 3) self.assertEqual(a.result(), b"012") self.assertEqual(b.result(), b"567") self.assertEqual((first.tell(), second.tell()), (3, 8)) first.close() second.close() def test_premature_empty_read_is_failure_and_disconnects(self): reader = self.transport.open("film.mp4") reader._fd.read.side_effect = lambda *args, **kwargs: (object(), lambda _: b"") with self.assertRaises(OSError) as caught: reader.read(1) self.assertEqual(caught.exception.errno, errno.EIO) self.assertTrue(self.transport._closed) def test_close_during_read_aborts_connection(self): reader = self.transport.open("film.mp4") reader._lock.acquire() reader.close() reader._lock.release() self.assertTrue(reader.closed) self.conn.disconnect.assert_called_with(close=False) def test_transport_close_does_not_reconnect(self): self.transport.stat(()) self.transport.close() with self.assertRaises(OSError): self.transport.stat(()) self.conn.connect.assert_called_once() class SmbDeadlineTests(unittest.TestCase): def connection(self): conn = _DeadlineConnection(SmbConnection("files.invalid", "media", timeout=.05)) conn._check_worker_running = Mock() conn.disconnect = Mock() return conn def request(self, status=NtStatus.STATUS_SUCCESS): response = SMB2HeaderResponse() response["status"] = status request = Mock() request.message = {"message_id": field(1)} request.response = response request.response_event = threading.Event() request.response_event_lock = threading.Lock() return request def test_timeout_cancels_and_disconnects_without_background_read(self): conn, request = self.connection(), self.request() with self.assertRaises(TimeoutError): conn.receive(request) request.cancel.assert_called_once() conn.disconnect.assert_called_once_with(close=False) def test_continuous_pending_does_not_extend_deadline(self): conn, request = self.connection(), self.request(NtStatus.STATUS_PENDING) request.response_event = Mock() request.response_event.wait.return_value = True with patch("peach.smb_transport.time.monotonic", side_effect=[10, 10.01, 10.03, 10.06]): with self.assertRaises(TimeoutError): conn.receive(request) waits = [call.args[0] for call in request.response_event.wait.call_args_list] self.assertLess(waits[1], waits[0]) request.cancel.assert_called_once() conn.disconnect.assert_called_once_with(close=False) def test_cancel_failure_still_disconnects(self): conn, request = self.connection(), self.request() request.cancel.side_effect = OSError("offline") with self.assertRaises(TimeoutError): conn.receive(request) conn.disconnect.assert_called_once_with(close=False) def test_receive_uses_deadline_set_before_send(self): conn, request = self.connection(), self.request() request._peach_deadline = 9 with patch("peach.smb_transport.time.monotonic", return_value=10): with self.assertRaises(TimeoutError): conn.receive(request) request.cancel.assert_called_once() conn.disconnect.assert_called_once_with(close=False) def test_authentication_uses_deadline_before_tree_connect(self): conn, request = self.connection(), self.request() conn.transport = SimpleNamespace(connected=True) config = SmbConnection("files.invalid", "media", "reader", "secret", timeout=.05) transport = SmbTransport(config) with patch("peach.smb_transport._DeadlineConnection", return_value=conn), \ patch.object(conn, "connect"), \ patch("smbclient._pool.Session.connect", autospec=True, side_effect=lambda session: session.connection.receive(request)), \ patch("peach.smb_transport.TreeConnect") as tree: with self.assertRaises(TimeoutError): transport.stat(()) request.cancel.assert_called_once() tree.assert_not_called() self.assertTrue(transport._closed) self.assertEqual(transport._cache, {}) def test_sdk_authentication_details_are_below_logging_threshold(self): for namespace in ("smbprotocol.session", "smbprotocol.connection", "smbclient._pool"): with self.subTest(namespace=namespace): self.assertFalse(logging.getLogger(namespace).isEnabledFor(logging.INFO)) self.assertFalse(logging.getLogger(namespace).isEnabledFor(logging.DEBUG)) def test_response_parser_never_resolves_symlinks(self): conn, request = self.connection(), self.request() request.response_event.set() conn.outstanding_requests[request.message["message_id"].get_value()] = request # 复用真实库的接收解析路径，只隔离网络事件与线程状态。 self.assertIs(conn.receive(request, resolve_symlinks=True), request.response) def test_real_open_read_parser_receives_explicit_offset(self): conn = Mock() conn.transport.connected = True conn.max_read_size = 65536 conn.sequence_window = {"high": 10, "low": 0} tree = SimpleNamespace(session=SimpleNamespace(connection=conn, session_id=7, username="reader"), tree_connect_id=9, share_name="media") fd = Open(tree, "film.mp4") fd.end_of_file = 6 fd.file_id = b"f" * 16 packet = SMB2ReadResponse() packet["buffer"] = b"def" response = SMB2HeaderResponse() response["data"] = packet.pack() conn.receive.return_value = response transport = SimpleNamespace(_closed=False, close=Mock()) reader = _SmbReader(transport, fd, []) reader.seek(3) self.assertEqual(reader.read(3), b"def") message = conn.send.call_args.args[0] self.assertEqual(message["offset"].get_value(), 3) self.assertEqual(message["length"].get_value(), 3) self.assertEqual(reader.tell(), 6) if __name__ == "__main__": unittest.main()