Repository: jborean93/smbprotocol

Issue:
On disconnectt wait for msg_worker thread to end.
This should not be an issue in production.

It's just something to help with the automated tests.

For the automated tests, I have some extra checks to make sure the tests are not generating "stray" threads.

Since smbprotocol is not waiting for the `msg_thread-` to end on disconnect my cleanup checks is "flaky" and sometimes it reports that the `msg_worker-` thread was left unclosed by the test.

The `msg_worker` will enventuall close

It would be nice to have something like `self._t_worker.join(timeout=5)` at the end of `smbprotocol.connection.Connection.disconnect`

https://github.com/jborean93/smbprotocol/blob/37512ee0648ad64f98755833382fea790d9b2df6/src/smbprotocol/connection.py#L934-L949

Thanks

Hints from issue comments:
@adiroiban: Let me know if this is something you would like to have in the code and I can send a PR.

Also, have you seen my comment from https://github.com/jborean93/smbprotocol/pull/80#issuecomment-1824970873  this is another case in which the automated tests are failing.@jborean93: Feel free to open a PR, I started a mini refactor of the TCP transport stack in the branch https://github.com/jborean93/smbprotocol/tree/socket-refactor. This was in preparation of supporting other transports like QUIC and to try and remove of the multithreading complexity that currently exists. I still plan on getting back to it at some point but I don't know when so if you have a PR to try and improve the current situation I'm definitely all for it.

Solution patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 9257879..5ff1bfb 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,7 @@
 * Ensure `DateTimeField` values are set to `UTC` timezones as FILETIME values are in UTC
 * Stop using `datetime.datetime.utcfromtimestamp()` as it has been deprecated
 * Added default timeout for disconnect operations for 60 seconds to ensure the process doesn't hang forever when closing a broken connection
+* `smbprotocol.connection.Connection.disconnect()` now waits (with a timeout) for the message processing threads to be stopped before returning.
 
 ## 1.12.0 - 2023-11-09
 
diff --git a/README.md b/README.md
index 3edd014..09f475e 100644
--- a/README.md
+++ b/README.md
@@ -222,7 +222,10 @@ To this module, you need to install some pre-requisites first. This can be done
 by running;
 
 ```bash
+# Install in current environment.
+# Recommend to have virtual environment installed at .venv path.
 pip install -r requirements-dev.txt
+pip install -e .
 
 # you can also run tox by installing tox
 pip install tox
@@ -231,12 +234,21 @@ pip install tox
 From there to run the basic tests run;
 
 ```bash
-py.test -v --pep8 --cov smbprotocol --cov-report term-missing
+py.test -v --cov smbprotocol --cov-report term-missing
 
