Restore rapid7#3738

Author: jvazquez-r7

lib/msf/core/exploit/smb/server.rb

@@ -18,8 +18,18 @@ def initialize(info = {})
         deregister_options('SSL', 'SSLCert')
         register_options(
           [
-            OptPort.new('SRVPORT', [ true, "The local port to listen on.", 445 ])
+            OptPort.new('SRVPORT', [ true, 'The local port to listen on.', 445 ])
           ], self.class)
+
+        register_advanced_options(
+          [
+            OptInt.new('SMBServerMaximumBuffer', [ true, 'The maximum number of data in megabytes to buffer', 2 ]),
+            OptInt.new('SMBServerIdleTimeout', [ true, 'The maximum amount of time to keep an idle session open in seconds', 120 ])
+          ], self.class)
+
+        @smb_server_last_pool_sweep = Time.now.to_f
+        @smb_server_pool_mutex = Mutex.new
+        @smb_server_request_counter = 0
       end
 
       def setup
@@ -44,16 +54,58 @@ def on_client_close(client)
 
       def smb_conn(c)
         @state[c] = {:name => "#{c.peerhost}:#{c.peerport}", :ip => c.peerhost, :port => c.peerport}
+        smb_pool_update(c)
       end
 
       def smb_stop(c)
+        # Make sure the socket is closed
+        begin
+          c.close
+            # Handle any number of errors that a double-close or failed shutdown can trigger
+        rescue ::IOError, ::EOFError,
+          ::Errno::ECONNRESET, ::Errno::ENOTCONN, ::Errno::ECONNABORTED,
+          ::Errno::ETIMEDOUT, ::Errno::ENETRESET, ::Errno::ESHUTDOWN
+        end
+
+        # Delete the state table entry
         @state.delete(c)
       end
 
       def smb_recv(c)
         smb = @state[c]
         smb[:data] ||= ''
-        smb[:data] << c.get_once
+
+        buff = ''
+        begin
+          buff = c.get_once(-1, 0.25)
+            # Handle any number of errors that a read can trigger depending on socket state
+        rescue ::IOError, ::EOFError,
+          ::Errno::ECONNRESET, ::Errno::ENOTCONN, ::Errno::ECONNABORTED,
+          ::Errno::ETIMEDOUT, ::Errno::ENETRESET, ::Errno::ESHUTDOWN
+          vprint_status("Dropping connection from #{smb[:name]} due to exception: #{$!.class} #{$!}")
+          smb_stop(c)
+          return
+        end
+
+        # The client said it had data, but lied, kill the session
+        unless buff and buff.length > 0
+          vprint_status("Dropping connection from #{smb[:name]} due to empty payload...")
+          smb_stop(c)
+          return
+        end
+
+        # Append the new data to the buffer
+        smb[:data] << buff
+
+        # Prevent a simplistic DoS if the buffer is too big
+        if smb[:data].length > (1024*1024*datastore['SMBServerMaximumBuffer'])
+          vprint_status("Dropping connection from #{smb[:name]} due to oversized buffer of #{smb[:data].length} bytes...")
+          smb_stop(c)
+          return
+        end
+
+        # Update the last-seen timestamp and purge old entries
+        smb_pool_update(c)
 
         while(smb[:data].length > 0)
 
@@ -95,10 +147,11 @@ def smb_recv(c)
           pkt = CONST::SMB_BASE_PKT.make_struct
           pkt.from_s(buff)
 
-          # Only response to requests, ignore server replies
+          # Only respond to requests, ignore server replies
           if (pkt['Payload']['SMB'].v['Flags1'] & 128 != 0)
-            print_status("Ignoring server response from #{smb[:name]}")
-            next
+            vprint_status("Dropping connection from #{smb[:name]} due to missing client request flag")
+            smb_stop(c)
+            return
           end
 
           cmd = pkt['Payload']['SMB'].v['Command']
@@ -149,6 +202,38 @@ def smb_error(cmd, c, errorclass, esn = false)
         pkt['Payload']['SMB'].v['ErrorClass'] = errorclass
         c.put(pkt.to_s)
       end
+
+      # Update the last-seen timestamp and purge old entries
+      def smb_pool_update(c)
+
+        @state[c][:last_action] = Time.now.to_f
+        @smb_server_request_counter += 1
+
+        unless @smb_server_request_counter % 100 == 0 ||
+          @smb_server_last_pool_sweep + datastore['SMBServerIdleTimeout'].to_f < Time.now.to_f
+          return
+        end
+
+        # Synchronize pool sweeps in case we move to threaded services
+        @smb_server_pool_mutex.synchronize do
+          purge_list = []
+
+          @smb_server_last_pool_sweep = Time.now.to_f
+
+          @state.keys.each do |sc|
+            if @state[sc][:last_action] + datastore['SMBServerIdleTimeout'].to_f < Time.now.to_f
+              purge_list << sc
+            end
+          end
+
+          # Purge any idle connections to rescue file descriptors
+          purge_list.each do |sc|
+            vprint_status("Dropping connection from #{@state[sc][:name]} due to idle timeout...")
+            smb_stop(sc)
+          end
+        end
+      end
+
     end
   end