rex 1.0.2 → 2.0.2

data/lib/rex/struct2/s_struct.rb

@@ -0,0 +1,111 @@
+# -*- coding: binary -*-
+
+# Rex::Struct2
+module Rex
+module Struct2
+
+class SStruct
+
+  require 'rex/struct2/element'
+  include Rex::Struct2::Element
+
+  attr_reader  :leftover, :elements
+  attr_writer  :leftover, :elements
+
+  private :elements, :elements=
+
+  # watch out!, leftover returns our copy of the string!  so don't do
+  # anything stupid like struct.leftover.slice! !!
+
+  def initialize(*opts)
+    self.elements = [ ]
+    self.add_element(*opts)
+  end
+
+
+  def reset
+    elements.each {|e| e.reset}
+    return self
+  end
+
+  def add_element(*objs)
+    objs.each { |o|
+      elements << o
+      o.container = self
+    }
+    return self
+  end
+
+  def <<(obj)
+    self.add_element(obj)
+  end
+
+  def to_s
+    # !!! what do we do on mix restraint issues? just fail?
+    # maybe throw an exception, because that is most likely
+    # a usage error
+
+    buff = ""
+    elements.each do |e|
+      buff << e.to_s
+    end
+
+    if restraint && restraint.max
+      return buff.slice(0, restraint.max)
+    else
+      return buff
+    end
+  end
+
+  def length
+    return elements.length
+  end
+
+  def [](obj)
+    return elements[obj]
+  end
+
+  def each(&block)
+    return elements.each(&block)
+  end
+
+  def from_s(obytes)
+    # make my own copy so I can chop it up
+    bytes = obytes.dup
+    length = 0
+
+    # I don't think we should call update_restraint here, but
+    # I could have mis thought or something
+
+    # if we have a restraint (and if there is a val) truncate
+    if restraint
+      max = restraint.max
+      bytes = bytes.slice(0, max) if max
+    end
+
+    elements.each { |e|
+      used = e.from_s(bytes)
+      return if !used
+      bytes.slice!(0, used)
+      length += used
+    }
+
+    # make sure we matched out min restraint, else return failure
+    if restraint
+      min = restraint.min
+      return if min && length < min
+    end
+
+    # I guess this is me getting "set", so I should have a value
+    # and I should update my restraints on set
+    self.value = obytes.slice(0, length)
+
+    self.leftover = bytes
+    return(length)
+  end
+
+end
+
+# end Rex::Struct2
+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
+

data/lib/rex/sync/ref.rb

@@ -0,0 +1,58 @@
+# -*- coding: binary -*-
+require 'thread'
+
+module Rex
+
+###
+#
+# This module provides a uniform reference counted interface for classes to
+# use.
+#
+###
+module Ref
+
+  #
+  # Initializes the reference count to one.
+  #
+  def refinit
+    @_references       = 1
+    @_references_mutex = Mutex.new
+
+    self
+  end
+
+  #
+  # Increments the total number of references.
+  #
+  def ref
+    @_references_mutex.synchronize {
+      @_references += 1
+    }
+
+    self
+  end
+
+  #
+  # Decrements the total number of references.  If the reference count
+  # reaches zero, true is returned.  Otherwise, false is returned.
+  #
+  def deref
+    @_references_mutex.synchronize {
+      if ((@_references -= 1) == 0)
+        cleanup
+
+        true
+      else
+        false
+      end
+    }
+  end
+
+  #
+  # Called to clean up resources once the ref count drops to zero.
+  #
+  def cleanup
+  end
+
+end
+end

data/lib/rex/sync/thread_safe.rb

@@ -0,0 +1,83 @@
+# -*- coding: binary -*-
+require 'timeout'
+
+module Rex
+
+###
+#
+# This module provides a set of methods for performing various blocking
+# operations in a manner that is compatible with ruby style threads.
+#
+###
+module ThreadSafe
+
+  DefaultCycle = 0.2
+
+  #
+  # Wraps calls to select with a lower timeout period and does the
+  # calculations to walk down to zero timeout.  This has a little room for
+  # improvement in that it should probably check how much time actually
+  # elapsed during the select call considering ruby threading wont be exactly
+  # accurate perhaps.
+  #
+  def self.select(rfd = nil, wfd = nil, efd = nil, t = nil)
+    left = t
+
+    # Immediately raise a StreamClosedError if the socket was closed. This
+    # prevents a bad fd from being passed downstream and solves an issue
+    # with Ruby on Windows.
+    rfd.each { |fd| raise StreamClosedError.new(fd) if (fd.closed?)	} if rfd
+
+    begin
+      orig_size = rfd.length if (rfd)
+
+      # Poll the set supplied to us at least once.
+      begin
+        rv = ::IO.select(rfd, wfd, efd, DefaultCycle)
+      rescue ::IOError, ::Errno::EBADF, ::Errno::ENOTSOCK
+        # If a stream was detected as being closed, re-raise the error as
+        # a StreamClosedError with the specific file descriptor that was
+        # detected as being closed.  This is to better handle the case of
+        # a closed socket being detected so that it can be cleaned up and
+        # removed.
+        rfd.each { |fd| raise StreamClosedError.new(fd) if (fd.closed?)	} if rfd
+
+        # If the original rfd length is not the same as the current
+        # length, then the list may have been altered and as such may not
+        # contain the socket that caused the IOError.  This is a bad way
+        # to do this since it's possible that the array length could be
+        # back to the size that it was originally and yet have had the
+        # socket that caused the IOError to be removed.
+        return nil if (rfd and rfd.length != orig_size)
+
+        # Re-raise the exception since we didn't handle it here.
+        raise $!
+#			rescue ::Exception => e
+#			$stderr.puts "SELECT(#{t}) #{[rfd,wfd,efd].inspect} #{e.class} #{e} #{e.backtrace}"
+      end
+
+      return rv if (rv)
+
+      # Decrement the amount of time left by the polling cycle
+      left -= DefaultCycle if (left)
+
+      # Keep chugging until we run out of time, if time was supplied.
+    end while ((left == nil) or (left > 0))
+
+    # Nothin.
+    nil
+  end
+
+  #
+  # Simulates a sleep operation by selecting on nil until a timeout period
+  # expires.
+  #
+  def self.sleep(seconds=nil)
+    self.select(nil, nil, nil, seconds)
+
+    seconds
+  end
+
+end
+
+end