Fix stackprof not installed error

Fixes malformed response when pp=flamegraph is used and stackprof is not installed. Also forwards headers and status code from app response error cases so that app logs are consistent with what is actually returned by middleware.
MiniProfiler · Feb 7, 2023 · 8a1285c · 8a1285c
1 parent edbcad3
commit 8a1285c
Show file tree

Hide file tree

Showing 3 changed files with 54 additions and 10 deletions.
diff --git a/lib/mini_profiler/profiler.rb b/lib/mini_profiler/profiler.rb
@@ -285,10 +285,12 @@ def call(env)
 
         unless defined?(MemoryProfiler) && MemoryProfiler.respond_to?(:report)
           message = "Please install the memory_profiler gem and require it: add gem 'memory_profiler' to your Gemfile"
-          _, _, body = @app.call(env)
+          status, headers, body = @app.call(env)
           body.close if body.respond_to? :close
 
-          return client_settings.handle_cookie(text_result(message))
+          return client_settings.handle_cookie(
+            text_result(message, status: status, headers: headers)
+          )
         end
 
         query_params = Rack::Utils.parse_nested_query(query_string)
@@ -348,10 +350,13 @@ def call(env)
 
         if query_string =~ /pp=(async-)?flamegraph/ || env['HTTP_REFERER'] =~ /pp=async-flamegraph/
           unless defined?(StackProf) && StackProf.respond_to?(:run)
-            headers = { 'Content-Type' => 'text/html' }
             message = "Please install the stackprof gem and require it: add gem 'stackprof' to your Gemfile"
+            status, headers, body = @app.call(env)
             body.close if body.respond_to? :close
-            return client_settings.handle_cookie([500, headers, message])
+
+            return client_settings.handle_cookie(
+              text_result(message, status: status, headers: headers)
+            )
           else
             # do not sully our profile with mini profiler timings
             current.measure = false
@@ -638,9 +643,9 @@ def analyze_memory
       text_result(body)
     end
 
-    def text_result(body)
-      headers = { 'Content-Type' => 'text/plain; charset=utf-8' }
-      [200, headers, [body]]
+    def text_result(body, status: 200, headers: nil)
+      headers = (headers || {}).merge('Content-Type' => 'text/plain; charset=utf-8')
+      [status, headers, [body]]
     end
 
     def make_link(postfix, env)

diff --git a/spec/integration/middleware_spec.rb b/spec/integration/middleware_spec.rb
@@ -43,7 +43,15 @@ def app
     def app
       Rack::Builder.new do
         use Rack::MiniProfiler
-        run lambda { |_env| [200, { 'Content-Type' => 'text/html' }, [+'<html><body><h1>Hi</h1></body></html>']] }
+        run(
+          lambda do |_env|
+            [
+              201,
+              { 'Content-Type' => 'text/html', 'X-CUSTOM' => "1" },
+              [+'<html><body><h1>Hi</h1></body></html>'],
+            ]
+          end
+        )
       end
     end
 
@@ -59,29 +67,37 @@ def app
     describe 'with profile-memory query' do
       it 'should return memory_profiler error message' do
         do_get(pp: 'profile-memory')
+
         expect(last_response.body).to eq(
           'Please install the memory_profiler gem and require it: add gem \'memory_profiler\' to your Gemfile'
         )
+        expect(last_response.headers['Content-Type']).to eq('text/plain; charset=utf-8')
+        expect(last_response.headers['X-CUSTOM']).to eq('1')
+        expect(last_response.status).to eq(201)
       end
     end
 
     describe 'with flamegraph query' do
       it 'should return stackprof error message' do
-        Rack::MiniProfiler.config.enable_advanced_debugging_tools = true
         do_get(pp: 'flamegraph')
         expect(last_response.body).to eq(
           'Please install the stackprof gem and require it: add gem \'stackprof\' to your Gemfile'
         )
+        expect(last_response.headers['Content-Type']).to eq('text/plain; charset=utf-8')
+        expect(last_response.headers['X-CUSTOM']).to eq('1')
+        expect(last_response.status).to eq(201)
       end
     end
 
     describe 'with async-flamegraph query' do
       it 'should return stackprof error message' do
-        Rack::MiniProfiler.config.enable_advanced_debugging_tools = true
         do_get(pp: 'async-flamegraph')
         expect(last_response.body).to eq(
           'Please install the stackprof gem and require it: add gem \'stackprof\' to your Gemfile'
         )
+        expect(last_response.headers['Content-Type']).to eq('text/plain; charset=utf-8')
+        expect(last_response.headers['X-CUSTOM']).to eq('1')
+        expect(last_response.status).to eq(201)
       end
     end
   end

diff --git a/spec/lib/profiler_spec.rb b/spec/lib/profiler_spec.rb
@@ -198,4 +198,27 @@ def self.bar(baz, boo)
       expect(Rack::MiniProfiler.snapshots_transporter?).to eq(true)
     end
   end
+
+  describe '#call' do
+    let(:app) { lambda { |env| [200, {}, ["OK"]] } }
+    let(:profiler) { Rack::MiniProfiler.new(app) }
+
+    it "returns well-formed response when stackprof isn't installed" do
+      response = profiler.call({ "PATH_INFO" => "/", "QUERY_STRING" => "pp=flamegraph" })
+      status, headers, body = response
+
+      expect(100..599).to include(status)
+      expect(headers.keys).to(all(be_a(String)))
+      expect(body).to(all(be_a(String)))
+    end
+
+    it "returns well-formed response when memory_profiler isn't installed" do
+      response = profiler.call({ "PATH_INFO" => "/", "QUERY_STRING" => "pp=profile-memory" })
+      status, headers, body = response
+
+      expect(100..599).to include(status)
+      expect(headers.keys).to(all(be_a(String)))
+      expect(body).to(all(be_a(String)))
+    end
+  end
 end