From a7b28cc0d689fdafff0c40c1ad2301f725c3f23f Mon Sep 17 00:00:00 2001
From: Genadi Samokovarov <gsamokovarov@gmail.com>
Date: Mon, 7 Sep 2015 20:41:56 +0300
Subject: [PATCH] Render the console inside the body tag

Before this commit, the console was rendered straight after the closing
</html> tag, which made the page HTML invalid. However, browsers are
used to be molested and they are really good at displaying whatever you
throw at them, so that wasn't big problem anyway. Still it was something
that bug me quite a bit. I'm trying to redeem myself here.

Because I had to change the newly introduce `WebConsole::Response` class
I had cleaned the workspace a bit:

* `WebConsole::Response` no longer depends on a logger. The detection
  should be better now, so we don't need to couple the query check in
  it.

* `WebConsole::Response` does one thing only. It used to check the
  acceptance of a given `Content-Type`. I moved this to the middleware
  class.

* I have reverted the `WebConsole::WhinyRequest` to keep the logging out
  of the logic in the request class. We no longer need to log in two
  separate classes, so no need to break `Liskov` or introduce logging
  into simple query methods.
---
 lib/web_console.rb                     |  1 +
 lib/web_console/middleware.rb          | 19 +++++---
 lib/web_console/request.rb             | 21 +--------
 lib/web_console/response.rb            | 42 ++++++-----------
 lib/web_console/whiny_request.rb       | 31 +++++++++++++
 test/web_console/helper_test.rb        | 24 ++++++----
 test/web_console/middleware_test.rb    | 38 ++++++++++++----
 test/web_console/request_test.rb       | 16 -------
 test/web_console/response_test.rb      | 63 --------------------------
 test/web_console/whiny_request_test.rb | 25 ++++++++++
 10 files changed, 126 insertions(+), 154 deletions(-)
 create mode 100644 lib/web_console/whiny_request.rb
 delete mode 100644 test/web_console/response_test.rb
 create mode 100644 test/web_console/whiny_request_test.rb

diff --git a/lib/web_console.rb b/lib/web_console.rb
index f15c6034..997a295f 100644
--- a/lib/web_console.rb
+++ b/lib/web_console.rb
@@ -15,6 +15,7 @@
 require 'web_console/request'
 require 'web_console/response'
 require 'web_console/view'
+require 'web_console/whiny_request'
 
 module WebConsole
   mattr_accessor :logger
diff --git a/lib/web_console/middleware.rb b/lib/web_console/middleware.rb
index 322a74a5..1ef281a2 100644
--- a/lib/web_console/middleware.rb
+++ b/lib/web_console/middleware.rb
@@ -15,11 +15,7 @@ def initialize(app)
     end
 
     def call(env)
-      logger = if whiny_requests
-                 env['action_dispatch.logger'] || WebConsole.logger
-               end
-
-      request = Request.new(env, logger)
+      request = create_regular_or_whiny_request(env)
       return @app.call(env) unless request.from_whitelited_ip?
 
       if id = id_for_repl_session_update(request)
@@ -29,7 +25,6 @@ def call(env)
       end
 
       status, headers, body = @app.call(env)
-      response = Response.new(body, status, headers, logger)
 
       if exception = env['web_console.exception']
         session = Session.from_exception(exception)
@@ -37,7 +32,8 @@ def call(env)
         session = Session.from_binding(binding)
       end
 
-      if session && response.acceptable_content_type?
+      if session && acceptable_content_type?(headers)
+        response = Response.new(body, status, headers)
         template = Template.new(env, session)
 
         response.headers["X-Web-Console-Session-Id"] = session.id
@@ -51,6 +47,10 @@ def call(env)
 
     private
 
+      def acceptable_content_type?(headers)
+        Mime::Type.parse(headers['Content-Type']).first == Mime::HTML
+      end
+
       def json_response(opts = {})
         status  = opts.fetch(:status, 200)
         headers = { 'Content-Type' => 'application/json; charset = utf-8' }
@@ -66,6 +66,11 @@ def json_response_with_session(id, request, opts = {})
         json_response(opts) { yield session }
       end
 
