Improved UID (User ID) allocation mechanism

SMBLibrary/Enums/NTStatus.cs

@@ -26,6 +26,7 @@ namespace SMBLibrary
         STATUS_MEDIA_WRITE_PROTECTED = 0xC00000A2,
         STATUS_FILE_IS_A_DIRECTORY = 0xC00000BA,
         STATUS_NOT_SUPPORTED = 0xC00000BB,
+        STATUS_TOO_MANY_SESSIONS = 0xC00000CE,
         STATUS_TOO_MANY_OPENED_FILES = 0xC000011F,
         STATUS_CANNOT_DELETE = 0xC0000121,
         STATUS_INSUFF_SERVER_RESOURCES = 0xC0000205,

SMBLibrary/Server/ConnectionState/SMB1ConnectionState.cs

@@ -45,21 +45,31 @@ namespace SMBLibrary.Server
         /// An open UID MUST be unique within an SMB connection.
         /// The value of 0xFFFE SHOULD NOT be used as a valid UID. All other possible values for a UID, excluding zero (0x0000), are valid.
         /// </summary>
-        private ushort AllocateUserID()
+        private ushort? AllocateUserID()
         {
-            while (m_connectedUsers.ContainsKey(m_nextUID) || m_nextUID == 0 || m_nextUID == 0xFFFE || m_nextUID == 0xFFFF)
+            for (ushort offset = 0; offset < UInt16.MaxValue; offset++)
             {
-                m_nextUID++;
+                ushort userID = (ushort)(m_nextUID + offset);
+                if (userID == 0 || userID == 0xFFFE || userID == 0xFFFF)
+                {
+                    continue;
+                }
+                if (!m_connectedUsers.ContainsKey(userID))
+                {
+                    m_nextUID = (ushort)(userID + 1);
+                    return userID;
+                }
             }
-            ushort userID = m_nextUID;
-            m_nextUID++;
-            return userID;
+            return null;
         }
 
-        public ushort AddConnectedUser(string userName)
+        public ushort? AddConnectedUser(string userName)
         {
-            ushort userID = AllocateUserID();
-            m_connectedUsers.Add(userID, userName);
+            ushort? userID = AllocateUserID();
+            if (userID.HasValue)
+            {
+                m_connectedUsers.Add(userID.Value, userName);
+            }
             return userID;
         }
 

SMBLibrary/Server/SMB1/SessionSetupHelper.cs

@@ -37,14 +37,26 @@ namespace SMBLibrary.Server.SMB1
 
             if (user != null)
             {
+                ushort? userID = state.AddConnectedUser(user.AccountName);
+                if (!userID.HasValue)
+                {
+                    header.Status = NTStatus.STATUS_TOO_MANY_SESSIONS;
+                    return new ErrorResponse(CommandName.SMB_COM_SESSION_SETUP_ANDX);
+                }
+                header.UID = userID.Value;
                 response.PrimaryDomain = request.PrimaryDomain;
-                header.UID = state.AddConnectedUser(user.AccountName);
             }
             else if (users.EnableGuestLogin)
             {
+                ushort? userID = state.AddConnectedUser("Guest");
+                if (!userID.HasValue)
+                {
+                    header.Status = NTStatus.STATUS_TOO_MANY_SESSIONS;
+                    return new ErrorResponse(CommandName.SMB_COM_SESSION_SETUP_ANDX);
+                }
+                header.UID = userID.Value;
                 response.Action = SessionSetupAction.SetupGuest;
                 response.PrimaryDomain = request.PrimaryDomain;
-                header.UID = state.AddConnectedUser("Guest");
             }
             else
             {
@@ -110,15 +122,26 @@ namespace SMBLibrary.Server.SMB1
                     return new ErrorResponse(CommandName.SMB_COM_SESSION_SETUP_ANDX);
                 }
 
-
                 if (user != null)
                 {
-                    header.UID = state.AddConnectedUser(user.AccountName);
+                    ushort? userID = state.AddConnectedUser(user.AccountName);
+                    if (!userID.HasValue)
+                    {
+                        header.Status = NTStatus.STATUS_TOO_MANY_SESSIONS;
+                        return new ErrorResponse(CommandName.SMB_COM_SESSION_SETUP_ANDX);
+                    }
+                    header.UID = userID.Value;
                 }
                 else if (users.EnableGuestLogin)
                 {
+                    ushort? userID = state.AddConnectedUser("Guest");
+                    if (!userID.HasValue)
+                    {
+                        header.Status = NTStatus.STATUS_TOO_MANY_SESSIONS;
+                        return new ErrorResponse(CommandName.SMB_COM_SESSION_SETUP_ANDX);
+                    }
+                    header.UID = userID.Value;
                     response.Action = SessionSetupAction.SetupGuest;
-                    header.UID = state.AddConnectedUser("Guest");
                 }
                 else
                 {