From a536cc8f1e31e09526b6b43c07ec255d66cd2049 Mon Sep 17 00:00:00 2001
From: Jon Moss <me@jonathanmoss.me>
Date: Wed, 24 Feb 2016 21:48:59 -0500
Subject: Only hijack Rack socket when first needed

Fixes #23471
---
 .../lib/action_cable/connection/client_socket.rb     |  1 +
 actioncable/lib/action_cable/connection/stream.rb    | 18 ++++++++----------
 actioncable/test/connection/base_test.rb             | 20 ++++++++++++++++++++
 3 files changed, 29 insertions(+), 10 deletions(-)

diff --git a/actioncable/lib/action_cable/connection/client_socket.rb b/actioncable/lib/action_cable/connection/client_socket.rb
index 95e1ac4c16..d7632b05fe 100644
--- a/actioncable/lib/action_cable/connection/client_socket.rb
+++ b/actioncable/lib/action_cable/connection/client_socket.rb
@@ -58,6 +58,7 @@ module ActionCable
 
       def start_driver
         return if @driver.nil? || @driver_started
+        @stream.hijack_rack_socket
         @driver_started = true
         @driver.start
       end
diff --git a/actioncable/lib/action_cable/connection/stream.rb b/actioncable/lib/action_cable/connection/stream.rb
index ace250cd16..2d97b28c09 100644
--- a/actioncable/lib/action_cable/connection/stream.rb
+++ b/actioncable/lib/action_cable/connection/stream.rb
@@ -4,15 +4,13 @@ module ActionCable
     # This class is heavily based on faye-websocket-ruby
     #
     # Copyright (c) 2010-2015 James Coglan
-    class Stream
+    class Stream # :nodoc:
       def initialize(event_loop, socket)
         @event_loop    = event_loop
         @socket_object = socket
         @stream_send   = socket.env['stream.send']
 
         @rack_hijack_io = nil
-
-        hijack_rack_socket
       end
 
       def each(&callback)
@@ -39,16 +37,16 @@ module ActionCable
         @socket_object.parse(data)
       end
 
-      private
-        def hijack_rack_socket
-          return unless @socket_object.env['rack.hijack']
+      def hijack_rack_socket
+        return unless @socket_object.env['rack.hijack']
 
-          @socket_object.env['rack.hijack'].call
-          @rack_hijack_io = @socket_object.env['rack.hijack_io']
+        @socket_object.env['rack.hijack'].call
+        @rack_hijack_io = @socket_object.env['rack.hijack_io']
 
-          @event_loop.attach(@rack_hijack_io, self)
-        end
+        @event_loop.attach(@rack_hijack_io, self)
+      end
 
+      private
         def clean_rack_hijack
           return unless @rack_hijack_io
           @event_loop.detach(@rack_hijack_io, self)
diff --git a/actioncable/test/connection/base_test.rb b/actioncable/test/connection/base_test.rb
index 3bef9e95a1..fb11f9be64 100644
--- a/actioncable/test/connection/base_test.rb
+++ b/actioncable/test/connection/base_test.rb
@@ -108,6 +108,26 @@ class ActionCable::Connection::BaseTest < ActionCable::TestCase
     end
   end
 
+  test "rejecting a connection causes a 404" do
+    run_in_eventmachine do
+      class CallMeMaybe
+        def call(*)
+          raise 'Do not call me!'
+        end
+      end
+
+      env = Rack::MockRequest.env_for(
+        "/test",
+        { 'HTTP_CONNECTION' => 'upgrade', 'HTTP_UPGRADE' => 'websocket',
+          'HTTP_ORIGIN' => 'http://rubyonrails.org', 'rack.hijack' => CallMeMaybe.new }
+      )
+
+      connection = ActionCable::Connection::Base.new(@server, env)
+      response = connection.process
+      assert_equal 404, response[0]
+    end
+  end
+
   private
     def open_connection
       env = Rack::MockRequest.env_for "/test", 'HTTP_CONNECTION' => 'upgrade', 'HTTP_UPGRADE' => 'websocket',
-- 
cgit v1.2.3