feat: enable last resort truncation strategy, delete (nearly) any key

lib/rollbar/truncation.rb

@@ -6,6 +6,7 @@
 require 'rollbar/truncation/min_body_strategy'
 require 'rollbar/truncation/remove_request_strategy'
 require 'rollbar/truncation/remove_extra_strategy'
+require 'rollbar/truncation/remove_any_key_strategy'
 
 module Rollbar
   module Truncation
@@ -17,7 +18,8 @@ module Truncation
                   StringsStrategy,
                   MinBodyStrategy,
                   RemoveRequestStrategy,
-                  RemoveExtraStrategy].freeze
+                  RemoveExtraStrategy,
+                  RemoveAnyKeyStrategy].freeze
 
     def self.truncate(payload, attempts = [])
       result = nil

lib/rollbar/truncation/remove_any_key_strategy.rb

@@ -0,0 +1,113 @@
+require 'rollbar/util'
+
+module Rollbar
+  module Truncation
+    class RemoveAnyKeyStrategy
+      include ::Rollbar::Truncation::Mixin
+
+      attr_accessor :payload, :data, :sizes, :extracted_title
+
+      def self.call(payload)
+        new(payload).call
+      end
+
+      def initialize(payload)
+        @payload = payload
+        @data = payload['data']
+        @extracted_title = extract_title(data['body']) if data['body']
+      end
+
+      def call
+        remove_unknown_root_keys
+
+        json_payload = remove_oversized_data_keys
+
+        return json_payload if json_payload
+
+        dump(payload)
+      end
+
+      def remove_unknown_root_keys
+        payload.keys.reject { |key| root_keys.include?(key) }.each do |key|
+          truncation_key['root'] ||= {}
+          size = dump(payload.delete(key)).bytesize
+          truncation_key['root'][key] = "unknown root key removed, size: #{size} bytes"
+        end
+      end
+
+      def remove_oversized_data_keys
+        data_keys.keys.sort { |a, b| data_keys[b] <=> data_keys[a] }.each do |key|
+          json_payload = remove_key_and_return_payload(key)
+
+          return json_payload unless truncate?(json_payload)
+        end
+
+        false
+      end
+
+      def remove_key_and_return_payload(key)
+        size = data_keys[key]
+
+        data.delete(key)
+
+        replace_message_body if key == 'body'
+
+        truncation_key[key] = "key removed, size: #{size} bytes"
+
+        dump(payload)
+      end
+
+      def replace_message_body
+        data['body'] = message_key
+        data['title'] ||= extracted_title if extracted_title
+      end
+
+      def truncation_key
+        @truncation_key ||=
+          # initialize the diagnostic key for truncation
+          (data['notifier']['diagnostic'] ||= {}) &&
+          (data['notifier']['diagnostic']['truncation'] ||= {})
+      end
+
+      def root_keys
+        # Valid keys in root of payload
+        %w[access_token data]
+      end
+
+      def skip_keys
+        # Don't try to truncate these data keys
+        %w[notifier uuid title platform language framework level]
+      end
+
+      def message_key
+        # use this message if data.body gets removed
+        {
+          'message' => {
+            'body' => 'Payload keys removed due to oversized payload. See diagnostic key'
+          }
+        }
+      end
+
+      def extract_title(body)
+        return body['message']['body'] if body['message'] && body['message']['body']
+        return extract_title_from_trace(body['trace']) if body['trace']
+        return extract_title_from_trace(body['trace_chain'][0]) if body['trace_chain'] && body['trace_chain'][0]
+      end
+
+      def extract_title_from_trace(trace)
+        exception = trace['exception']
+
+        "#{exception['class']}: #{exception['message']}"
+      end
+
+      def data_keys
+        @data_keys ||= {}.tap do |hash|
+          data.keys.reject { |key| skip_keys.include?(key) }.each do |key|
+            size = dump(data[key]).bytesize
+            hash[key] = size
+          end
+        end
+      end
+    end
+  end
+end

spec/rollbar/truncation/remove_any_key_strategy_spec.rb

@@ -0,0 +1,65 @@
+require 'spec_helper'
+require 'rollbar/truncation/remove_any_key_strategy'
+
+describe Rollbar::Truncation::RemoveAnyKeyStrategy do
+  describe '.call' do
+    let(:exception_class) { 'ExceptionClass' }
+    let(:exception_message) { 'Exception message' }
+    let(:body) do
+      {
+        'trace' => {
+          'exception' => {
+            'class' => exception_class,
+            'message' => exception_message
+          }
+        },
+        'foo' => 'bar' * 999
+      }
+    end
+    let(:request) { { 'bar' => 'baz' } }
+    let(:payload) do
+      {
+        'data' => {
+          'body' => body,
+          'request' => request,
+          'notifier' => {}
+        },
+        'unknown_root_key' => { 'foo' => 'bar' }
+      }
+    end
+    let(:truncation_message) do
+      {
+        'message' => {
+          'body' => 'Payload keys removed due to oversized payload. See diagnostic key'
+        }
+      }
+    end
+    let(:diagnostic) do
+      {
+        'diagnostic' => {
+          'truncation' => {
+            'body' => 'key removed, size: 3086 bytes',
+            'root' => {
+              'unknown_root_key' => 'unknown root key removed, size: 13 bytes'
+            }
+          }
+        }
+      }
+    end
+
+    it 'should remove unknown and oversized keys in the payload' do
+      result = Rollbar::JSON.load(described_class.call(Rollbar::Util.deep_copy(payload)))
+
+      original_payload_size = Rollbar::Truncation::MAX_PAYLOAD_SIZE
+      Rollbar::Truncation::MAX_PAYLOAD_SIZE = 200
+
+      expect(result['data']['body']).to be_eql(truncation_message)
+      expect(result['data']['request']).to be_eql(request)
+      expect(result['data']['title']).to be_eql([exception_class, exception_message].join(': '))
+      expect(result['unknown_root_key']).to be_nil
+      expect(result['data']['notifier']).to be_eql(diagnostic)
+
+      Rollbar::Truncation::MAX_PAYLOAD_SIZE = original_payload_size
+    end
+  end
+end