Improved open-search management

SMBLibrary/SMBLibrary.csproj

@@ -106,6 +106,7 @@
     <Compile Include="RPC\Structures\SyntaxID.cs" />
     <Compile Include="Server\ConnectionState\ConnectionState.cs" />
     <Compile Include="Server\ConnectionState\OpenFileObject.cs" />
+    <Compile Include="Server\ConnectionState\OpenSearch.cs" />
     <Compile Include="Server\ConnectionState\ProcessStateObject.cs" />
     <Compile Include="Server\ConnectionState\SMB1ConnectionState.cs" />
     <Compile Include="Server\ConnectionState\SMB1Session.cs" />

SMBLibrary/Server/ConnectionState/OpenSearch.cs

@@ -0,0 +1,24 @@
+/* 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,
+ * either version 3 of the License, or (at your option) any later version.
+ */
+using System;
+using System.Collections.Generic;
+using Utilities;
+
+namespace SMBLibrary.Server
+{
+    public class OpenSearch
+    {
+        public List<FileSystemEntry> Entries;
+        public int EnumerationLocation;
+
+        public OpenSearch(List<FileSystemEntry> entries, int enumerationLocation)
+        {
+            Entries = entries;
+            EnumerationLocation = enumerationLocation;
+        }
+    }
+}

SMBLibrary/Server/ConnectionState/SMB1Session.cs

@@ -26,7 +26,7 @@ namespace SMBLibrary.Server
         private Dictionary<ushort, OpenFileObject> m_openFiles = new Dictionary<ushort, OpenFileObject>();
 
         // Key is search handle a.k.a. Search ID
-        public Dictionary<ushort, List<FileSystemEntry>> OpenSearches = new Dictionary<ushort, List<FileSystemEntry>>();
+        private Dictionary<ushort, OpenSearch> m_openSearches = new Dictionary<ushort, OpenSearch>();
         private ushort m_nextSearchHandle = 1;
 
         public SMB1Session(SMB1ConnectionState connection, ushort userID, string userName)
@@ -102,7 +102,7 @@ namespace SMBLibrary.Server
             m_openFiles.Remove(fileID);
         }
 
-        public ushort? AllocateSearchHandle()
+        private ushort? AllocateSearchHandle()
         {
             for (ushort offset = 0; offset < UInt16.MaxValue; offset++)
             {
@@ -111,7 +111,7 @@ namespace SMBLibrary.Server
                 {
                     continue;
                 }
-                if (!OpenSearches.ContainsKey(searchHandle))
+                if (!m_openSearches.ContainsKey(searchHandle))
                 {
                     m_nextSearchHandle = (ushort)(searchHandle + 1);
                     return searchHandle;
@@ -120,12 +120,27 @@ namespace SMBLibrary.Server
             return null;
         }
 
-        public void ReleaseSearchHandle(ushort searchHandle)
+        public ushort? AddOpenSearch(List<FileSystemEntry> entries, int enumerationLocation)
         {
-            if (OpenSearches.ContainsKey(searchHandle))
+            ushort? searchHandle = AllocateSearchHandle();
+            if (searchHandle.HasValue)
             {
-                OpenSearches.Remove(searchHandle);
+                OpenSearch openSearch = new OpenSearch(entries, enumerationLocation);
+                m_openSearches.Add(searchHandle.Value, openSearch);
             }
+            return searchHandle;
+        }
+
+        public OpenSearch GetOpenSearch(ushort searchHandle)
+        {
+            OpenSearch openSearch;
+            m_openSearches.TryGetValue(searchHandle, out openSearch);
+            return openSearch;
+        }
+
+        public void RemoveOpenSearch(ushort searchHandle)
+        {
+            m_openSearches.Remove(searchHandle);
         }
 
         public ushort UserID

SMBLibrary/Server/SMB1/ServerResponseHelper.cs

@@ -45,7 +45,7 @@ namespace SMBLibrary.Server.SMB1
         internal static SMB1Command GetFindClose2Request(SMB1Header header, FindClose2Request request, SMB1ConnectionState state)
         {
             SMB1Session session = state.GetSession(header.UID);
-            session.ReleaseSearchHandle(request.SearchHandle);
+            session.RemoveOpenSearch(request.SearchHandle);
             return new FindClose2Response();
         }
 

SMBLibrary/Server/SMB1/Transaction2SubcommandHelper.cs

@@ -110,7 +110,7 @@ namespace SMBLibrary.Server.SMB1
             int returnCount = findInformationList.Count;
             Transaction2FindFirst2Response response = new Transaction2FindFirst2Response();
             response.SetFindInformationList(findInformationList, header.UnicodeFlag);
-            response.EndOfSearch = (returnCount == entries.Count) && (entries.Count <= subcommand.SearchCount);
+            response.EndOfSearch = (returnCount == entries.Count);
             bool closeAtEndOfSearch = (subcommand.Flags & FindFlags.SMB_FIND_CLOSE_AT_EOS) > 0;
             bool closeAfterRequest = (subcommand.Flags & FindFlags.SMB_FIND_CLOSE_AFTER_REQUEST) > 0;
             // If [..] the search fit within a single response and SMB_FIND_CLOSE_AT_EOS is set in the Flags field,
@@ -123,15 +123,13 @@ namespace SMBLibrary.Server.SMB1
             }
             else
             {
-                ushort? searchHandle = session.AllocateSearchHandle();
+                ushort? searchHandle = session.AddOpenSearch(entries, returnCount);
                 if (!searchHandle.HasValue)
                 {
                     header.Status = NTStatus.STATUS_OS2_NO_MORE_SIDS;
                     return null;
                 }
                 response.SID = searchHandle.Value;
-                entries.RemoveRange(0, returnCount);
-                session.OpenSearches.Add(response.SID, entries);
             }
             return response;
         }
