Optimized connection receive buffer allocation

SMBLibrary/Client/ConnectionState.cs

@@ -19,10 +19,10 @@ namespace SMBLibrary.Client
         private Socket m_clientSocket;
         private NBTConnectionReceiveBuffer m_receiveBuffer;
 
-        public ConnectionState(Socket clientSocket, bool isLargeMTU)
+        public ConnectionState(Socket clientSocket)
         {
             m_clientSocket = clientSocket;
-            m_receiveBuffer = new NBTConnectionReceiveBuffer(isLargeMTU);
+            m_receiveBuffer = new NBTConnectionReceiveBuffer();
         }
 
         public Socket ClientSocket

SMBLibrary/Client/SMB1Client.cs

@@ -143,7 +143,7 @@ namespace SMBLibrary.Client
                 return false;
             }
 
-            ConnectionState state = new ConnectionState(m_clientSocket, false);
+            ConnectionState state = new ConnectionState(m_clientSocket);
             NBTConnectionReceiveBuffer buffer = state.ReceiveBuffer;
             m_clientSocket.BeginReceive(buffer.Buffer, buffer.WriteOffset, buffer.AvailableLength, SocketFlags.None, new AsyncCallback(OnClientSocketReceive), state);
             return true;

SMBLibrary/Client/SMB2Client.cs

@@ -136,7 +136,7 @@ namespace SMBLibrary.Client
                 return false;
             }
 
-            ConnectionState state = new ConnectionState(m_clientSocket, true);
+            ConnectionState state = new ConnectionState(m_clientSocket);
             NBTConnectionReceiveBuffer buffer = state.ReceiveBuffer;
             m_clientSocket.BeginReceive(buffer.Buffer, buffer.WriteOffset, buffer.AvailableLength, SocketFlags.None, new AsyncCallback(OnClientSocketReceive), state);
             return true;
@@ -392,6 +392,21 @@ namespace SMBLibrary.Client
 
                 m_availableCredits += command.Header.Credits;
 
+                if (m_transport == SMBTransportType.DirectTCPTransport && command is NegotiateResponse)
+                {
+                    NegotiateResponse negotiateResponse = (NegotiateResponse)command;
+                    if ((negotiateResponse.Capabilities & Capabilities.LargeMTU) > 0)
+                    {
+                        // [MS-SMB2] 3.2.5.1 Receiving Any Message - If the message size received exceeds Connection.MaxTransactSize, the client MUST disconnect the connection.
+                        // Note: Windows clients do not enforce the MaxTransactSize value, we add 256 bytes.
+                        int maxPacketSize = SessionPacket.HeaderLength + (int)Math.Min(negotiateResponse.MaxTransactSize, ClientMaxTransactSize) + 256;
+                        if (maxPacketSize > state.ReceiveBuffer.Buffer.Length)
+                        {
+                            state.ReceiveBuffer.IncreaseBufferSize(maxPacketSize);
+                        }
+                    }
+                }
+
                 // [MS-SMB2] 3.2.5.1.2 - If the MessageId is 0xFFFFFFFFFFFFFFFF, this is not a reply to a previous request,
                 // and the client MUST NOT attempt to locate the request, but instead process it as follows:
                 // If the command field in the SMB2 header is SMB2 OPLOCK_BREAK, it MUST be processed as specified in 3.2.5.19.

SMBLibrary/NetBios/NBTConnectionReceiveBuffer.cs

@@ -18,7 +18,7 @@ namespace SMBLibrary.NetBios
         private int m_bytesInBuffer = 0;
         private int? m_packetLength;
 
-        public NBTConnectionReceiveBuffer(bool isLargeMTU) : this(isLargeMTU ? SessionPacket.MaxDirectTcpPacketLength : SessionPacket.MaxSessionPacketLength)
+        public NBTConnectionReceiveBuffer() : this(SessionPacket.MaxSessionPacketLength)
         {
         }
 
@@ -32,6 +32,17 @@ namespace SMBLibrary.NetBios
             m_buffer = new byte[bufferLength];
         }
 
