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

---
 .../Server/ConnectionState/SMB1Session.cs     | 19 +++++++++++++--
 .../Server/ConnectionState/SMB2Session.cs     | 24 +++++++++++++------
 SMBLibrary/Server/SMB1/TreeConnectHelper.cs   |  2 +-
 SMBLibrary/Server/SMBServer.SMB2.cs           |  2 +-
 4 files changed, 36 insertions(+), 11 deletions(-)

diff --git a/SMBLibrary/Server/ConnectionState/SMB1Session.cs b/SMBLibrary/Server/ConnectionState/SMB1Session.cs
index 3551eb2..32f8b4d 100644
--- a/SMBLibrary/Server/ConnectionState/SMB1Session.cs
+++ b/SMBLibrary/Server/ConnectionState/SMB1Session.cs
@@ -57,9 +57,24 @@ namespace SMBLibrary.Server
             return share;
         }
 
-        public void RemoveConnectedTree(ushort treeID)
+        public void DisconnectTree(ushort treeID)
         {
-            m_connectedTrees.Remove(treeID);
+            ISMBShare share;
+            m_connectedTrees.TryGetValue(treeID, out share);
+            if (share != null)
+            {
+                List<ushort> fileIDList = new List<ushort>(m_openFiles.Keys);
+                foreach (ushort fileID in fileIDList)
+                {
+                    OpenFileObject openFile = m_openFiles[fileID];
+                    if (openFile.TreeID == treeID)
+                    {
+                        share.FileStore.CloseFile(openFile.Handle);
+                        m_openFiles.Remove(fileID);
+                    }
+                }
+                m_connectedTrees.Remove(treeID);
+            }
         }
 
         public bool IsTreeConnected(ushort treeID)
diff --git a/SMBLibrary/Server/ConnectionState/SMB2Session.cs b/SMBLibrary/Server/ConnectionState/SMB2Session.cs
index d2e2e2c..71aa942 100644
--- a/SMBLibrary/Server/ConnectionState/SMB2Session.cs
+++ b/SMBLibrary/Server/ConnectionState/SMB2Session.cs
@@ -79,14 +79,24 @@ namespace SMBLibrary.Server
             }
         }
 
-        public void RemoveConnectedTree(uint treeID)
+        public void DisconnectTree(uint treeID)
         {
-            m_connectedTrees.Remove(treeID);
-        }
-
-        public void RemoveConnectedTrees()
-        {
-            m_connectedTrees.Clear();
+            ISMBShare share;
+            m_connectedTrees.TryGetValue(treeID, out share);
+            if (share != null)
+            {
+                List<ulong> fileIDList = new List<ulong>(m_openFiles.Keys);
+                foreach (ushort fileID in fileIDList)
+                {
+                    OpenFileObject openFile = m_openFiles[fileID];
+                    if (openFile.TreeID == treeID)
+                    {
+                        share.FileStore.CloseFile(openFile.Handle);
+                        m_openFiles.Remove(fileID);
+                    }
+                }
+                m_connectedTrees.Remove(treeID);
+            }
         }
 
         public bool IsTreeConnected(uint treeID)
diff --git a/SMBLibrary/Server/SMB1/TreeConnectHelper.cs b/SMBLibrary/Server/SMB1/TreeConnectHelper.cs
index 856fcab..0ab6aaa 100644
--- a/SMBLibrary/Server/SMB1/TreeConnectHelper.cs
+++ b/SMBLibrary/Server/SMB1/TreeConnectHelper.cs
@@ -98,7 +98,7 @@ namespace SMBLibrary.Server.SMB1
                 return new ErrorResponse(request.CommandName);
             }
 
-            session.RemoveConnectedTree(header.TID);
+            session.DisconnectTree(header.TID);
             return new TreeDisconnectResponse();
         }
     }
diff --git a/SMBLibrary/Server/SMBServer.SMB2.cs b/SMBLibrary/Server/SMBServer.SMB2.cs
index 78eaa8d..ed15417 100644
--- a/SMBLibrary/Server/SMBServer.SMB2.cs
+++ b/SMBLibrary/Server/SMBServer.SMB2.cs
@@ -158,7 +158,7 @@ namespace SMBLibrary.Server
 
                     if (command is TreeDisconnectRequest)
                     {
-                        session.RemoveConnectedTree(command.Header.TreeID);
+                        session.DisconnectTree(command.Header.TreeID);
                         return new TreeDisconnectResponse();
                     }
                     else if (command is CreateRequest)