Render the console inside the body tag

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

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,15 +25,15 @@ 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)
       elsif binding = env['web_console.binding']
         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

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

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

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

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

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
 

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