+        public void IncreaseBufferSize(int bufferLength)
+        {
+            byte[] buffer = new byte[bufferLength];
+            if (m_bytesInBuffer > 0)
+            {
+                Array.Copy(m_buffer, m_readOffset, buffer, 0, m_bytesInBuffer);
+                m_readOffset = 0;
+            }
+            m_buffer = buffer;
+        }
+
         public void SetNumberOfBytesReceived(int numberOfBytesReceived)
         {
             m_bytesInBuffer += numberOfBytesReceived;

SMBLibrary/Server/ConnectionState/ConnectionState.cs

@@ -29,11 +29,11 @@ namespace SMBLibrary.Server
         public SMBDialect Dialect;
         public GSSContext AuthenticationContext;
 
-        public ConnectionState(Socket clientSocket, IPEndPoint clientEndPoint, SMBTransportType transportType, LogDelegate logToServerHandler)
+        public ConnectionState(Socket clientSocket, IPEndPoint clientEndPoint, LogDelegate logToServerHandler)
         {
             m_clientSocket = clientSocket;
             m_clientEndPoint = clientEndPoint;
-            m_receiveBuffer = new NBTConnectionReceiveBuffer(transportType == SMBTransportType.DirectTCPTransport);
+            m_receiveBuffer = new NBTConnectionReceiveBuffer();
             m_sendQueue = new BlockingQueue<SessionPacket>();
             m_creationDT = DateTime.UtcNow;
             m_lastReceiveDT = DateTime.UtcNow;

SMBLibrary/Server/SMB2/NegotiateHelper.cs

@@ -7,6 +7,7 @@
 using System;
 using System.Collections.Generic;
 using SMBLibrary.Authentication.GSSAPI;
+using SMBLibrary.NetBios;
 using SMBLibrary.SMB2;
 using Utilities;
 
@@ -53,6 +54,12 @@ namespace SMBLibrary.Server.SMB2
                 response.MaxTransactSize = ServerMaxTransactSizeLargeMTU;
                 response.MaxReadSize = ServerMaxReadSizeLargeMTU;
                 response.MaxWriteSize = ServerMaxWriteSizeLargeMTU;
+                // [MS-SMB2] 3.3.5.2 Receiving Any Message - If the length of the message exceeds Connection.MaxTransactSize + 256, the server MUST disconnect the connection.
+                int maxPacketSize = SessionPacket.HeaderLength + (int)ServerMaxTransactSize + 256;
+                if (maxPacketSize > state.ReceiveBuffer.Buffer.Length)
+                {
+                    state.ReceiveBuffer.IncreaseBufferSize(maxPacketSize);
+                }
             }
             else
             {
@@ -92,6 +99,12 @@ namespace SMBLibrary.Server.SMB2
                 response.MaxTransactSize = ServerMaxTransactSizeLargeMTU;
                 response.MaxReadSize = ServerMaxReadSizeLargeMTU;
                 response.MaxWriteSize = ServerMaxWriteSizeLargeMTU;
+                // [MS-SMB2] 3.3.5.2 Receiving Any Message - If the length of the message exceeds Connection.MaxTransactSize + 256, the server MUST disconnect the connection.
+                int maxPacketSize = SessionPacket.HeaderLength + (int)ServerMaxTransactSize + 256;
+                if (maxPacketSize > state.ReceiveBuffer.Buffer.Length)
+                {
+                    state.ReceiveBuffer.IncreaseBufferSize(maxPacketSize);
+                }
             }
             else
             {

SMBLibrary/Server/SMBServer.cs

@@ -160,7 +160,7 @@ namespace SMBLibrary.Server
 
             if (acceptConnection)
             {
-                ConnectionState state = new ConnectionState(clientSocket, clientEndPoint, m_transport, Log);
+                ConnectionState state = new ConnectionState(clientSocket, clientEndPoint, Log);
                 state.LogToServer(Severity.Verbose, "New connection request accepted");
                 Thread senderThread = new Thread(delegate()
                 {