Merge pull request #1852 from limhoff-r7/bug/migrations

[Delivers #50179803]

lib/msf/core/db_manager.rb

@@ -3,6 +3,7 @@
 require 'msf/base/config'
 require 'msf/core'
 require 'msf/core/db'
+require 'msf/core/db_manager/migration'
 require 'msf/core/task_manager'
 require 'fileutils'
 require 'shellwords'
@@ -17,6 +18,9 @@ module Msf
 ###
 
 class DBManager
+	# Provides :framework and other accessors
+	include Msf::DBManager::Migration
+	include Msf::Framework::Offspring
 
 	# Mainly, it's Ruby 1.9.1 that cause a lot of problems now, along with Ruby 1.8.6.
 	# Ruby 1.8.7 actually seems okay, but why tempt fate? Let's say 1.9.3 and beyond.
@@ -28,9 +32,6 @@ class DBManager
 		end
 	end
 
-	# Provides :framework and other accessors
-	include Framework::Offspring
-
 	# Returns true if we are ready to load/store data
 	def active
 		return false if not @usable
@@ -53,9 +54,6 @@ class DBManager
 	# Stores a TaskManager for serializing database events
 	attr_accessor :sink
 
-	# Flag to indicate database migration has completed
-	attr_accessor :migrated
-
 	# Flag to indicate that modules are cached
 	attr_accessor :modules_cached
 
@@ -287,33 +285,6 @@ class DBManager
 		end
 	end
 
-	# Migrate database to latest schema version.
-	#
-	# @param verbose [Boolean] see ActiveRecord::Migration.verbose
-	# @return [Array<ActiveRecord::MigrationProxy] List of migrations that ran.
-	#
-	# @see ActiveRecord::Migrator.migrate
-	def migrate(verbose=false)
-		ran = []
-		ActiveRecord::Migration.verbose = verbose
-
-		ActiveRecord::Base.connection_pool.with_connection do
-			begin
-				ran = ActiveRecord::Migrator.migrate(
-						ActiveRecord::Migrator.migrations_paths
-				)
-			# ActiveRecord::Migrator#migrate rescues all errors and re-raises them as
-			# StandardError
-			rescue StandardError => error
-				self.error = error
-				elog("DB.migrate threw an exception: #{error}")
-				dlog("Call stack:\n#{error.backtrace.join "\n"}")
-			end
-		end
-
-		return ran
-	end
-
 	def workspace=(workspace)
 		@workspace_name = workspace.name
 	end

lib/msf/core/db_manager/migration.rb

@@ -0,0 +1,58 @@
+module Msf
+	class DBManager
+		module Migration
+			# Migrate database to latest schema version.
+			#
+			# @param verbose [Boolean] see ActiveRecord::Migration.verbose
+			# @return [Array<ActiveRecord::MigrationProxy] List of migrations that
+			#   ran.
+			#
+			# @see ActiveRecord::Migrator.migrate
+			def migrate(verbose=false)
+				ran = []
+				ActiveRecord::Migration.verbose = verbose
+
+				ActiveRecord::Base.connection_pool.with_connection do
+					begin
+						ran = ActiveRecord::Migrator.migrate(
+								ActiveRecord::Migrator.migrations_paths
+						)
+					# ActiveRecord::Migrator#migrate rescues all errors and re-raises them
+					# as StandardError
+					rescue StandardError => error
+						self.error = error
+						elog("DB.migrate threw an exception: #{error}")
+						dlog("Call stack:\n#{error.backtrace.join "\n"}")
+					end
+				end
+
+				# Since the connections that existed before the migrations ran could
+				# have outdated column information, reset column information for all
+				# ActiveRecord::Base descendents to prevent missing method errors for
+				# column methods for columns created in migrations after the column
+				# information was cached.
+				reset_column_information
+
+				return ran
+			end
+
+			# Flag to indicate database migration has completed
+			#
+			# @return [Boolean]
+			attr_accessor :migrated
+
+			private
+
+			# Resets the column information for all descendants of ActiveRecord::Base
+			# since some of the migrations may have cached column information that
+			# has been updated by later migrations.
+			#
+			# @return [void]
+			def reset_column_information
+				ActiveRecord::Base.descendants.each do |descendant|
+					descendant.reset_column_information
+				end
+			end
+		end
+	end
+end

spec/lib/msf/db_manager_spec.rb

@@ -18,6 +18,7 @@ describe Msf::DBManager do
 		db_manager
 	end
 
+	it_should_behave_like 'Msf::DBManager::Migration'
 	it_should_behave_like 'Msf::DBManager::ImportMsfXml'
 
 	context '#initialize_metasploit_data_models' do

spec/support/shared/examples/msf/db_manager/migration.rb

@@ -0,0 +1,143 @@
+shared_examples_for 'Msf::DBManager::Migration' do
+	it { should be_a Msf::DBManager::Migration }
+
+	context '#migrate' do
+		def migrate
+			db_manager.migrate
+		end
+
+		it 'should create a connection' do
+			ActiveRecord::Base.connection_pool.should_receive(:with_connection).twice
+
+			migrate
+		end
+
+		it 'should call ActiveRecord::Migrator.migrate' do
+			ActiveRecord::Migrator.should_receive(:migrate).with(
+					ActiveRecord::Migrator.migrations_paths
+			)
+
+			migrate
+		end
+
+		it 'should return migrations that were ran from ActiveRecord::Migrator.migrate' do
+			migrations = [mock('Migration 1')]
+			ActiveRecord::Migrator.stub(:migrate => migrations)
+
+			migrate.should == migrations
+		end
+
+		it 'should reset the column information' do
+			db_manager.should_receive(:reset_column_information)
+
+			migrate
+		end
+
+		context 'with StandardError from ActiveRecord::Migration.migrate' do
+			let(:error) do
+				StandardError.new(message)
+			end
+
+			let(:message) do
+				"Error during migration"
+			end
+
+			before(:each) do
+				ActiveRecord::Migrator.stub(:migrate).and_raise(error)
+			end
+
+			it 'should set Msf::DBManager#error' do
+				migrate
+
+				db_manager.error.should == error
+			end
+
+			it 'should log error message at error level' do
+				db_manager.should_receive(:elog) do |error_message|
+					error_message.should include(error.to_s)
+				end
+
+				migrate
+			end
+
+			it 'should log error backtrace at debug level' do
+				db_manager.should_receive(:dlog) do |debug_message|
+					debug_message.should include('Call stack')
+				end
+
+				migrate
+			end
+		end
+
+		context 'with verbose' do
+			def migrate
+				db_manager.migrate(verbose)
+			end
+
+			context 'false' do
+				let(:verbose) do
+					false
+				end
+
+				it 'should set ActiveRecord::Migration.verbose to false' do
+					ActiveRecord::Migration.should_receive(:verbose=).with(verbose)
+
+					migrate
+				end
+			end
+
+			context 'true' do
+				let(:verbose) do
+					true
+				end
+
+				it 'should set ActiveRecord::Migration.verbose to true' do
+					ActiveRecord::Migration.should_receive(:verbose=).with(verbose)
+
+					migrate
+				end
+			end
+		end
+
+		context 'without verbose' do
+			it 'should set ActiveRecord::Migration.verbose to false' do
+				ActiveRecord::Migration.should_receive(:verbose=).with(false)
+
+				db_manager.migrate
+			end
+		end
+	end
+
+	context '#migrated' do
+		it { should respond_to :migrated }
+		it { should respond_to :migrated= }
+	end
+
+	context '#reset_column_information' do
+		def reset_column_information
+			db_manager.send(:reset_column_information)
+		end
+
+		it 'should use ActiveRecord::Base.descendants to find both direct and indirect subclasses' do
+			ActiveRecord::Base.should_receive(:descendants).and_return([])
+
+			reset_column_information
+		end
+
+		it 'should reset column information on each descendant of ActiveRecord::Base' do
+			descendants = []
+
+			1.upto(2) do |i|
+				descendants << mock("Descendant #{i}")
+			end
+
+			ActiveRecord::Base.stub(:descendants => descendants)
+
+			descendants.each do |descendant|
+				descendant.should_receive(:reset_column_information)
+			end
+
+			reset_column_information
+		end
+	end
+end