From 41a772a31ce39c6989d6d6e8c736e72247f2b620 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 | 37 ++++---------------
 2 files changed, 14 insertions(+), 44 deletions(-)

diff --git a/lib/lograge/rails_ext/action_cable/channel/base.rb b/lib/lograge/rails_ext/action_cable/channel/base.rb
index 31286629..d024e021 100644
--- a/lib/lograge/rails_ext/action_cable/channel/base.rb
+++ b/lib/lograge/rails_ext/action_cable/channel/base.rb
@@ -1,22 +1,13 @@
-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
 
@@ -28,3 +19,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 8f7c512f..73c2b626 100644
--- a/lib/lograge/rails_ext/action_cable/connection/base.rb
+++ b/lib/lograge/rails_ext/action_cable/connection/base.rb
@@ -1,42 +1,19 @@
-module ActionCable
-  module Connection
-    class Base
-      # rubocop:disable Metrics/MethodLength
+module Lograge
+  module ActionCable
+    module ConnectionInstrumentation
       def handle_open
-        ActiveSupport::Notifications.instrument('connect.action_cable', notification_payload('connect')) do
-          begin
-            @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
-        end
+        ActiveSupport::Notifications.instrument('connect.action_cable', notification_payload('connect')) { super }
       end
-      # rubocop:enable Metrics/MethodLength
 
       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)