service allocation draft

git-svn-id: file:///home/svn/incoming/trunk@2844 4d416f70-5f16-0410-b530-b9f4589650da

lib/rex/proto.rb

@@ -1 +1,9 @@
 require 'rex/proto/http'
+
+module Rex
+module Proto
+
+attr_accessor :alias
+
+end
+end

lib/rex/proto/http/client.rb

@@ -16,6 +16,8 @@ module Http
 ###
 class Client
 
+	include Proto
+
 	#
 	# Performs a block-based HTTP operation
 	#
@@ -29,6 +31,9 @@ class Client
 		end
 	end
 
+	#
+	# Initializes a GET request and returns it to the caller.
+	#
 	def self.get(uri = '/', proto = DefaultProtocol)
 		return init_request(Request::Get.new(uri, proto))
 	end
@@ -38,6 +43,13 @@ class Client
 		self.port     = port.to_i
 	end
 
+	#
+	# HTTP client.
+	#
+	def alias
+		"HTTP Client"
+	end
+
 	#
 	# Connects to the remote server if possible.
 	#

lib/rex/proto/http/server.rb

@@ -72,6 +72,8 @@ end
 ###
 class Server
 
+	include Proto
+
 	DefaultServer = "Rex"
 
 	def initialize(port = 80, listen_host = '0.0.0.0')
@@ -81,6 +83,13 @@ class Server
 		self.resources   = {}
 	end
 
+	#
+	# HTTP server.
+	#
+	def alias
+		super || "HTTP Server"
+	end
+
 	#
 	# Listens on the defined port and host and starts monitoring for clients.
 	#

lib/rex/ref.rb

@@ -0,0 +1,48 @@
+require 'thread'
+
+module Rex
+
+###
+#
+# Ref
+# ---
+#
+# 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 {
+			((@_references -= 1) == 0) ? true : false
+		}
+	end
+
+end
+end

lib/rex/service.rb

@@ -0,0 +1,22 @@
+require 'rex'
+
+require 'rex'
+require 'rex/proto'
+
+module Rex
+
+###
+#
+# Service
+# -------
+#
+# The service module is used to extend classes that are passed into the
+# service manager start routine.  It provides extra methods, such as reference
+# counting, that are used to track the service instances more uniformly.
+#
+###
+module Service
+	include Ref
+end
+
+end

lib/rex/service_manager.rb

@@ -0,0 +1,95 @@
+require 'singleton'
+require 'rex'
+require 'rex/service'
+
+module Rex
+
+###
+#
+# ServiceManager
+# --------------
+#
+# This class manages service allocation and interaction.  This class can be
+# used to start HTTP servers and manage them and all that stuff.  Yup.
+#
+###
+class ServiceManager < Hash
+
+	#
+	# This class is a singleton.
+	#
+	include Singleton
+
+	#
+	# Calls the instance method to start a service.
+	#
+	def self.start(klass, *args)
+		self.instance.start(klass, *args)
+	end
+	
+	#
+	# Calls the instance method to stop a service.
+	#
+	def self.stop(als)
+		self.instance.stop(als)
+	end
+
+	#
+	# Starts a service and assigns it a unique name in the service hash.
+	#
+	def start(klass, *args)
+		# Get the hardcore alias.
+		hals = "__#{klass.name}#{args.to_s}"
+
+		# Has a service already been constructed for this guy?  If so, increment
+		# its reference count like it aint no thang.
+		if (inst = self[hals])
+			inst.ref
+			return inst
+		end
+
+		inst = klass.new(*args)
+		als  = inst.alias
+
+		# Find an alias that isn't taken.
+		if (self[als])
+			cnt  = 1
+			cnt += 1 while (self[als + " #{cnt}"])
+			als += " #{cnt}"
+		end
+
+		# Extend the instance as a service.
+		inst.extend(Rex::Service)
+
+		# Re-aliases the instance.
+		inst.alias = als
+
+		# Fire up the engines.  If an error occurs an exception will be 
+		# raised.
+		inst.start
+
+		# Alias associate and initialize reference counting
+		self[als] = self[hals] = inst.refinit
+	end
+
+	#
+	# Stops a service using the provided alias
+	#
+	def stop(als)
+		# Stop the service and be done wif it, but only if the number of
+		# references has dropped to zero
+		if ((inst = self[als]) and
+		    (inst.deref))
+			inst.stop
+
+			# Since the instance may have multiple aliases, scan through
+			# all the pairs for matching stuff.
+			self.each_pair { |cals, cinst|
+				self.delete(cals) if (inst == cinst)
+			}
+		end
+	end
+	
+end
+
+end

lib/rex/service_manager.rb.ut.rb

@@ -0,0 +1,32 @@
+#!/usr/bin/ruby
+
+$:.unshift(File.join(File.dirname(__FILE__), '..'))
+
+require 'test/unit'
+require 'rex/service_manager'
+
+class Rex::ServiceManager::UnitTest < Test::Unit::TestCase
+
+	Klass = Rex::ServiceManager
+
+	def test_svcm
+		begin
+			c = Klass
+			s = c.start(Rex::Proto::Http::Server, 8090)
+			assert_not_nil(s)
+			t = c.start(Rex::Proto::Http::Server, 8090)
+			assert_not_nil(t)
+			assert_equal(s, t)
+			z = c.start(Rex::Proto::Http::Server, 8091)
+			assert_not_equal(t, z)
+			assert_equal("HTTP Server", s.alias)
+			assert_equal("HTTP Server 1", z.alias)
+		ensure
+			c.stop(s.alias) if (s)
+			c.stop(z.alias) if (z)
+			c.stop(t.alias) if (t)
+		end
+
+	end
+
+end