@@ -197,34 +195,38 @@ namespace SMBLibrary.Server.SMB1
         internal static Transaction2FindNext2Response GetSubcommandResponse(SMB1Header header, Transaction2FindNext2Request subcommand, FileSystemShare share, SMB1ConnectionState state)
         {
             SMB1Session session = state.GetSession(header.UID);
-            if (!session.OpenSearches.ContainsKey(subcommand.SID))
+            OpenSearch openSearch = session.GetOpenSearch(subcommand.SID);
+            if (openSearch == null)
             {
                 header.Status = NTStatus.STATUS_INVALID_HANDLE;
                 return null;
             }
 
             bool returnResumeKeys = (subcommand.Flags & FindFlags.SMB_FIND_RETURN_RESUME_KEYS) > 0;
-            List<FileSystemEntry> entries = session.OpenSearches[subcommand.SID];
             FindInformationList findInformationList = new FindInformationList();
-            for (int index = 0; index < entries.Count; index++)
+            for (int index = openSearch.EnumerationLocation; index < openSearch.Entries.Count; index++)
             {
-                FindInformation infoEntry = InfoHelper.FromFileSystemEntry(entries[index], subcommand.InformationLevel, header.UnicodeFlag, returnResumeKeys);
+                FindInformation infoEntry = InfoHelper.FromFileSystemEntry(openSearch.Entries[index], subcommand.InformationLevel, header.UnicodeFlag, returnResumeKeys);
                 findInformationList.Add(infoEntry);
                 if (findInformationList.GetLength(header.UnicodeFlag) > state.GetMaxDataCount(header.PID))
                 {
                     findInformationList.RemoveAt(findInformationList.Count - 1);
                     break;
                 }
+
+                if (findInformationList.Count == subcommand.SearchCount)
+                {
+                    break;
+                }
             }
             int returnCount = findInformationList.Count;
             Transaction2FindNext2Response response = new Transaction2FindNext2Response();
             response.SetFindInformationList(findInformationList, header.UnicodeFlag);
-            entries.RemoveRange(0, returnCount);
-            session.OpenSearches[subcommand.SID] = entries;
-            response.EndOfSearch = (returnCount == entries.Count) && (entries.Count <= subcommand.SearchCount);
+            openSearch.EnumerationLocation += returnCount;
+            response.EndOfSearch = (openSearch.EnumerationLocation == openSearch.Entries.Count);
             if (response.EndOfSearch)
             {
-                session.ReleaseSearchHandle(subcommand.SID);
+                session.RemoveOpenSearch(subcommand.SID);
             }
             return response;
         }