feat: action cable connection callbacks

actioncable/CHANGELOG.md

@@ -1,3 +1,7 @@
+*   Added command callbacks to `ActionCable::Base::Connection`.
 
+    Now you can define `before_command`, `after_command`, and `around_command` to be invoked before, after or around any command received by a client respectively.
+
+    *Vladimir Dementyev*
 
 Please check [7-0-stable](https://github.com/rails/rails/blob/7-0-stable/actioncable/CHANGELOG.md) for previous changes.

actioncable/lib/action_cable/connection/base.rb

@@ -47,6 +47,7 @@ module ActionCable
       include Identification
       include InternalChannel
       include Authorization
+      include Callbacks
       include ActiveSupport::Rescuable
 
       attr_reader :server, :env, :subscriptions, :logger, :worker_pool, :protocol
@@ -86,12 +87,18 @@ module ActionCable
 
       def dispatch_websocket_message(websocket_message) # :nodoc:
         if websocket.alive?
-          subscriptions.execute_command decode(websocket_message)
+          handle_channel_command decode(websocket_message)
         else
           logger.error "Ignoring message processed after the WebSocket was closed: #{websocket_message.inspect})"
         end
       end
 
+      def handle_channel_command(payload)
+        run_callbacks :command do
+          subscriptions.execute_command payload
+        end
+      end
+
       def transmit(cable_message) # :nodoc:
         websocket.transmit encode(cable_message)
       end

actioncable/lib/action_cable/connection/callbacks.rb

@@ -0,0 +1,30 @@
+# frozen_string_literal: true
+
+require "active_support/callbacks"
+
+module ActionCable
+  module Connection
+    module Callbacks
+      extend  ActiveSupport::Concern
+      include ActiveSupport::Callbacks
+
+      included do
+        define_callbacks :command
+      end
+
+      module ClassMethods
+        def before_command(*methods, &block)
+          set_callback(:command, :before, *methods, &block)
+        end
+
+        def after_command(*methods, &block)
+          set_callback(:command, :after, *methods, &block)
+        end
+
+        def around_command(*methods, &block)
+          set_callback(:command, :around, *methods, &block)
+        end
+      end
+    end
+  end
+end

actioncable/test/connection/callbacks_test.rb

@@ -0,0 +1,100 @@
+# frozen_string_literal: true
+
+require "test_helper"
+require "stubs/test_server"
+
+class ActionCable::Connection::CallbacksTest < ActionCable::TestCase
+  class Connection < ActionCable::Connection::Base
+    identified_by :context
+
+    attr_reader :commands_counter
+
+    before_command do
+      throw :abort unless context.nil?
+    end
+
+    around_command :set_current_context
+    after_command :increment_commands_counter
+
+    def initialize(*)
+      super
+      @commands_counter = 0
+    end
+
+    private
+      def set_current_context
+        self.context = request.params["context"]
+        yield
+      ensure
+        self.context = nil
+      end
+
+      def increment_commands_counter
+        @commands_counter += 1
+      end
+  end
+
+  class ChatChannel < ActionCable::Channel::Base
+    class << self
+      attr_accessor :words_spoken, :subscribed_count
+    end
+
+    self.words_spoken = []
+    self.subscribed_count = 0
+
+    def subscribed
+      self.class.subscribed_count += 1
+    end
+
+    def speak(data)
+      self.class.words_spoken << { data: data, context: context }
+    end
+  end
+
+  setup do
+    @server = TestServer.new
+    @env = Rack::MockRequest.env_for "/test", "HTTP_HOST" => "localhost", "HTTP_CONNECTION" => "upgrade", "HTTP_UPGRADE" => "websocket"
+    @connection = Connection.new(@server, @env)
+    @identifier = { channel: "ActionCable::Connection::CallbacksTest::ChatChannel" }.to_json
+  end
+
+  attr_reader :server, :env, :connection, :identifier
+
+  test "before and after callbacks" do
+    result = assert_difference -> { ChatChannel.subscribed_count }, +1 do
+      assert_difference -> { connection.commands_counter }, +1 do
+        connection.handle_channel_command({ "identifier" => identifier, "command" => "subscribe" })
+      end
+    end
+    assert result
+  end
+
+  test "before callback halts" do
+    connection.context = "non_null"
+    result = assert_no_difference -> { ChatChannel.subscribed_count } do
+      connection.handle_channel_command({ "identifier" => identifier, "command" => "subscribe" })
+    end
+    assert_not result
+  end
+
+  test "around_command callback" do
+    env["QUERY_STRING"] = "context=test"
+    connection = Connection.new(server, env)
+
+    assert_difference -> { ChatChannel.words_spoken.size }, +1 do
+      # We need to add subscriptions first
+      connection.handle_channel_command({
+        "identifier" => identifier,
+        "command" => "subscribe"
+      })
+      connection.handle_channel_command({
+        "identifier" => identifier,
+        "command" => "message",
+        "data" =>  { "action" => "speak", "message" => "hello" }.to_json
+      })
+    end
+
+    message = ChatChannel.words_spoken.last
+    assert_equal({ data: { "action" => "speak", "message" => "hello" }, context: "test" }, message)
+  end
+end