From 2c1a849d68a567e85e68d7fd8fa3ba4184a4da06 Mon Sep 17 00:00:00 2001
From: Vladimir Dementyev <dementiev.vm@gmail.com>
Date: Wed, 6 May 2020 09:56:54 +0300
Subject: [PATCH] Use prepend to patch Action Cable classes

To preserve the original functionality a better monkey-patching technique than copying the code should be used, e.g., Module#prepend

Fixes #304
---
 .../rails_ext/action_cable/channel/base.rb    | 21 ++++--------
 .../rails_ext/action_cable/connection/base.rb | 33 ++++---------------
 2 files changed, 14 insertions(+), 40 deletions(-)

diff --git a/lib/lograge/rails_ext/action_cable/channel/base.rb b/lib/lograge/rails_ext/action_cable/channel/base.rb
index 1387d023..6ad7ff22 100644
--- a/lib/lograge/rails_ext/action_cable/channel/base.rb
+++ b/lib/lograge/rails_ext/action_cable/channel/base.rb
@@ -1,24 +1,15 @@
 # frozen_string_literal: true
 
-module ActionCable
-  module Channel
-    class Base
+module Lograge
+  module ActionCable
+    module ChannelInstrumentation
       def subscribe_to_channel
-        ActiveSupport::Notifications.instrument('subscribe.action_cable', notification_payload('subscribe')) do
-          run_callbacks :subscribe do
-            subscribed
-          end
-
-          reject_subscription if subscription_rejected?
-          ensure_confirmation_sent
-        end
+        ActiveSupport::Notifications.instrument('subscribe.action_cable', notification_payload('subscribe')) { super }
       end
 
       def unsubscribe_from_channel
         ActiveSupport::Notifications.instrument('unsubscribe.action_cable', notification_payload('unsubscribe')) do
-          run_callbacks :unsubscribe do
-            unsubscribed
-          end
+          super
         end
       end
 
@@ -30,3 +21,5 @@ def notification_payload(method_name)
     end
   end
 end
+
+ActionCable::Channel::Base.prepend(Lograge::ActionCable::ChannelInstrumentation)
diff --git a/lib/lograge/rails_ext/action_cable/connection/base.rb b/lib/lograge/rails_ext/action_cable/connection/base.rb
index 253cf07b..9ce285c2 100644
--- a/lib/lograge/rails_ext/action_cable/connection/base.rb
+++ b/lib/lograge/rails_ext/action_cable/connection/base.rb
@@ -1,40 +1,21 @@
 # frozen_string_literal: true
 
-module ActionCable
-  module Connection
-    class Base
+module Lograge
+  module ActionCable
+    module ConnectionInstrumentation
       def handle_open
-        ActiveSupport::Notifications.instrument('connect.action_cable', notification_payload('connect')) do
-          @protocol = websocket.protocol
-          connect if respond_to?(:connect)
-          subscribe_to_internal_channel
-          send_welcome_message
-
-          message_buffer.process!
-          server.add_connection(self)
-        rescue ActionCable::Connection::Authorization::UnauthorizedError
-          respond_to_invalid_request
-        end
+        ActiveSupport::Notifications.instrument('connect.action_cable', notification_payload('connect')) { super }
       end
 
       def handle_close
-        ActiveSupport::Notifications.instrument('disconnect.action_cable', notification_payload('disconnect')) do
-          logger.info finished_request_message if Lograge.lograge_config.keep_original_rails_log
-
-          server.remove_connection(self)
-
-          subscriptions.unsubscribe_from_all
-          unsubscribe_from_internal_channel
-
-          disconnect if respond_to?(:disconnect)
-        end
+        ActiveSupport::Notifications.instrument('disconnect.action_cable', notification_payload('disconnect')) { super }
       end
 
-      private
-
       def notification_payload(method_name)
         { connection_class: self.class.name, action: method_name, data: request.params }
       end
     end
   end
 end
+
+ActionCable::Connection::Base.prepend(Lograge::ActionCable::ConnectionInstrumentation)