From d7e31c1f011d9dc016e6c4b40e53c911bd995095 Mon Sep 17 00:00:00 2001
From: Dmitry Vorotilin <d.vorotilin@gmail.com>
Date: Thu, 4 Jan 2024 16:58:47 +0300
Subject: [PATCH] fix: Close connections

---
 lib/ferrum/browser.rb            | 1 +
 lib/ferrum/browser/web_socket.rb | 3 ++-
 lib/ferrum/context.rb            | 8 ++++++++
 lib/ferrum/contexts.rb           | 5 +++++
 lib/ferrum/page.rb               | 8 +++++++-
 spec/spec_helper.rb              | 2 +-
 6 files changed, 24 insertions(+), 3 deletions(-)

diff --git a/lib/ferrum/browser.rb b/lib/ferrum/browser.rb
index 929a5ef5..544943b8 100644
--- a/lib/ferrum/browser.rb
+++ b/lib/ferrum/browser.rb
@@ -208,6 +208,7 @@ def restart
     def quit
       return unless @client
 
+      contexts.close_connections
       @client.close
       @process.stop
       @client = @process = @contexts = nil
diff --git a/lib/ferrum/browser/web_socket.rb b/lib/ferrum/browser/web_socket.rb
index 97af3010..aa04a795 100644
--- a/lib/ferrum/browser/web_socket.rb
+++ b/lib/ferrum/browser/web_socket.rb
@@ -52,6 +52,7 @@ def on_message(event)
 
       def on_close(_event)
         @messages.close
+        @sock.close
         @thread.kill
       end
 
@@ -65,7 +66,7 @@ def send_message(data)
 
       def write(data)
         @sock.write(data)
-      rescue EOFError, Errno::ECONNRESET, Errno::EPIPE
+      rescue EOFError, Errno::ECONNRESET, Errno::EPIPE, IOError # rubocop:disable Lint/ShadowedException
         @messages.close
       end
 
diff --git a/lib/ferrum/context.rb b/lib/ferrum/context.rb
index a22c583f..f5cc190f 100644
--- a/lib/ferrum/context.rb
+++ b/lib/ferrum/context.rb
@@ -71,6 +71,14 @@ def delete_target(target_id)
       @targets.delete(target_id)
     end
 
+    def close_targets_connection
+      @targets.each_value do |target|
+        next unless target.attached?
+
+        target.page.close_connection
+      end
+    end
+
     def dispose
       @contexts.dispose(@id)
     end
diff --git a/lib/ferrum/contexts.rb b/lib/ferrum/contexts.rb
index 40b888e0..9753c0e7 100644
--- a/lib/ferrum/contexts.rb
+++ b/lib/ferrum/contexts.rb
@@ -45,11 +45,16 @@ def create(**options)
 
     def dispose(context_id)
       context = @contexts[context_id]
+      context.close_targets_connection
       @client.command("Target.disposeBrowserContext", browserContextId: context.id)
       @contexts.delete(context_id)
       true
     end
 
+    def close_connections
+      @contexts.each_value(&:close_targets_connection)
+    end
+
     def reset
       @default_context = nil
       @contexts.each_key { |id| dispose(id) }
diff --git a/lib/ferrum/page.rb b/lib/ferrum/page.rb
index 5d0101b4..73707449 100644
--- a/lib/ferrum/page.rb
+++ b/lib/ferrum/page.rb
@@ -120,7 +120,13 @@ def go_to(url = nil)
     def close
       @headers.clear
       client(browser: true).command("Target.closeTarget", targetId: @target_id)
-      @page_client.close
+      close_connection
+
+      true
+    end
+
+    def close_connection
+      @page_client&.close
     end
 
     #
diff --git a/spec/spec_helper.rb b/spec/spec_helper.rb
index 4c944a29..9f94527b 100644
--- a/spec/spec_helper.rb
+++ b/spec/spec_helper.rb
@@ -46,7 +46,7 @@
   end
 
   config.after(:all) do
-    @browser&.quit
+    @browser.quit
   end
 
   config.before(:each) do