Specs for activerecord patch

[#46141013]

Spec the desired behavior for ConnectionPool prior to removing the patch
to sync with upstream 3.2.12.

Author: limhoff-r7

Rakefile

@@ -2,12 +2,19 @@ require 'bundler/setup'
 
 require 'metasploit_data_models'
 
+pathname = Pathname.new(__FILE__)
+root = pathname.parent
+
+# add metasploit-framework/lib to load paths so rake files can just require
+# files normally without having to use __FILE__ and recalculating root and the
+# path to lib
+lib_pathname = root.join('lib')
+$LOAD_PATH.unshift(lib_pathname.to_s)
+
 #
 # load rake files like a rails engine
 #
 
-pathname = Pathname.new(__FILE__)
-root = pathname.parent
 rakefile_glob = root.join('lib', 'tasks', '**', '*.rake').to_path
 
 Dir.glob(rakefile_glob) do |rakefile|

lib/metasploit/framework.rb

@@ -0,0 +1,45 @@
+# Top-level namespace that is shared between {Metasploit::Framework
+# metasploit-framework} and pro, which uses Metasploit::Pro.
+module Metasploit
+	# Supports Rails and Rails::Engine like access to metasploit-framework so it
+	# works in compatible manner with activerecord's rake tasks and other
+	# railties.
+	module Framework
+		# Returns the environment for {Metasploit::Framework}.  Checks
+		# `METASPLOIT_FRAMEWORK_ENV` environment variable for value.  Defaults to
+		# `'development'` if `METASPLOIT_FRAMEWORK_ENV` is not set in the
+		# environment variables.
+		#
+		# {env} is a ActiveSupport::StringInquirer like `Rails.env` so it can be
+		# queried for its value.
+		#
+		# @example check if environment is development
+		#   if Metasploit::Framework.env.development?
+		#     # runs only when in development
+		#   end
+		#
+		# @return [ActiveSupport::StringInquirer] the environment name
+		def self.env
+			unless instance_variable_defined? :@env
+				name = ENV['METASPLOIT_FRAMEWORK_ENV']
+				name ||= 'development'
+				@env = ActiveSupport::StringInquirer.new(name)
+			end
+
+			@env
+		end
+
+		# Returns the root of the metasploit-framework project.  Use in place of
+		# `Rails.root`.
+		#
+		# @return [Pathname]
+		def self.root
+			unless instance_variable_defined? :@root
+				pathname = Pathname.new(__FILE__)
+				@root = pathname.parent.parent.parent
+			end
+
+			@root
+		end
+	end
+end

lib/tasks/database.rake

@@ -1,28 +1,6 @@
 load 'active_record/railties/databases.rake'
 
-module Metasploit
-	module Framework
-		def self.env
-			unless instance_variable_defined? :@env
-				name = ENV['METASPLOIT_FRAMEWORK_ENV']
-				name ||= 'development'
-				@env = ActiveSupport::StringInquirer.new(name)
-		  end
-
-			@env
-		end
-
-		def self.root
-			unless instance_variable_defined? :@root
-				pathname = Pathname.new(__FILE__)
-				@root = pathname.parent.parent.parent
-			end
-
-			@root
-		end
-	end
-end
-
+require 'metasploit/framework'
 
 # A modification to remove dependency on Rails.env
 #

spec/lib/active_record/connection_adapters/abstract_adapter/connection_pool_spec.rb

@@ -0,0 +1,251 @@
+require 'spec_helper'
+
+# include patch under test
+require 'msf/core/patches/active_record'
+# helps with environment configuration to use for connection to database
+require 'metasploit/framework'
+
+# load Mdm::Host for testing
+MetasploitDataModels.require_models
+
+describe ActiveRecord::ConnectionAdapters::ConnectionPool do
+	def database_configurations
+	  YAML.load_file(database_configurations_pathname)
+	end
+
+	def database_configurations_pathname
+		Metasploit::Framework.root.join('config', 'database.yml')
+	end
+
+	subject(:connection_pool) do
+		ActiveRecord::Base.connection_pool
+	end
+
+	# Not all specs require a database connection, and railties aren't being
+	# used, so have to manually establish connection.
+	before(:each) do
+		ActiveRecord::Base.configurations = database_configurations
+		spec = ActiveRecord::Base.configurations[Metasploit::Framework.env]
+		ActiveRecord::Base.establish_connection(spec)
+	end
+
+	after(:each) do
+		ActiveRecord::Base.clear_all_connections!
+	end
+
+	context '#active_thread_connection?' do
+		subject(:active_thread_connection?) do
+			connection_pool.active_thread_connection?
+		end
+
+		# Let! so that Thread is captured before creating and entering new Threads
+		let!(:main_thread) do
+			Thread.current
+		end
+
+		before(:each) do
+			ActiveRecord::Base.connection_pool.connection
+		end
+
+		context 'with connection id' do
+			subject(:active_thread_connection?) do
+				connection_pool.active_thread_connection?(connection_id)
+			end
+
+			context 'with connection id from other thread that has connection' do
+				let(:connection_id) do
+					main_thread.object_id
+				end
+
+				it 'should be true' do
+					thread = Thread.new do
+						Thread.current.should_not == main_thread
+
+						expect(active_thread_connection?).to be_true
+					end
+
+					thread.join
+				end
+			end
+		end
+
+		context 'without connection id' do
+			context 'in thread with connection' do
+				it { should be_true }
+			end
+
+			context 'in thread without connection' do
+				it 'should be false' do
+					thread = Thread.new do
+						Thread.current.should_not == main_thread
+						expect(active_thread_connection?).to be_false
+					end
+
+					thread.join
+				end
+			end
+		end
+	end
+
+	context 'checkout' do
+
+	end
+
+	context '#with_connection' do
+		def reserved_connection_count
+			connection_pool.instance_variable_get(:@reserved_connections).length
+		end
+
+		let(:connection_id) do
+			main_thread.object_id
+		end
+
+		it 'should call #current_connection_id' do
+			connection_pool.should_receive(
+					:current_connection_id
+			).at_least(
+					:once
+			).and_call_original
+
+			connection_pool.with_connection { }
+		end
+
+		it 'should call #active_thread_connection? with #current_connection_id' do
+			current_connection_id = mock('Connection ID')
+			connection_pool.stub(:current_connection_id => current_connection_id)
+
+			connection_pool.should_receive(
+					:active_thread_connection?
+			).with(
+					current_connection_id
+			).and_call_original
+
+			connection_pool.with_connection { }
+		end
+
+		it 'should yield #connection' do
+			connection = mock('Connection')
+			connection_pool.stub(:connection => connection)
+
+			expect { |block|
+				connection_pool.with_connection(&block)
+			}.to yield_with_args(connection)
+		end
+
+		context 'with active thread connection' do
+			let!(:connection) do
+				connection_pool.connection
+			end
+
+			after(:each) do
+				connection_pool.checkin connection
+			end
+
+			it 'should return true from #active_thread_connection?' do
+				expect(connection_pool.active_thread_connection?).to be_true
+			end
+
+			context 'with error' do
+				it 'should not release connection' do
+					expect {
+						# capture error so it doesn't stop example
+					  expect {
+							connection_pool.with_connection do
+								# raise error to trigger with_connection's ensure
+								raise ArgumentError, 'bad arguments'
+							end
+						}.to raise_error(ArgumentError)
+					}.to change {
+					  reserved_connection_count
+					}.by(0)
+				end
+			end
+
+			context 'without error' do
+				it 'should not release connection' do
+					expect {
+						connection_pool.with_connection { }
+					}.to change{
+				    reserved_connection_count
+					}.by(0)
+				end
+			end
+		end
+
+		context 'without active thread connection' do
+			it 'should return false from #active_thread_connection?' do
+				expect(connection_pool.active_thread_connection?).to be_false
+			end
+
+			context 'with error' do
+				it 'should not leave connection created for block' do
+					expect {
+						# capture error so it doesn't stop example
+						expect {
+							connection_pool.with_connection do
+								# raise error to trigger with_connection's ensure
+								raise ArgumentError, 'bad arguments'
+							end
+						}.to raise_error(ArgumentError)
+					}.to change {
+					  reserved_connection_count
+					}.by(0)
+				end
+			end
+
+			context 'without error' do
+				it 'should not leave connection created for block' do
+					expect {
+						connection_pool.with_connection { }
+					}.to change{
+						reserved_connection_count
+					}.by(0)
+				end
+			end
+
+			context 'with nested' do
+				it 'should not reserve another connection in the nested block' do
+				  before_count = reserved_connection_count
+
+					connection_pool.with_connection do
+						child_count = reserved_connection_count
+						count_change = child_count - before_count
+
+						count_change.should == 1
+
+						connection_pool.with_connection do
+						  grandchild_count = reserved_connection_count
+
+						  grandchild_count.should == child_count
+						end
+					end
+
+					after_count = reserved_connection_count
+
+					after_count.should == before_count
+				end
+			end
+
+			context 'without with_connection first' do
+				it 'should use connection reserved outside with_connection' do
+					# Using query methods without a block is expected to retain the
+					# reserved connection
+					expect {
+						# access database outside with_connection block
+						Mdm::Host.count
+					}.to change {
+						reserved_connection_count
+					}.by(1)
+
+					outside = reserved_connection_count
+
+					connection_pool.with_connection do
+						inside = reserved_connection_count
+
+						inside.should == outside
+					end
+				end
+			end
+		end
+	end
+end