raise exceptions on header set after response committed

3 files changed, 43 insertions, 0 deletions

diff --git a/actionpack/lib/action_controller/metal/live.rb b/actionpack/lib/action_controller/metal/live.rb
index aea00849d3..efeeefda9d 100644
--- a/actionpack/lib/action_controller/metal/live.rb
+++ b/actionpack/lib/action_controller/metal/live.rb
@@ -1,4 +1,5 @@
 require 'action_dispatch/http/response'
+require 'delegate'
 
 module ActionController
   module Live
@@ -30,6 +31,26 @@ module ActionController
         end
       end
 
+      class Header < DelegateClass(Hash)
+        def initialize(response, header)
+          @response = response
+          super(header)
+        end
+
+        def []=(k,v)
+          if @response.committed?
+            raise ActionDispatch::IllegalStateError, 'header already sent'
+          end
+
+          super
+        end
+      end
+
+      def initialize(status = 200, header = {}, body = [])
+        header = Header.new self, header
+        super(status, header, body)
+      end
+
       private
 
       def build_buffer(response, body)
diff --git a/actionpack/lib/action_dispatch.rb b/actionpack/lib/action_dispatch.rb
index 1e4ac70f3d..c259b865cc 100644
--- a/actionpack/lib/action_dispatch.rb
+++ b/actionpack/lib/action_dispatch.rb
@@ -36,6 +36,9 @@ end
 module ActionDispatch
   extend ActiveSupport::Autoload
 
+  class IllegalStateError < StandardError
+  end
+
   autoload_under 'http' do
     autoload :Request
     autoload :Response
diff --git a/actionpack/test/dispatch/live_response_test.rb b/actionpack/test/dispatch/live_response_test.rb
index 70dd7524aa..d17d516a0b 100644
--- a/actionpack/test/dispatch/live_response_test.rb
+++ b/actionpack/test/dispatch/live_response_test.rb
@@ -40,6 +40,25 @@ module ActionController
         @response.stream.write 'omg'
         assert_nil @response.headers['Content-Length']
       end
+
+      def test_headers_cannot_be_written_after_write
+        @response.stream.write 'omg'
+
+        e = assert_raises(ActionDispatch::IllegalStateError) do
+          @response.headers['Content-Length'] = "zomg"
+        end
+
+        assert_equal 'header already sent', e.message
+      end
+
+      def test_headers_cannot_be_written_after_close
+        @response.stream.close
+
+        e = assert_raises(ActionDispatch::IllegalStateError) do
+          @response.headers['Content-Length'] = "zomg"
+        end
+        assert_equal 'header already sent', e.message
+      end
     end
   end
 end