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 · Apr 5, 2023 · c83b710 · c83b710 · nateberkopec · Apr 5, 2023
1 parent c13e925
commit c83b710
Show file tree

Hide file tree

Showing 3 changed files with 66 additions and 9 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)
@@ -347,7 +349,15 @@ def call(env)
         env['HTTP_ACCEPT_ENCODING'] = 'identity' if config.suppress_encoding
 
         if query_string =~ /pp=(async-)?flamegraph/ || env['HTTP_REFERER'] =~ /pp=async-flamegraph/
-          if defined?(StackProf) && StackProf.respond_to?(:run)
+          unless defined?(StackProf) && StackProf.respond_to?(:run)
+            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(
+              text_result(message, status: status, headers: headers)
+            )
+          else
             # do not sully our profile with mini profiler timings
             current.measure = false
             match_data      = query_string.match(/flamegraph_sample_rate=([\d\.]+)/)
@@ -642,9 +652,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
@@ -89,7 +89,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
 
@@ -105,29 +113,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,35 @@ 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 error response when stackprof isn't installed" do
+      response = profiler.call({ "PATH_INFO" => "/", "QUERY_STRING" => "pp=flamegraph" })
+
+      expect(response).to eq([
+        200,
+        { "Content-Type" => "text/plain; charset=utf-8", "Set-Cookie"=>"__profilin=p%3Dt; path=/; HttpOnly; SameSite=Lax" },
+        ["Please install the stackprof gem and require it: add gem 'stackprof' to your Gemfile"],
+      ])
+    end
+
+    it "returns error response when memory_profiler isn't installed" do
+      original_enable_advanced_debugging_tools = Rack::MiniProfiler.config.enable_advanced_debugging_tools
+      Rack::MiniProfiler.config.enable_advanced_debugging_tools = true
+
+      response = profiler.call({ "PATH_INFO" => "/", "QUERY_STRING" => "pp=profile-memory" })
+
+      expect(response).to eq([
+        200,
+        { "Content-Type" => "text/plain; charset=utf-8", "Set-Cookie"=>"__profilin=p%3Dt; path=/; HttpOnly; SameSite=Lax" },
+        ["Please install the memory_profiler gem and require it: add gem 'memory_profiler' to your Gemfile"],
+      ])
+
+    ensure
+      Rack::MiniProfiler.config.enable_advanced_debugging_tools = original_enable_advanced_debugging_tools
+    end
+  end
 end