rex-core 0.1.1 → 0.1.2

data/lib/rex/io/stream_abstraction.rb

@@ -0,0 +1,32 @@
+# -*- coding: binary -*-
+
+require 'rex/io/socket_abstraction'
+
+module Rex
+module IO
+
+###
+#
+# This class provides an abstraction to a stream based
+# connection through the use of a streaming socketpair.
+#
+###
+module StreamAbstraction
+  include Rex::IO::SocketAbstraction
+
+  #
+  # This method creates a streaming socket pair and initializes it.
+  #
+  def initialize_abstraction
+    self.lsock, self.rsock = Rex::Socket.tcp_socket_pair()
+    self.lsock.extend(Rex::IO::Stream)
+    self.lsock.extend(Ext)
+    self.rsock.extend(Rex::IO::Stream)
+
+    self.monitor_rsock("StreamMonitorRemote")
+  end
+
+end
+
+end; end
+

data/lib/rex/io/stream_server.rb

@@ -0,0 +1,221 @@
+# -*- coding: binary -*-
+require 'thread'
+
+module Rex
+module IO
+
+###
+#
+# This mixin provides the framework and interface for implementing a streaming
+# server that can listen for and accept stream client connections.  Stream
+# servers extend this class and are required to implement the following
+# methods:
+#
+#   accept
+#   fd
+#
+###
+module StreamServer
+
+  ##
+  #
+  # Abstract methods
+  #
+  ##
+
+  ##
+  #
+  # Default server monitoring and client management implementation follows
+  # below.
+  #
+  ##
+
+  #
+  # This callback is notified when a client connects.
+  #
+  def on_client_connect(client)
+    if (on_client_connect_proc)
+      on_client_connect_proc.call(client)
+    end
+  end
+
+  #
+  # This callback is notified when a client connection has data that needs to
+  # be processed.
+  #
+  def on_client_data(client)
+    if (on_client_data_proc)
+      on_client_data_proc.call(client)
+    end
+  end
+
+  #
+  # This callback is notified when a client connection has closed.
+  #
+  def on_client_close(client)
+    if (on_client_close_proc)
+      on_client_close_proc.call(client)
+    end
+  end
+
+  #
+  # Start monitoring the listener socket for connections and keep track of
+  # all client connections.
+  #
+  def start
+    self.clients = []
+    self.client_waiter = ::Queue.new
+
+    self.listener_thread = Rex::ThreadFactory.spawn("StreamServerListener", false) {
+      monitor_listener
+    }
+    self.clients_thread = Rex::ThreadFactory.spawn("StreamServerClientMonitor", false) {
+      monitor_clients
+    }
+  end
+
+  #
+  # Terminates the listener monitoring threads and closes all active clients.
+  #
+  def stop
+    self.listener_thread.kill
+    self.clients_thread.kill
+
+    self.clients.each { |cli|
+      close_client(cli)
+    }
+  end
+
+  #
+  # This method closes a client connection and cleans up the resources
+  # associated with it.
+  #
+  def close_client(client)
+    if (client)
+      clients.delete(client)
+
+      begin
+        client.close
+      rescue IOError
+      end
+    end
+  end
+
+  #
+  # This method waits on the server listener thread
+  #
+  def wait
+    self.listener_thread.join if self.listener_thread
+  end
+
+  ##
+  #
+  # Callback procedures.
+  #
+  ##
+
+  #
+  # This callback procedure can be set and will be called when new clients
+  # connect.
+  #
+  attr_accessor :on_client_connect_proc
+  #
+  # This callback procedure can be set and will be called when clients
+  # have data to be processed.
+  #
+  attr_accessor :on_client_data_proc
+  #
+  # This callback procedure can be set and will be called when a client
+  # disconnects from the server.
+  #
+  attr_accessor :on_client_close_proc
+
+  attr_accessor :clients # :nodoc:
+  attr_accessor :listener_thread, :clients_thread # :nodoc:
+  attr_accessor :client_waiter
+
+protected
+
+  #
+  # This method monitors the listener socket for new connections and calls
+  # the +on_client_connect+ callback routine.
+  #
+  def monitor_listener
+
+    while true
+      begin
+        cli = accept
+        if not cli
+          elog("The accept() returned nil in stream server listener monitor:  #{fd.inspect}")
+          ::IO.select(nil, nil, nil, 0.10)
+          next
+        end
+
+        # Append to the list of clients
+        self.clients << cli
+
+        # Initialize the connection processing
+        on_client_connect(cli)
+
+        # Notify the client monitor
+        self.client_waiter.push(cli)
+
+      # Skip exceptions caused by accept() [ SSL ]
+      rescue ::EOFError, ::Errno::ECONNRESET, ::Errno::ENOTCONN, ::Errno::ECONNABORTED
+      rescue ::Interrupt
+        raise $!
+      rescue ::Exception
+        elog("Error in stream server server monitor: #{$!}")
+        rlog(ExceptionCallStack)
+        break
+      end
+    end
+  end
+
+  #
+  # This method monitors client connections for data and calls the
+  # +on_client_data+ routine when new data arrives.
+  #
+  def monitor_clients
+    begin
+
+      # Wait for a notify if our client list is empty
+      if (clients.length == 0)
+        self.client_waiter.pop
+        next
+      end
+
+      sd = Rex::ThreadSafe.select(clients, nil, nil, nil)
+
+      sd[0].each { |cfd|
+        begin
+          on_client_data(cfd)
+        rescue ::EOFError, ::Errno::ECONNRESET, ::Errno::ENOTCONN, ::Errno::ECONNABORTED
+          on_client_close(cfd)
+          close_client(cfd)
+        rescue ::Interrupt
+          raise $!
+        rescue ::Exception
+          close_client(cfd)
+          elog("Error in stream server client monitor: #{$!}")
+          rlog(ExceptionCallStack)
+
+        end
+      }
+
+    rescue ::Rex::StreamClosedError => e
+      # Remove the closed stream from the list
+      clients.delete(e.stream)
+    rescue ::Interrupt
+      raise $!
+    rescue ::Exception
+      elog("Error in stream server client monitor: #{$!}")
+      rlog(ExceptionCallStack)
+    end while true
+  end
+
+end
+
+end
+end
+

