From 7d6abeb9e4240a03f7336eb95b8a23e39b109391 Mon Sep 17 00:00:00 2001
From: Robbert-Brand <142212423+Robbert-Brand@users.noreply.github.com>
Date: Mon, 8 Jan 2024 02:03:50 +0100
Subject: [PATCH] Added default timeout on connection.receive() (#257)

* Added default timeout on receive() to prevent disconnect hanging against server that died after connecting.

* Added a default timeout when connecting

* Mistakenly forgot to change back the default of connection.receive()
---
 src/smbclient/_pool.py        | 4 ++--
 src/smbprotocol/connection.py | 4 ++--
 src/smbprotocol/session.py    | 4 ++--
 3 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/src/smbclient/_pool.py b/src/smbclient/_pool.py
index c5c8cf1f..3a7dc505 100644
--- a/src/smbclient/_pool.py
+++ b/src/smbclient/_pool.py
@@ -210,7 +210,7 @@ def dfs_request(tree: TreeConnect, path: str) -> DFSReferralResponse:
     return dfs_response
 
 
-def delete_session(server, port=445, connection_cache=None):
+def delete_session(server, port=445, connection_cache=None, timeout=60):
     """
     Deletes the connection in the connection pool for the server specified. This will also close all sessions
     associated with the connection.
@@ -226,7 +226,7 @@ def delete_session(server, port=445, connection_cache=None):
     connection = connection_cache.get(connection_key, None)
     if connection:
         del connection_cache[connection_key]
-        connection.disconnect(close=True)
+        connection.disconnect(close=True, timeout=timeout)
 
 
 def get_smb_tree(
diff --git a/src/smbprotocol/connection.py b/src/smbprotocol/connection.py
index be06a403..de7fc404 100644
--- a/src/smbprotocol/connection.py
+++ b/src/smbprotocol/connection.py
@@ -931,7 +931,7 @@ def connect(self, dialect=None, timeout=60, preferred_encryption_algos=None, pre
                 elif context_type == NegotiateContextType.SMB2_SIGNING_CAPABILITIES:
                     self.signing_algorithm_id = context["data"]["signing_algorithms"][0]
 
-    def disconnect(self, close=True):
+    def disconnect(self, close=True, timeout=None):
         """
         Closes the connection as well as logs off any of the
         Disconnects the TCP connection and shuts down the socket listener
@@ -943,7 +943,7 @@ def disconnect(self, close=True):
         # We cannot close the session or tree if the socket has been closed.
         if close and self.transport.connected:
             for session in list(self.session_table.values()):
-                session.disconnect(True)
+                session.disconnect(True, timeout=timeout)
 
         log.info("Disconnecting transport connection")
         self.transport.close()
diff --git a/src/smbprotocol/session.py b/src/smbprotocol/session.py
index f521173f..d9a3481f 100644
--- a/src/smbprotocol/session.py
+++ b/src/smbprotocol/session.py
@@ -394,7 +394,7 @@ def connect(self):
             log.info("Verifying the SMB Setup Session signature as auth is successful")
             self.connection.verify_signature(response, self.session_id, force=True)
 
-    def disconnect(self, close=True):
+    def disconnect(self, close=True, timeout=None):
         """
         Logs off the session
 
@@ -418,7 +418,7 @@ def disconnect(self, close=True):
         request = self.connection.send(logoff, sid=self.session_id)
 
         log.info("Session: %s - Receiving Logoff response", self.username)
-        res = self.connection.receive(request)
+        res = self.connection.receive(request, timeout=timeout)
         res_logoff = SMB2Logoff()
         res_logoff.unpack(res["data"].get_value())
         log.debug(res_logoff)