From 9141665ac64bb3c08bb3697a1f0e86fade61efea Mon Sep 17 00:00:00 2001
From: Tal Aloni <tal.aloni.il@gmail.com>
Date: Sat, 4 Mar 2017 14:12:02 +0200
Subject: [PATCH] All matching opened files will now be closed during logoff

---
 .../Server/ConnectionState/SMB1ConnectionState.cs    |  8 +++++++-
 SMBLibrary/Server/ConnectionState/SMB1Session.cs     | 12 ++++++++++++
 .../Server/ConnectionState/SMB2ConnectionState.cs    |  8 +++++++-
 SMBLibrary/Server/ConnectionState/SMB2Session.cs     | 12 ++++++++++++
 4 files changed, 38 insertions(+), 2 deletions(-)

diff --git a/SMBLibrary/Server/ConnectionState/SMB1ConnectionState.cs b/SMBLibrary/Server/ConnectionState/SMB1ConnectionState.cs
index 5f4c73e..1926f06 100644
--- a/SMBLibrary/Server/ConnectionState/SMB1ConnectionState.cs
+++ b/SMBLibrary/Server/ConnectionState/SMB1ConnectionState.cs
@@ -79,7 +79,13 @@ namespace SMBLibrary.Server
 
         public void RemoveSession(ushort userID)
         {
-            m_sessions.Remove(userID);
+            SMB1Session session;
+            m_sessions.TryGetValue(userID, out session);
+            if (session != null)
+            {
+                session.Close();
+                m_sessions.Remove(userID);
+            }
         }
 
         /// <summary>
diff --git a/SMBLibrary/Server/ConnectionState/SMB1Session.cs b/SMBLibrary/Server/ConnectionState/SMB1Session.cs
index 32f8b4d..351c3a2 100644
--- a/SMBLibrary/Server/ConnectionState/SMB1Session.cs
+++ b/SMBLibrary/Server/ConnectionState/SMB1Session.cs
@@ -152,6 +152,18 @@ namespace SMBLibrary.Server
             m_openSearches.Remove(searchHandle);
         }
 
+        /// <summary>
+        /// Free all resources used by this session
+        /// </summary>
+        public void Close()
+        {
+            List<ushort> treeIDList = new List<ushort>(m_connectedTrees.Keys);
+            foreach (ushort treeID in treeIDList)
+            {
+                DisconnectTree(treeID);
+            }
+        }
+
         public ushort UserID
         {
             get
diff --git a/SMBLibrary/Server/ConnectionState/SMB2ConnectionState.cs b/SMBLibrary/Server/ConnectionState/SMB2ConnectionState.cs
index 3840fac..c16b1cb 100644
--- a/SMBLibrary/Server/ConnectionState/SMB2ConnectionState.cs
+++ b/SMBLibrary/Server/ConnectionState/SMB2ConnectionState.cs
@@ -60,7 +60,13 @@ namespace SMBLibrary.Server
 
         public void RemoveSession(ulong sessionID)
         {
-            m_sessions.Remove(sessionID);
+            SMB2Session session;
+            m_sessions.TryGetValue(sessionID, out session);
+            if (session != null)
+            {
+                session.Close();
+                m_sessions.Remove(sessionID);
+            }
         }
 
         public void ClearSessions()
diff --git a/SMBLibrary/Server/ConnectionState/SMB2Session.cs b/SMBLibrary/Server/ConnectionState/SMB2Session.cs
index 71aa942..34e8fe7 100644
--- a/SMBLibrary/Server/ConnectionState/SMB2Session.cs
+++ b/SMBLibrary/Server/ConnectionState/SMB2Session.cs
@@ -152,6 +152,18 @@ namespace SMBLibrary.Server
             m_openSearches.Remove(fileID);
         }
 
+        /// <summary>
+        /// Free all resources used by this session
+        /// </summary>
+        public void Close()
+        {
+            List<uint> treeIDList = new List<uint>(m_connectedTrees.Keys);
+            foreach (uint treeID in treeIDList)
+            {
+                DisconnectTree(treeID);
+            }
+        }
+
         public byte[] SessionKey
         {
             get