data/lib/rex/sync.rb

@@ -0,0 +1,6 @@
+# -*- coding: binary -*-
+
+require 'rex/sync/thread_safe'
+require 'rex/sync/ref'
+require 'rex/sync/read_write_lock'
+require 'rex/sync/event'

data/lib/rex/sync/event.rb

@@ -0,0 +1,85 @@
+# -*- coding: binary -*-
+require 'thread'
+
+module Rex
+module Sync
+
+###
+#
+# This class wraps the logical ConditionVariable class to make it an easier to
+# work with interface that is similar to Windows' synchronization events.
+#
+###
+class Event
+
+  Infinite = 10000
+
+  #
+  # Initializes a waitable event.  The state parameter initializes the
+  # default state of the event.  If auto_reset is true, any calls to set()
+  # will automatically reset the event back to an unset state.
+  #
+  def initialize(state = false, auto_reset = true, param = nil)
+    self.state      = state
+    self.auto_reset = auto_reset
+    self.param      = param
+    self.mutex      = Mutex.new
+    self.cond       = ConditionVariable.new
+  end
+
+  #
+  # Sets the event and wakes up anyone who was waiting.
+  #
+  def set(param = nil)
+    self.param = param
+
+    self.mutex.synchronize {
+      # If this event does not automatically reset its state,
+      # set the state to true
+      if (auto_reset == false)
+        self.state = true
+      end
+
+      self.cond.broadcast
+    }
+  end
+
+  #
+  # Resets the signaled state to false.
+  #
+  def reset
+    self.param = nil
+    self.state = false
+  end
+
+  #
+  # Alias notify with set.
+  #
+  alias notify set
+
+  #
+  # Waits for the event to become signaled.  Timeout is measured in
+  # seconds.  Raises TimeoutError if the condition does not become signaled.
+  #
+  def wait(t = Infinite)
+    self.mutex.synchronize {
+      break if (self.state == true)
+
+      Timeout.timeout(t) {
+        self.cond.wait(self.mutex)
+      }
+    }
+
+    return self.param
+  end
+
+protected
+
+  attr_accessor :state, :auto_reset # :nodoc:
+  attr_accessor :param, :mutex, :cond # :nodoc:
+
+end
+
+end
+end
+

data/lib/rex/sync/read_write_lock.rb

