lib/hypertrace/instrumentation/rack_compatible.rb in hypertrace-agent-0.1.0 vs lib/hypertrace/instrumentation/rack_compatible.rb in hypertrace-agent-0.1.1

@@ -31,23 +31,50 @@
         content_type = content_type_arr[1]
         Hypertrace::Instrumentation::DataCapture.can_capture?(content_type, Hypertrace::Instrumentation::DataCapture::TYPE_RESPONSE)
       end
 
       def self.extract_response_body rack_response
+        if rack_response.is_a?(Rack::Response)
+          body = rack_response.respond_to?(:body) ? rack_response.body : nil
+          return Hypertrace::Instrumentation::DataCapture.capturable_body(body)
+        end
+
         if rack_response.is_a?(Rack::BodyProxy)
-          Hypertrace::Instrumentation::DataCapture.capturable_body(rack_response.join(''))
+          count = 0
+          while count < 15
+            rack_response = rack_response.instance_variable_get(:'@body')
+            if rack_response.is_a?(Array)
+              body = rack_response.respond_to?(:join) ? rack_response.join('') : nil
+              return Hypertrace::Instrumentation::DataCapture.capturable_body(body)
+            end
+            if action_dispatch_defined? && rack_response.is_a?(ActionDispatch::Response::RackBody)
+              return Hypertrace::Instrumentation::DataCapture.capturable_body(rack_response.body) if rack_response.respond_to?(:body)
+            end
+            unless rack_response.is_a?(Rack::BodyProxy)
+              return
+            end
+            count += 1
+          end
         elsif rack_response.is_a?(Array) && rack_response.length == 3
-          Hypertrace::Instrumentation::DataCapture.capturable_body(rack_response[2].join(''))
+          body = rack_response[2]
+          if body.is_a?(Array)
+            return Hypertrace::Instrumentation::DataCapture.capturable_body(body.join(''))
+          end
         end
       end
 
       def self.extract_response_headers_as_attributes rack_response, attr_prefix
         return unless Hypertrace::RubyAgent.instance.config.data_capture.http_headers.response
         return if rack_response.is_a?(Rack::BodyProxy)
 
         rack_response[1].map do |header, value|
           ["#{attr_prefix}.#{header}".downcase, value]
         end.to_h
+      end
+
+      # memoize so we don't check action dispatch
+      def self.action_dispatch_defined?
+        @_action_dispatch_defined ||= defined?(ActionDispatch) && defined?(ActionDispatch::Response) && defined?(ActionDispatch::Response::RackBody)
       end
 
       def self.extract_headers_from_env env
         result = env.select{|x|x.start_with?("HTTP_")}
         result.transform_keys!{|k|k[5..-1]}
\ No newline at end of file