From e93dc2817491f69e77535fb5565d70797fc789fd Mon Sep 17 00:00:00 2001 From: zinduolis Date: Mon, 25 Aug 2025 23:30:54 +1000 Subject: [PATCH] Add disconnect_all_active_record! calls before forking in multiple specs for improved SQL connection handling safety --- .../main/autorun_engine/autorun_engine_spec.rb | 3 +++ .../handlers/browser_details_handler_spec.rb | 3 +++ .../handlers/dynamic_reconstruction_spec.rb | 4 ++++ .../network_stack/handlers/redirector_spec.rb | 5 +++++ .../extensions/websocket_hooked_browser_spec.rb | 5 +++++ spec/beef/modules/debug/test_beef_debugs_spec.rb | 3 +++ spec/spec_helper.rb | 16 +++++++++++----- 7 files changed, 34 insertions(+), 5 deletions(-) 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 aef99a958..6b42e8eb6 100644 --- a/spec/beef/core/main/autorun_engine/autorun_engine_spec.rb +++ b/spec/beef/core/main/autorun_engine/autorun_engine_spec.rb @@ -68,6 +68,9 @@ RSpec.describe 'AutoRunEngine Test', run_on_browserstack: true do # Generate a token for the server to respond with @token = BeEF::Core::Crypto.api_token + # ***** IMPORTANT: close any and all AR/OTR connections before forking ***** + disconnect_all_active_record! + # Initiate server start-up @pids = fork do BeEF::API::Registrar.instance.fire(BeEF::API::Server, 'pre_http_start', http_hook_server) 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 4c54f2d37..eca9e33bc 100644 --- a/spec/beef/core/main/handlers/browser_details_handler_spec.rb +++ b/spec/beef/core/main/handlers/browser_details_handler_spec.rb @@ -63,6 +63,9 @@ RSpec.describe 'Browser Details Handler', run_on_browserstack: true do # Generate a token for the server to respond with @token = BeEF::Core::Crypto.api_token + # ***** IMPORTANT: close any and all AR/OTR connections before forking ***** + disconnect_all_active_record! + # Initiate server start-up @pids = fork do BeEF::API::Registrar.instance.fire(BeEF::API::Server, 'pre_http_start', http_hook_server) 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 458bc3d70..3cfd652fe 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 @@ -12,6 +12,10 @@ RSpec.describe 'BeEF Dynamic Reconsturction' do @server = Thin::Server.new('127.0.0.1', @port.to_s, @rackApp) trap("INT") { @server.stop } trap("TERM") { @server.stop } + + # ***** IMPORTANT: close any and all AR/OTR connections before forking ***** + disconnect_all_active_record! + @pid = fork do @server.start! end 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 8333ab459..e17293127 100644 --- a/spec/beef/core/main/network_stack/handlers/redirector_spec.rb +++ b/spec/beef/core/main/network_stack/handlers/redirector_spec.rb @@ -12,6 +12,11 @@ RSpec.describe 'BeEF Redirector' do @server = Thin::Server.new('127.0.0.1', @port.to_s, @rackApp) trap("INT") { @server.stop } trap("TERM") { @server.stop } + + + # ***** IMPORTANT: close any and all AR/OTR connections before forking ***** + disconnect_all_active_record! + @pid = fork do @server.start! end diff --git a/spec/beef/extensions/websocket_hooked_browser_spec.rb b/spec/beef/extensions/websocket_hooked_browser_spec.rb index 6acf55641..45b02cc89 100644 --- a/spec/beef/extensions/websocket_hooked_browser_spec.rb +++ b/spec/beef/extensions/websocket_hooked_browser_spec.rb @@ -59,6 +59,11 @@ RSpec.describe 'Browser hooking with Websockets', run_on_browserstack: true do http_hook_server.prepare # Generate a token for the server to respond with @token = BeEF::Core::Crypto.api_token + + + # ***** IMPORTANT: close any and all AR/OTR connections before forking ***** + disconnect_all_active_record! + # Initiate server start-up @pids = fork do BeEF::API::Registrar.instance.fire(BeEF::API::Server, 'pre_http_start', http_hook_server) diff --git a/spec/beef/modules/debug/test_beef_debugs_spec.rb b/spec/beef/modules/debug/test_beef_debugs_spec.rb index 1bc89d813..054cf287c 100644 --- a/spec/beef/modules/debug/test_beef_debugs_spec.rb +++ b/spec/beef/modules/debug/test_beef_debugs_spec.rb @@ -60,6 +60,9 @@ RSpec.describe 'BeEF Debug Command Modules:', run_on_browserstack: true do # Generate a token for the server to respond with @token = BeEF::Core::Crypto.api_token + # ***** IMPORTANT: close any and all AR/OTR connections before forking ***** + disconnect_all_active_record! + # Initiate server start-up @pids = fork do BeEF::API::Registrar.instance.fire(BeEF::API::Server, 'pre_http_start', http_hook_server) diff --git a/spec/spec_helper.rb b/spec/spec_helper.rb index 9571463f9..323a72938 100644 --- a/spec/spec_helper.rb +++ b/spec/spec_helper.rb @@ -193,14 +193,20 @@ RSpec.configure do |config| # --- HARD fork-safety: disconnect every pool/adapter we can find --- def disconnect_all_active_record! + print_info "Entering disconnect_all_active_record!" if defined?(ActiveRecord::Base) - # Disconnect every connection pool explicitly + print_info "Disconnecting ActiveRecord connections" handler = ActiveRecord::Base.connection_handler - handler.connection_pool_list.each { |pool| pool.disconnect! } if handler.respond_to?(:connection_pool_list) - ActiveRecord::Base.clear_active_connections! - ActiveRecord::Base.clear_all_connections! + if handler.respond_to?(:connection_pool_list) + print_info "Using connection_pool_list" + handler.connection_pool_list.each { |pool| pool.disconnect! } + elsif handler.respond_to?(:connection_pools) + print_info "Using connection_pools" + handler.connection_pools.each_value { |pool| pool.disconnect! } + end + else + print_info "ActiveRecord::Base not defined" end - OTR::ActiveRecord.disconnect! if defined?(OTR::ActiveRecord) end def start_beef_server