From bd67b8ba2a990c795569f9c79c3110bc901792bc Mon Sep 17 00:00:00 2001 From: kaitoozawa Date: Tue, 16 Dec 2025 10:51:55 +1000 Subject: [PATCH] AR snapshot/restore helpers for test isolation --- .../autorun_engine/autorun_engine_spec.rb | 2 + .../handlers/browser_details_handler_spec.rb | 3 +- .../handlers/dynamic_reconstruction_spec.rb | 2 + .../network_stack/handlers/redirector_spec.rb | 2 + spec/beef/extensions/requester_spec.rb | 2 + .../websocket_hooked_browser_spec.rb | 2 + .../modules/debug/test_beef_debugs_spec.rb | 2 + spec/spec_helper.rb | 56 +++++++++++++++++++ 8 files changed, 70 insertions(+), 1 deletion(-) diff --git a/spec/beef/core/main/autorun_engine/autorun_engine_spec.rb b/spec/beef/core/main/autorun_engine/autorun_engine_spec.rb index 02258042d..e8bf3e7e3 100644 --- a/spec/beef/core/main/autorun_engine/autorun_engine_spec.rb +++ b/spec/beef/core/main/autorun_engine/autorun_engine_spec.rb @@ -12,6 +12,7 @@ require_relative '../../../../support/beef_test' RSpec.describe 'AutoRunEngine Test', run_on_browserstack: true do before(:all) do + @__ar_config_snapshot = SpecActiveRecordConnection.snapshot @config = BeEF::Core::Configuration.instance # Grab DB file and regenerate if requested @@ -110,6 +111,7 @@ RSpec.describe 'AutoRunEngine Test', run_on_browserstack: true do after(:all) do server_teardown(@driver, @pid, @pids) disconnect_all_active_record! + SpecActiveRecordConnection.restore!(@__ar_config_snapshot) end it 'AutoRunEngine is working' do diff --git a/spec/beef/core/main/handlers/browser_details_handler_spec.rb b/spec/beef/core/main/handlers/browser_details_handler_spec.rb index c0c5fe9d7..9d7ff1080 100644 --- a/spec/beef/core/main/handlers/browser_details_handler_spec.rb +++ b/spec/beef/core/main/handlers/browser_details_handler_spec.rb @@ -12,7 +12,7 @@ require_relative '../../../../support/beef_test' RSpec.describe 'Browser Details Handler', run_on_browserstack: true do before(:all) do - + @__ar_config_snapshot = SpecActiveRecordConnection.snapshot @config = BeEF::Core::Configuration.instance db_file = @config.get('beef.database.file') print_info 'Resetting the database for BeEF.' @@ -106,6 +106,7 @@ RSpec.describe 'Browser Details Handler', run_on_browserstack: true do after(:all) do server_teardown(@driver, @pid, @pids) disconnect_all_active_record! + @__ar_config_snapshot = SpecActiveRecordConnection.snapshot end it 'can successfully hook a browser' do diff --git a/spec/beef/core/main/network_stack/handlers/dynamic_reconstruction_spec.rb b/spec/beef/core/main/network_stack/handlers/dynamic_reconstruction_spec.rb index 3cfd652fe..4a0899648 100644 --- a/spec/beef/core/main/network_stack/handlers/dynamic_reconstruction_spec.rb +++ b/spec/beef/core/main/network_stack/handlers/dynamic_reconstruction_spec.rb @@ -1,6 +1,7 @@ RSpec.describe 'BeEF Dynamic Reconsturction' do before(:all) do + @__ar_config_snapshot = SpecActiveRecordConnection.snapshot @port = 2001 config = {} config[:BindAddress] = '127.0.0.1' @@ -25,6 +26,7 @@ RSpec.describe 'BeEF Dynamic Reconsturction' do after(:all) do Process.kill("INT",@pid) + SpecActiveRecordConnection.restore!(@__ar_config_snapshot) end it 'delete' do diff --git a/spec/beef/core/main/network_stack/handlers/redirector_spec.rb b/spec/beef/core/main/network_stack/handlers/redirector_spec.rb index 415d2ff39..df78232ca 100644 --- a/spec/beef/core/main/network_stack/handlers/redirector_spec.rb +++ b/spec/beef/core/main/network_stack/handlers/redirector_spec.rb @@ -1,6 +1,7 @@ RSpec.describe 'BeEF Redirector' do before(:all) do + @__ar_config_snapshot = SpecActiveRecordConnection.snapshot @port = 2002 config = {} config[:BindAddress] = '127.0.0.1' @@ -26,6 +27,7 @@ RSpec.describe 'BeEF Redirector' do after(:all) do Process.kill("INT",@pid) + SpecActiveRecordConnection.restore!(@__ar_config_snapshot) end it 'redirects' do diff --git a/spec/beef/extensions/requester_spec.rb b/spec/beef/extensions/requester_spec.rb index 8aabfbd23..cbe5a88d6 100644 --- a/spec/beef/extensions/requester_spec.rb +++ b/spec/beef/extensions/requester_spec.rb @@ -24,6 +24,7 @@ RSpec.describe 'BeEF Extension Requester' do xit 'requester works' do begin + ar_snapshot = SpecActiveRecordConnection.snapshot # Start beef server @config = BeEF::Core::Configuration.instance @config.set('beef.credentials.user', 'beef') @@ -77,6 +78,7 @@ RSpec.describe 'BeEF Extension Requester' do BeEF::Core::Models::Http.where(hooked_browser_id: hb_session).delete_all if defined? hb_session Process.kill('KILL', @pid) if defined? @pid Process.kill('KILL', @pids) if defined? @pids + SpecActiveRecordConnection.restore!(ar_snapshot) end end end diff --git a/spec/beef/extensions/websocket_hooked_browser_spec.rb b/spec/beef/extensions/websocket_hooked_browser_spec.rb index 34156645e..4058429b4 100644 --- a/spec/beef/extensions/websocket_hooked_browser_spec.rb +++ b/spec/beef/extensions/websocket_hooked_browser_spec.rb @@ -13,6 +13,7 @@ require 'websocket-client-simple' RSpec.describe 'Browser hooking with Websockets', run_on_browserstack: true do before(:all) do + @__ar_config_snapshot = SpecActiveRecordConnection.snapshot @config = BeEF::Core::Configuration.instance # Grab DB file and regenerate if requested print_info 'Loading database' @@ -104,6 +105,7 @@ RSpec.describe 'Browser hooking with Websockets', run_on_browserstack: true do after(:all) do server_teardown(@driver, @pid, @pids) disconnect_all_active_record! + @__ar_config_snapshot = SpecActiveRecordConnection.snapshot end it 'confirms a websocket server has been started' do diff --git a/spec/beef/modules/debug/test_beef_debugs_spec.rb b/spec/beef/modules/debug/test_beef_debugs_spec.rb index 075d13490..6a178078f 100644 --- a/spec/beef/modules/debug/test_beef_debugs_spec.rb +++ b/spec/beef/modules/debug/test_beef_debugs_spec.rb @@ -12,6 +12,7 @@ require_relative '../../../support/beef_test' RSpec.describe 'BeEF Debug Command Modules:', run_on_browserstack: true do before(:all) do + @__ar_config_snapshot = SpecActiveRecordConnection.snapshot # Grab config and set creds in variables for ease of access @config = BeEF::Core::Configuration.instance @pids = [] # ensure defined for teardown consistency @@ -128,6 +129,7 @@ RSpec.describe 'BeEF Debug Command Modules:', run_on_browserstack: true do after(:all) do server_teardown(@driver, @pid, @pids) disconnect_all_active_record! + SpecActiveRecordConnection.restore!(@__ar_config_snapshot) end it 'The Test_beef.debug() command module successfully executes' do diff --git a/spec/spec_helper.rb b/spec/spec_helper.rb index aa23cbfbb..73aa41101 100644 --- a/spec/spec_helper.rb +++ b/spec/spec_helper.rb @@ -322,4 +322,60 @@ require 'socket' pid = nil end +end + +# ------------------------------------------------------------------- +# ActiveRecord connection snapshot/restore helpers (test isolation) +# Some specs disconnect ActiveRecord (fork safety), destroying the SQLite in-memory DB. +# These helpers restore it for later specs. +# ------------------------------------------------------------------- +module SpecActiveRecordConnection + module_function + + def snapshot + # Capture the current AR connection configuration hash if possible. + if ActiveRecord::Base.respond_to?(:connection_db_config) && ActiveRecord::Base.connection_db_config + ActiveRecord::Base.connection_db_config.configuration_hash + else + ActiveRecord::Base.connection_config + end + rescue StandardError + nil + end + + def restore!(config_hash) + # Ensure we don't leave AR disconnected for subsequent specs. + begin + handler = ActiveRecord::Base.connection_handler + if handler.respond_to?(:connection_pool_list) + handler.connection_pool_list.each { |pool| pool.disconnect! } + elsif handler.respond_to?(:connection_pools) + handler.connection_pools.each_value { |pool| pool.disconnect! } + else + ActiveRecord::Base.connection_pool.disconnect! + end + rescue StandardError + # ignore + end + + if config_hash + OTR::ActiveRecord.configure_from_hash!(config_hash) + else + # Fallback to suite default + OTR::ActiveRecord.configure_from_hash!(adapter: 'sqlite3', database: ':memory:') + end + + if Gem.loaded_specs['otr-activerecord'].version > Gem::Version.create('1.4.2') + OTR::ActiveRecord.establish_connection! + end + ActiveRecord::Schema.verbose = false + + # Run migrations if the restored DB is empty/outdated + ActiveRecord::Migration.verbose = false + ActiveRecord::Migrator.migrations_paths = [File.join('core', 'main', 'ar-migrations')] + context = ActiveRecord::MigrationContext.new(ActiveRecord::Migrator.migrations_paths) + if context.needs_migration? + ActiveRecord::Migrator.new(:up, context.migrations, context.schema_migration, context.internal_metadata).migrate + end + end end \ No newline at end of file