+      def create_regular_or_whiny_request(env)
+        request = Request.new(env)
+        whiny_requests ? WhinyRequest.new(request) : request
+      end
+
       def repl_sessions_re
         @_repl_sessions_re ||= %r{#{mount_point}/repl_sessions/(?<id>[^/]+)}
       end
diff --git a/lib/web_console/request.rb b/lib/web_console/request.rb
index 5f6df2dc..ad2d80b3 100644
--- a/lib/web_console/request.rb
+++ b/lib/web_console/request.rb
@@ -8,22 +8,12 @@ class Request < ActionDispatch::Request
     # Define a vendor MIME type. We can call it using Mime::WEB_CONSOLE_V2 constant.
     Mime::Type.register 'application/vnd.web-console.v2', :web_console_v2
 
-    def initialize(env, logger = nil)
-      @logger = logger
-      super(env)
-    end
-
     # Returns whether a request came from a whitelisted IP.
     #
     # For a request to hit Web Console features, it needs to come from a white
     # listed IP.
     def from_whitelited_ip?
-      if whitelisted_ips.include?(strict_remote_ip)
-        true
-      else
-        log_whilelisted_ip_restriction
-        false
-      end
+      whitelisted_ips.include?(strict_remote_ip)
     end
 
     # Determines the remote IP using our much stricter whitelist.
@@ -38,15 +28,6 @@ def acceptable?
 
     private
 
-      attr_reader :logger
-
-      def log_whilelisted_ip_restriction
-        if logger
-          logger.info "Cannot render console from #{remote_ip}! " \
-            "Allowed networks: #{whitelisted_ips}"
-        end
-      end
-
       class GetSecureIp < ActionDispatch::RemoteIp::GetIp
         def initialize(req, proxies)
           # After rails/rails@07b2ff0 ActionDispatch::RemoteIp::GetIp initializes
diff --git a/lib/web_console/response.rb b/lib/web_console/response.rb
index e047e874..0fc3d947 100644
--- a/lib/web_console/response.rb
+++ b/lib/web_console/response.rb
@@ -1,37 +1,23 @@
 module WebConsole
-  # Web Console tailored response object.
-  class Response < Rack::Response
-    def initialize(body, status, headers, logger = nil)
-      @logger = logger
-      super(body, status, headers)
-    end
+  # A response object that writes content before the closing </body> tag, if
+  # possible.
+  #
+  # The object quacks like Rack::Response.
+  class Response < Struct.new(:body, :status, :headers)
+    def write(content)
+      raw_body = Array(body).first.to_s
 
-    # Returns whether the response is from an acceptable content type.
-    #
-    # We can render a console only for HTML responses.
-    def acceptable_content_type?
-      if content_type == Mime::HTML
-        true
+      if position = raw_body.rindex('</body>')
+        raw_body.insert(position, content)
       else
-        log_not_acceptable_content_type
-        false
+        raw_body << content
       end
-    end
 
-    def content_type
-      formats = Mime::Type.parse(headers['Content-Type'])
-      formats.first
+      self.body = raw_body
     end
 
-    private
-
-      attr_reader :logger
-
-      def log_not_acceptable_content_type
-        if logger
-          logger.info "Cannot render console with content type " \
-            "#{content_type}. Console can be rendered only in HTML responses"
-        end
-      end
+    def finish
+      Rack::Response.new(body, status, headers).finish
+    end
   end
 end
diff --git a/lib/web_console/whiny_request.rb b/lib/web_console/whiny_request.rb
new file mode 100644
index 00000000..75e9e4b6
--- /dev/null
+++ b/lib/web_console/whiny_request.rb
@@ -0,0 +1,31 @@
+module WebConsole
+  # Noisy wrapper around +Request+.
+  #
+  # If any calls to +from_whitelisted_ip?+ and +acceptable_content_type?+
+  # return false, an info log message will be displayed in users' logs.
+  class WhinyRequest < SimpleDelegator
+    def from_whitelited_ip?
+      whine_unless request.from_whitelited_ip? do
+        "Cannot render console from #{request.remote_ip}! " \
+          "Allowed networks: #{request.whitelisted_ips}"
+      end
+    end
+
+    private
+
+      def whine_unless(condition)
+        unless condition
+          logger.info { yield }
+        end
+        condition
+      end
+
+      def logger
+        env['action_dispatch.logger'] || WebConsole.logger
+      end
+
+      def request
+        __getobj__
+      end
+  end
+end
diff --git a/test/web_console/helper_test.rb b/test/web_console/helper_test.rb
index eef8df30..f2f1fc83 100644
--- a/test/web_console/helper_test.rb
+++ b/test/web_console/helper_test.rb
@@ -6,16 +6,7 @@ class BaseApplication
       include Helper
 
       def call(env)
-        Rack::Response.new(<<-HTML.strip_heredoc, status, headers).finish
-          <html>
-            <head>
-              <title>Hello world</title>
-            </head>
-            <body>
-              <p id="hello-world">Hello world</p>
-            </body>
-          </html>
-        HTML
+        [ status, headers, body ] 
       end
 
       private
@@ -31,6 +22,19 @@ def status
         def headers
           { 'Content-Type' => "#{Mime::HTML}; charset=utf-8" }
         end
+
+        def body
+          Array(<<-HTML.strip_heredoc)
+            <html>
+              <head>
+                <title>Hello world</title>
+              </head>
+              <body>
+                <p id="hello-world">Hello world</p>
+              </body>
+            </html>
+          HTML
+        end
     end
 
     class SingleConsoleApplication < BaseApplication
diff --git a/test/web_console/middleware_test.rb b/test/web_console/middleware_test.rb
index 866276ca..94272091 100644
--- a/test/web_console/middleware_test.rb
+++ b/test/web_console/middleware_test.rb
@@ -8,16 +8,7 @@ def initialize(options = {})
       end
 
       def call(env)
-        Rack::Response.new(<<-HTML.strip_heredoc, status, headers).finish
-          <html>
-            <head>
-              <title>Hello world</title>
-            </head>
-            <body>
-              <p id="hello-world">Hello world</p>
-            </body>
-          </html>
-        HTML
+        [ status, headers, body ]
       end
 
       private
@@ -29,6 +20,19 @@ def status
         def headers
           { 'Content-Type' => "#{@response_content_type}; charset=utf-8" }
         end
+
+        def body
+          Array(<<-HTML.strip_heredoc)
+            <html>
+              <head>
+                <title>Hello world</title>
+              </head>
+              <body>
+                <p id="hello-world">Hello world</p>
+              </body>
+            </html>
+          HTML
+        end
     end
 
     setup do
@@ -47,9 +51,23 @@ def headers
     test 'render console in an html application from web_console.exception' do
       get '/', nil, 'web_console.exception' => raise_exception
 
+      assert_select 'body > #console'
+    end
+
+    test 'render console if response format is HTML' do
+      @app = Middleware.new(Application.new(response_content_type: Mime::HTML))
+      get '/', nil, 'web_console.binding' => binding
+
       assert_select '#console'
     end
 
+    test 'does not render console if response format is not HTML' do
+      @app = Middleware.new(Application.new(response_content_type: Mime::JSON))
+      get '/', nil, 'web_console.binding' => binding
+
+      assert_select '#console', 0
+    end
+
     test 'returns X-Web-Console-Session-Id as response header' do
       get '/', nil, 'web_console.binding' => binding
 
diff --git a/test/web_console/request_test.rb b/test/web_console/request_test.rb
index ee5aafc4..5b013fe2 100644
--- a/test/web_console/request_test.rb
+++ b/test/web_console/request_test.rb
@@ -42,13 +42,6 @@ class RequestTest < ActiveSupport::TestCase
       assert_not req.from_whitelited_ip?
     end
 
-    test '#from_whitelited_ip? logs out to stderr' do
-      assert_output_to_stderr do
-        req = request_with_logger('http://example.com', 'REMOTE_ADDR' => '0.0.0.0')
-        assert_not req.from_whitelited_ip?
-      end
-    end
-
     test '#acceptable? is truthy for current version' do
       req = xhr('http://example.com', 'HTTP_ACCEPT' => "#{Mime::WEB_CONSOLE_V2}")
 
@@ -67,10 +60,6 @@ def request(*args)
         Request.new(mock_env(*args))
       end
 
-      def request_with_logger(*args)
-        Request.new(mock_env(*args), WebConsole.logger)
-      end
-
       def mock_env(*args)
         Rack::MockRequest.env_for(*args)
       end
@@ -79,10 +68,5 @@ def xhr(*args)
         args[1]['HTTP_X_REQUESTED_WITH'] ||= 'XMLHttpRequest'
         request(*args)
       end
-
-      def assert_output_to_stderr
-        output = capture(:stderr) { yield }
-        assert_not output.blank?
-      end
   end
 end
diff --git a/test/web_console/response_test.rb b/test/web_console/response_test.rb
deleted file mode 100644
index 5fadfe54..00000000
--- a/test/web_console/response_test.rb
+++ /dev/null
@@ -1,63 +0,0 @@
-require 'test_helper'
-
-module WebConsole
-  class ResponseTest < ActiveSupport::TestCase
-    test '#acceptable_content_type? is truthy if response format is HTML' do
-      res = response('Content-Type' => 'text/html; charset=utf-8')
-
-      assert res.acceptable_content_type?
-    end
-
-    test '#acceptable_content_type? is falsy if response format is not HTML' do
-      res = response('Content-Type' => 'application/json; charset=utf-8')
-
-      refute res.acceptable_content_type?
-    end
-
-    test '#acceptable_content_type? logs out to stderr if falsy' do
-      assert_output_to_stderr do
-        res = response({ 'Content-Type' => 'application/json; charset=utf-8' }, default_logger)
-        res.acceptable_content_type?
-      end
-    end
-
-    test '#acceptable_content_type? does not log out to stderr if truthy' do
-      assert_not_output_to_stderr do
-        res = response({ 'Content-Type' => 'text/html; charset=utf-8' }, default_logger)
-        res.acceptable_content_type?
-      end
-    end
-
-    test '#content_type returns Mime::HTML if content type is HTML' do
-      res = response('Content-Type' => 'text/html; charset=utf-8')
-
-      assert_equal Mime::HTML, res.content_type
-    end
-
-    test '#content_type returns Mime::JSON if content type is JSON' do
-      res = response('Content-Type' => 'application/json; charset=utf-8')
-
-      assert_equal Mime::JSON, res.content_type
-    end
-
-    private
-
-      def response(headers, logger = nil)
-        Response.new([], 200, headers, logger)
-      end
-
-      def default_logger
-        WebConsole.logger
-      end
-
-      def assert_output_to_stderr
-        output = capture(:stderr) { yield }
-        assert_not output.blank?
-      end
-
-      def assert_not_output_to_stderr
-        output = capture(:stderr) { yield }
-        assert output.blank?
-      end
-  end
-end
diff --git a/test/web_console/whiny_request_test.rb b/test/web_console/whiny_request_test.rb
new file mode 100644
index 00000000..07ce1d88
--- /dev/null
+++ b/test/web_console/whiny_request_test.rb
@@ -0,0 +1,25 @@
+require 'test_helper'
+
+module WebConsole
+  class WhinyRequestTest < ActiveSupport::TestCase
+    test '#from_whitelited_ip? logs out to stderr' do
+      Request.stubs(:whitelisted_ips).returns(IPAddr.new('127.0.0.1'))
+      assert_output_to_stderr do
+        req = request('http://example.com', 'REMOTE_ADDR' => '0.0.0.0')
+        assert_not req.from_whitelited_ip?
+      end
+    end
+
+    private
+
+      def assert_output_to_stderr
+        output = capture(:stderr) { yield }
+        assert_not output.blank?
+      end
+
+      def request(*args)
+        request = Request.new(Rack::MockRequest.env_for(*args))
+        WhinyRequest.new(request)
+      end
+  end
+end