From 617e5fb85447d9d0776c83cb36ecf79fd745037a Mon Sep 17 00:00:00 2001
From: Tal Aloni <tal.aloni.il@gmail.com>
Date: Fri, 13 Jan 2017 14:19:42 +0200
Subject: [PATCH] Improved TID (TreeID) allocation mechanism

---
 SMBLibrary/Enums/NTStatus.cs                  |  1 +
 .../ConnectionState/SMB1ConnectionState.cs    | 28 +++++++++++++------
 SMBLibrary/Server/SMB1/TreeConnectHelper.cs   | 18 ++++++++++--
 3 files changed, 35 insertions(+), 12 deletions(-)

diff --git a/SMBLibrary/Enums/NTStatus.cs b/SMBLibrary/Enums/NTStatus.cs
index bf67104..7cd50be 100644
--- a/SMBLibrary/Enums/NTStatus.cs
+++ b/SMBLibrary/Enums/NTStatus.cs
@@ -27,6 +27,7 @@ namespace SMBLibrary
         STATUS_FILE_IS_A_DIRECTORY = 0xC00000BA,
         STATUS_NOT_SUPPORTED = 0xC00000BB,
         STATUS_CANNOT_DELETE = 0xC0000121,
+        STATUS_INSUFF_SERVER_RESOURCES = 0xC0000205,
         
         STATUS_INVALID_SMB = 0x00010002,     // CIFS/SMB1: A corrupt or invalid SMB request was received
         STATUS_SMB_BAD_COMMAND = 0x00160002, // CIFS/SMB1: An unknown SMB command code was received by the server
diff --git a/SMBLibrary/Server/ConnectionState/SMB1ConnectionState.cs b/SMBLibrary/Server/ConnectionState/SMB1ConnectionState.cs
index b44f01a..fb67928 100644
--- a/SMBLibrary/Server/ConnectionState/SMB1ConnectionState.cs
+++ b/SMBLibrary/Server/ConnectionState/SMB1ConnectionState.cs
@@ -89,21 +89,31 @@ namespace SMBLibrary.Server
         /// An open TID MUST be unique within an SMB connection.
         /// The value 0xFFFF MUST NOT be used as a valid TID. All other possible values for TID, including zero (0x0000), are valid.
         /// </summary>
-        private ushort AllocateTreeID()
+        private ushort? AllocateTreeID()
         {
-            while (m_connectedTrees.ContainsKey(m_nextTID) || m_nextTID == 0 || m_nextTID == 0xFFFF)
+            for (ushort offset = 0; offset < UInt16.MaxValue; offset++)
             {
-                m_nextTID++;
+                ushort treeID = (ushort)(m_nextTID + offset);
+                if (treeID == 0 || treeID == 0xFFFF)
+                {
+                    continue;
+                }
+                if (!m_connectedTrees.ContainsKey(treeID))
+                {
+                    m_nextTID = (ushort)(treeID + 1);
+                    return treeID;
+                }
             }
-            ushort treeID = m_nextTID;
-            m_nextTID++;
-            return treeID;
+            return null;
         }
 
-        public ushort AddConnectedTree(string relativePath)
+        public ushort? AddConnectedTree(string relativePath)
         {
-            ushort treeID = AllocateTreeID();
-            m_connectedTrees.Add(treeID, relativePath);
+            ushort? treeID = AllocateTreeID();
+            if (treeID.HasValue)
+            {
+                m_connectedTrees.Add(treeID.Value, relativePath);
+            }
             return treeID;
         }
 
diff --git a/SMBLibrary/Server/SMB1/TreeConnectHelper.cs b/SMBLibrary/Server/SMB1/TreeConnectHelper.cs
index b58eb34..703a4a7 100644
--- a/SMBLibrary/Server/SMB1/TreeConnectHelper.cs
+++ b/SMBLibrary/Server/SMB1/TreeConnectHelper.cs
@@ -1,4 +1,4 @@
-/* Copyright (C) 2014 Tal Aloni <tal.aloni.il@gmail.com>. All rights reserved.
+/* Copyright (C) 2014-2017 Tal Aloni <tal.aloni.il@gmail.com>. All rights reserved.
  * 
  * You can redistribute this program and/or modify it under the terms of
  * the GNU Lesser Public License as published by the Free Software Foundation,
@@ -20,7 +20,13 @@ namespace SMBLibrary.Server.SMB1
             string relativePath = ServerPathUtils.GetRelativeServerPath(request.Path);
             if (String.Equals(relativePath, "\\IPC$", StringComparison.InvariantCultureIgnoreCase))
             {
-                header.TID = state.AddConnectedTree(relativePath);
+                ushort? treeID = state.AddConnectedTree(relativePath);
+                if (!treeID.HasValue)
+                {
+                    header.Status = NTStatus.STATUS_INSUFF_SERVER_RESOURCES;
+                    return new ErrorResponse(CommandName.SMB_COM_TREE_CONNECT_ANDX);
+                }
+                header.TID = treeID.Value;
                 if (isExtended)
                 {
                     return CreateTreeConnectResponseExtended(ServiceName.NamedPipe);
@@ -48,7 +54,13 @@ namespace SMBLibrary.Server.SMB1
                     }
                     else
                     {
-                        header.TID = state.AddConnectedTree(relativePath);
+                        ushort? treeID = state.AddConnectedTree(relativePath);
+                        if (!treeID.HasValue)
+                        {
+                            header.Status = NTStatus.STATUS_INSUFF_SERVER_RESOURCES;
+                            return new ErrorResponse(CommandName.SMB_COM_TREE_CONNECT_ANDX);
+                        }
+                        header.TID = treeID.Value;
                         if (isExtended)
                         {
                             return CreateTreeConnectResponseExtended(ServiceName.DiskShare);