@@ -0,0 +1,177 @@
+# -*- coding: binary -*-
+require 'thread'
+
+module Rex
+
+###
+#
+# This class implements a read/write lock synchronization
+# primitive.  It is meant to allow for more efficient access to
+# resources that are more often read from than written to and many
+# times can have concurrent reader threads.  By allowing the reader
+# threads to lock the resource concurrently rather than serially,
+# a large performance boost can be seen.  Acquiring a write lock
+# results in exclusive access to the resource and thereby prevents
+# any read operations during the time that a write lock is acquired.
+# Only one write lock may be acquired at a time.
+#
+###
+class ReadWriteLock
+
+  #
+  # Initializes a reader/writer lock instance.
+  #
+  def initialize
+    @read_sync_mutex  = Mutex.new
+    @write_sync_mutex = Mutex.new
+    @exclusive_mutex  = Mutex.new
+    @readers          = 0
+    @writer           = false
+  end
+
+  #
+  # Acquires the read lock for the calling thread.
+  #
+  def lock_read
+    read_sync_mutex.lock
+
+    begin
+      # If there are a non-zero number of readers and a
+      # writer is waiting to acquire the exclusive lock,
+      # free up the sync mutex temporarily and lock/unlock
+      # the exclusive lock.  This is to give the writer
+      # thread a chance to acquire the lock and prevents
+      # it from being constantly starved.
+      if ((@readers > 0) and
+          (@writer))
+        read_sync_mutex.unlock
+        exclusive_mutex.lock
+        exclusive_mutex.unlock
+        read_sync_mutex.lock
+      end
+
+      # Increment the active reader count
+      @readers += 1
+
+      # If we now have just one reader, acquire the exclusive
+      # lock.  Track the thread owner so that we release the
+      # lock from within the same thread context later on.
+      if (@readers == 1)
+        exclusive_mutex.lock
+
+        @owner = Thread.current
+      end
+    ensure
+      read_sync_mutex.unlock
+    end
+  end
+
+  #
+  # Releases the read lock for the calling thread.
+  #
+  def unlock_read
+    read_sync_mutex.lock
+
+    begin
+      unlocked = false
+
+      # Keep looping until we've lost this thread's reader
+      # lock
+      while (!unlocked)
+        # If there are no more readers left after this one
+        if (@readers - 1 == 0)
+          # If the calling thread is the owner of the exclusive
+          # reader lock, then let's release it
+          if (Thread.current == @owner)
+            @owner = nil
+
+            exclusive_mutex.unlock
+          end
+        # If there is more than one reader left and this thread is
+        # the owner of the exclusive lock, then keep looping so that
+        # we can eventually unlock the exclusive mutex in this thread's
+        # context
+        elsif (Thread.current == @owner)
+          read_sync_mutex.unlock
+
+          next
+        end
+
+        # Unlocked!
+        unlocked = true
+
+        # Decrement the active reader count
+        @readers -= 1
+      end
+    ensure
+      read_sync_mutex.unlock
+    end
+  end
+
+  #
+  # Acquire the exclusive write lock.
+  #
+  def lock_write
+    write_sync_mutex.lock
+
+    begin
+      @writer = true
+
+      exclusive_mutex.lock
+
+      @owner  = Thread.current
+    ensure
+      write_sync_mutex.unlock
+    end
+  end
+
+  #
+  # Release the exclusive write lock.
+  #
+  def unlock_write
+    # If the caller is not the owner of the write lock, then someone is
+    # doing something broken, let's let them know.
+    if (Thread.current != @owner)
+      raise RuntimeError, "Non-owner calling thread attempted to release write lock", caller
+    end
+
+    # Otherwise, release the exclusive write lock
+    @writer = false
+
+    exclusive_mutex.unlock
+  end
+
+  #
+  # Synchronize a block for read access.
+  #
+  def synchronize_read
+    lock_read
+    begin
+      yield
+    ensure
+      unlock_read
+    end
+  end
+
+  #
+  # Synchronize a block for write access.
+  #
+  def synchronize_write
+    lock_write
+    begin
+      yield
+    ensure
+      unlock_write
+    end
+  end
+
+protected
+
+  attr_accessor :read_sync_mutex # :nodoc:
+  attr_accessor :write_sync_mutex # :nodoc:
+  attr_accessor :exclusive_mutex # :nodoc:
+
+end
+
+end
+