Land #3738, SMBServer file descriptor updates

lib/msf/core/exploit/smb.rb

@@ -661,6 +661,16 @@ module Exploit::Remote::SMBServer
       [
         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
@@ -685,16 +695,59 @@ module Exploit::Remote::SMBServer
 
   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)
 
@@ -736,10 +789,11 @@ module Exploit::Remote::SMBServer
       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']
@@ -784,6 +838,37 @@ module Exploit::Remote::SMBServer
     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