-# or with tox 2.7, 2.7, 3.4, 3.5, and 3.6
+# or with tox for dedicated virtual environments and multiple Python versions.
 tox
 ```
 
+Before sending the code for review, besides making sure all the test pass,
+check that the code complies with the coding standards:
+
+```bash
+source ./build_helpers/lib.sh
+
+lib::sanity::run
+```
+
 There are extra tests that only run when certain environment variables are set.
 To run these tests set the following variables;
 
diff --git a/src/smbprotocol/connection.py b/src/smbprotocol/connection.py
index de7fc40..423bbd1 100644
--- a/src/smbprotocol/connection.py
+++ b/src/smbprotocol/connection.py
@@ -820,6 +820,9 @@ class Connection:
         # Keep track of the message processing thread's potential traceback that it may raise.
         self._t_exc = None
 
+        # The thread that will handle message processing.
+        self._t_worker = None
+
     def connect(self, dialect=None, timeout=60, preferred_encryption_algos=None, preferred_signing_algos=None):
         """
         Will connect to the target server and negotiate the capabilities
@@ -861,11 +864,11 @@ class Connection:
         log.info("Setting up transport connection")
         self.transport = Tcp(self.server_name, self.port, timeout)
         self.transport.connect()
-        t_worker = threading.Thread(
+        self._t_worker = threading.Thread(
             target=self._process_message_thread, name=f"msg_worker-{self.server_name}:{self.port}"
         )
-        t_worker.daemon = True
-        t_worker.start()
+        self._t_worker.daemon = True
+        self._t_worker.start()
 
         log.info("Starting negotiation with SMB server")
         enc_algos = preferred_encryption_algos or [
@@ -947,6 +950,8 @@ class Connection:
 
         log.info("Disconnecting transport connection")
         self.transport.close()
+        if self._t_worker:
+            self._t_worker.join(timeout=2)
 
     def send(
         self, message, sid=None, tid=None, credit_request=None, message_id=None, async_id=None, force_signature=False
diff --git a/src/smbprotocol/transport.py b/src/smbprotocol/transport.py
index 052b43a..2288239 100644
--- a/src/smbprotocol/transport.py
+++ b/src/smbprotocol/transport.py
@@ -65,7 +65,7 @@ class Tcp:
                 try:
                     self._sock = socket.create_connection((self.server, self.port), timeout=self.timeout)
                 except (OSError, socket.gaierror) as err:
-                    raise ValueError(f"Failed to connect to '{self.server}:{self.port}'") from err
+                    raise ValueError(f"Failed to connect to '{self.server}:{self.port}': {err}") from err
                 self._sock.settimeout(None)  # Make sure the socket is in blocking mode.
                 self.connected = True
 


Test patch:
diff --git a/tests/test_connection.py b/tests/test_connection.py
index fe61232..47e3202 100644
--- a/tests/test_connection.py
+++ b/tests/test_connection.py
@@ -2,6 +2,7 @@
 # MIT License (see LICENSE or https://opensource.org/licenses/MIT)
 
 import os
+import threading
 import uuid
 from datetime import datetime, timezone
 
@@ -34,6 +35,22 @@ from smbprotocol.ioctl import SMB2IOCTLRequest
 from smbprotocol.session import Session
 
 
+@pytest.fixture
+def connected_session(smb_real):
+    """
+    Return a (connection, session) tuple that is already connected to the
+    SMB server and that will disconnect at the end.
+    """
+    connection = Connection(uuid.uuid4(), smb_real[2], smb_real[3])
+    connection.connect()
+    session = Session(connection, smb_real[0], smb_real[1])
+    session.connect()
+
+    yield connection, session
+
+    connection.disconnect(True)
+
+
 class TestSMB2NegotiateRequest:
     def test_create_message(self):
         message = SMB2NegotiateRequest()
@@ -1183,6 +1200,9 @@ class TestConnection:
             connection.disconnect()
 
     def test_connection_echo(self, smb_real):
+        """
+        SMB2 Echo can be used to request extra credits from the server.
+        """
         connection = Connection(uuid.uuid4(), smb_real[2], smb_real[3])
         connection.connect()
         session = Session(connection, smb_real[0], smb_real[1])
@@ -1278,3 +1298,39 @@ class TestConnection:
         actual = connection._encrypt(b"\x01\x02\x03\x04", session)
         assert isinstance(actual, SMB2TransformHeader)
         assert actual.pack() == expected.pack()
+
+    def test_connection_echo_more(self, connected_session):
+        """
+        When requesting more credit from the server, it will return the
+        allocated credits.
+        """
+        connection, session = connected_session
+
+        actual = connection.echo(sid=session.session_id, credit_request=65535)
+
+        assert actual == 8129
+
+    def test_disconnect_wait_for_thread(self, connected_session):
+        """
+        Disconnect will continue after the message thread is stopped.
+        """
+        connection, _ = connected_session
+        msg_thread = f"msg_worker-{connection.server_name}:{connection.port}"
+
+        current_threads = [t.name for t in threading.enumerate()]
+        assert msg_thread in current_threads
+
+        connection.disconnect()
+
+        current_threads = [t.name for t in threading.enumerate()]
+        assert msg_thread not in current_threads
+
+    def test_disconnect_already_disconnected(self, connected_session):
+        """
+        Disconnect can be called multiple times, and is a noop if
+        already disconnected.
+        """
+        connection, _ = connected_session
+
+        connection.disconnect()
+        connection.disconnect()
diff --git a/tests/test_transport.py b/tests/test_transport.py
index 960350c..0108dea 100644
--- a/tests/test_transport.py
+++ b/tests/test_transport.py
@@ -72,8 +72,15 @@ class TestTcp:
         )
 
     def test_invalid_host(self):
+        """
+        Raises ValueError when failing to connect to the remote server.
+
+        The error message contains the low-level OS error details.
+        """
         tcp = Tcp("fake-host", 445)
-        with pytest.raises(ValueError, match=re.escape("Failed to connect to 'fake-host:445'")):
+        # We just check for OSError marker, as the actual error message
+        # might be different based on current OS.
+        with pytest.raises(ValueError, match=r"Failed to connect to 'fake-host:445': \[Errno .*"):
             tcp.connect()