From 1c4f79507f97fa57facc680d4b4ab2b49cf0f8e2 Mon Sep 17 00:00:00 2001
From: John Davenport <johns10davenport@gmail.com>
Date: Mon, 7 Mar 2022 18:57:24 -0500
Subject: [PATCH 1/2] Add handle disconnect to websocket module

---
 lib/websocket.ex | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/lib/websocket.ex b/lib/websocket.ex
index dd089c1..5586d06 100644
--- a/lib/websocket.ex
+++ b/lib/websocket.ex
@@ -11,4 +11,13 @@ defmodule ChromeRemoteInterface.Websocket do
     send(state, {:message, frame_data})
     {:ok, state}
   end
+
+  def handle_disconnect(_status, state) do
+    Process.exit(state, :remote_closed)
+    {:ok, state}
+  end
+
+  def terminate({:remote, :closed}, _state) do
+    :stop
+  end
 end

From bcf12bd6a4693bb09b3e0cb7f409d1b1760e57e8 Mon Sep 17 00:00:00 2001
From: John Davenport <johns10davenport@gmail.com>
Date: Mon, 7 Mar 2022 18:57:42 -0500
Subject: [PATCH 2/2] Trap exits in PageSession

---
 lib/page_session.ex | 13 +++++++++++--
 1 file changed, 11 insertions(+), 2 deletions(-)

diff --git a/lib/page_session.ex b/lib/page_session.ex
index c39ad67..4b9bece 100644
--- a/lib/page_session.ex
+++ b/lib/page_session.ex
@@ -119,6 +119,7 @@ defmodule ChromeRemoteInterface.PageSession do
   # ---
 
   def init(url) do
+    Process.flag(:trap_exit, true)
     {:ok, socket} = ChromeRemoteInterface.Websocket.start_link(url)
 
     state = %__MODULE__{
@@ -282,8 +283,16 @@ defmodule ChromeRemoteInterface.PageSession do
     |> Enum.each(&send(&1, event))
   end
 
-  def terminate(_reason, state) do
-    Process.exit(state.socket, :kill)
+  # handle the trapped exit call
+  def handle_info({:EXIT, from, reason}, %{socket: socket} = state) do
+    cleanup(reason, state)
+    {:stop, reason, state} # see GenServer docs for other return types
+  end
+
+  def terminate(reason, state) do
+    cleanup(reason, state)
     :stop
   end
+
+  def cleanup(reason, state), do: Process.exit(state.socket, reason)
 end