Truncate large Reports

Gemfile

@@ -7,3 +7,4 @@ gem "minitest", "~> 5.23"
 gem "rake", "~> 13.2"
 
 gem "standard", group: %i[development test]
+gem "base64"

lib/telebugs.rb

@@ -17,6 +17,7 @@
 require_relative "telebugs/code_hunk"
 require_relative "telebugs/middleware"
 require_relative "telebugs/middleware_stack"
+require_relative "telebugs/truncator"
 
 module Telebugs
   # The general error that this library uses when it wants to raise.

lib/telebugs/report.rb

@@ -16,16 +16,38 @@ class Report
     # The maxium size of the JSON data in bytes
     MAX_REPORT_SIZE = 64000
 
+    # The maximum size of hashes, arrays and strings in the report.
+    DATA_MAX_SIZE = 10000
+
     attr_reader :data
     attr_accessor :ignored
 
     def initialize(error)
       @ignored = false
+      @truncator = Truncator.new(DATA_MAX_SIZE)
+
       @data = {
         errors: errors_as_json(error)
       }
     end
 
+    # Converts the report to JSON. Calls +to_json+ on each object inside the
+    # reports's payload. Truncates report data, JSON representation of which is
+    # bigger than {MAX_REPORT_SIZE}.
+    def to_json(*_args)
+      loop do
+        begin
+          json = @payload.to_json
+        rescue *JSON_EXCEPTIONS
+          # TODO: log the error
+        else
+          return json if json && json.bytesize <= MAX_REPORT_SIZE
+        end
+
+        break if truncate == 0
+      end
+    end
+
     private
 
     def errors_as_json(error)
@@ -38,8 +60,8 @@ def errors_as_json(error)
       end
     end
 
-    def attach_code(b)
-      b.each do |frame|
+    def attach_code(backtrace)
+      backtrace.each do |frame|
         next unless frame[:file]
         next unless File.exist?(frame[:file])
         next unless frame[:line]
@@ -55,7 +77,11 @@ def frame_belogns_to_root_directory?(frame)
     end
 
     def truncate
-      0
+      @data.each_key do |key|
+        @data[key] = @truncator.truncate(@data[key])
+      end
+
+      @truncator.reduce_max_size
     end
   end
 end

lib/telebugs/truncator.rb

@@ -0,0 +1,111 @@
+# frozen_string_literal: true
+
+module Telebugs
+  # This class is responsible for the truncation of too-big objects. Mainly, you
+  # should use it for simple objects such as strings, hashes, & arrays.
+  class Truncator
+    # The options for +String#encode+
+    ENCODING_OPTIONS = {invalid: :replace, undef: :replace}.freeze
+
+    # The temporary encoding to be used when fixing invalid strings with
+    # +ENCODING_OPTIONS+
+    TEMP_ENCODING = "utf-16"
+
+    # Encodings that are eligible for fixing invalid characters
+    SUPPORTED_ENCODINGS = [Encoding::UTF_8, Encoding::ASCII].freeze
+
+    # What to append when something is a circular reference
+    CIRCULAR = "[Circular]"
+
+    # What to append when something is truncated
+    TRUNCATED = "[Truncated]"
+
+    # The types that can contain references to itself
+    CIRCULAR_TYPES = [Array, Hash, Set].freeze
+
+    # Maximum size of hashes, arrays and strings
+    def initialize(max_size)
+      @max_size = max_size
+    end
+
+    # Performs deep truncation of arrays, hashes, sets & strings. Uses a
+    # placeholder for recursive objects (`[Circular]`).
+    def truncate(object, seen = Set.new)
+      if seen.include?(object.object_id)
+        return CIRCULAR if CIRCULAR_TYPES.any? { |t| object.is_a?(t) }
+
+        return object
+      end
+      truncate_object(object, seen << object.object_id)
+    end
+
+    # Reduces maximum allowed size of hashes, arrays, sets & strings by half.
+    def reduce_max_size
+      @max_size /= 2
+    end
+
+    private
+
+    def truncate_object(object, seen)
+      case object
+      when Hash then truncate_hash(object, seen)
+      when Array then truncate_array(object, seen)
+      when Set then truncate_set(object, seen)
+      when String then truncate_string(object)
+      when Numeric, TrueClass, FalseClass, Symbol, NilClass then object
+      else
+        truncate_string(stringify_object(object))
+      end
+    end
+
+    def truncate_string(str)
+      fixed_str = replace_invalid_characters(str)
+      return fixed_str if fixed_str.length <= @max_size
+
+      (fixed_str.slice(0, @max_size) + TRUNCATED).freeze
+    end
+
+    def stringify_object(object)
+      object.to_json
+    rescue *Report::JSON_EXCEPTIONS
+      object.to_s
+    end
+
+    def truncate_hash(hash, seen)
+      truncated_hash = {}
+      hash.each_with_index do |(key, val), idx|
+        break if idx + 1 > @max_size
+
+        truncated_hash[key] = truncate(val, seen)
+      end
+
+      truncated_hash.freeze
+    end
+
+    def truncate_array(array, seen)
+      array.slice(0, @max_size).map! { |elem| truncate(elem, seen) }.freeze
+    end
+
+    def truncate_set(set, seen)
+      truncated_set = Set.new
+
+      set.each do |elem|
+        truncated_set << truncate(elem, seen)
+        break if truncated_set.size >= @max_size
+      end
+
+      truncated_set.freeze
+    end
+
+    # Replaces invalid characters in a string with arbitrary encoding.
+    # @see https://github.com/flori/json/commit/3e158410e81f94dbbc3da6b7b35f4f64983aa4e3
+    def replace_invalid_characters(str)
+      utf8_string = SUPPORTED_ENCODINGS.include?(str.encoding)
+      return str if utf8_string && str.valid_encoding?
+
+      temp_str = str.dup
+      temp_str.encode!(TEMP_ENCODING, **ENCODING_OPTIONS) if utf8_string
+      temp_str.encode!("utf-8", **ENCODING_OPTIONS)
+    end
+  end
+end

test/test_helper.rb

@@ -5,6 +5,7 @@
 
 require "minitest/autorun"
 require "webmock/minitest"
+require "base64"
 
 class OCIError < StandardError; end
 

test/test_truncator.rb

@@ -0,0 +1,199 @@
+# frozen_string_literal: true
+
+require "test_helper"
+
+class TelebugsTruncatorTest < Minitest::Test
+  def multiply_by_2_max_len(chr, max_len)
+    chr * 2 * max_len
+  end
+
+  def setup
+    @max_size = 3
+    @truncated = "[Truncated]"
+    @max_len = @max_size + @truncated.length
+  end
+
+  def test_truncate_frozen_string
+    object = multiply_by_2_max_len("a", @max_size)
+    truncator = Telebugs::Truncator.new(@max_size).truncate(object)
+
+    assert_equal @max_len, truncator.length
+    assert truncator.frozen?
+  end
+
+  def test_truncate_frozen_hash_of_strings
+    object = {
+      banana: multiply_by_2_max_len("a", @max_size),
+      kiwi: multiply_by_2_max_len("b", @max_size),
+      strawberry: "c",
+      shrimp: "d"
+    }.freeze
+
+    truncator = Telebugs::Truncator.new(@max_size).truncate(object)
+
+    assert_equal @max_size, truncator.size
+    assert truncator.frozen?
+    assert_equal({banana: "aaa[Truncated]", kiwi: "bbb[Truncated]", strawberry: "c"}, truncator)
+    assert truncator[:banana].frozen?
+    assert truncator[:kiwi].frozen?
+    assert truncator[:strawberry].frozen?
+  end
+
+  def test_truncate_frozen_array_of_strings
+    object = [
+      multiply_by_2_max_len("a", @max_size),
+      "b",
+      multiply_by_2_max_len("c", @max_size),
+      "d"
+    ].freeze
+
+    truncator = Telebugs::Truncator.new(@max_size).truncate(object)
+
+    assert_equal @max_size, truncator.size
+    assert truncator.frozen?
+    assert_equal ["aaa[Truncated]", "b", "ccc[Truncated]"], truncator
+    assert truncator[0].frozen?
+    assert truncator[1].frozen?
+    assert truncator[2].frozen?
+  end
+
+  def test_truncate_frozen_set_of_strings
+    object = Set.new([
+      multiply_by_2_max_len("a", @max_size),
+      "b",
+      multiply_by_2_max_len("c", @max_size),
+      "d"
+    ]).freeze
+
+    truncator = Telebugs::Truncator.new(@max_size).truncate(object)
+
+    assert_equal @max_size, truncator.size
+    assert truncator.frozen?
+    assert_equal Set.new(["aaa[Truncated]", "b", "ccc[Truncated]"]), truncator
+  end
+
+  def test_truncate_frozen_object_with_to_json
+    object = Object.new
+    def object.to_json
+      '{"object":"shrimp"}'
+    end
+    object.freeze
+
+    truncator = Telebugs::Truncator.new(@max_size).truncate(object)
+
+    assert_equal @max_len, truncator.length
+    assert truncator.frozen?
+    assert_equal '{"o[Truncated]', truncator
+  end
+
+  def test_truncate_object_without_to_json
+    object = Object.new
+    def object.to_json
+      raise Telebugs::Report::JSON_EXCEPTIONS.first
+    end
+
+    truncator = Telebugs::Truncator.new(@max_size).truncate(object)
+
+    assert_equal @max_len, truncator.length
+    assert_equal "#<O[Truncated]", truncator
+  end
+
+  def test_self_returning_objects
+    [1, true, false, :symbol].each do |object|
+      assert_equal object, Telebugs::Truncator.new(@max_size).truncate(object)
+    end
+
+    assert_nil Telebugs::Truncator.new(@max_size).truncate(nil)
+  end
+
+  def test_recursive_array
+    a = %w[aaaaa bb]
+    a << a
+    a << "c"
+
+    truncator = Telebugs::Truncator.new(@max_size).truncate(a)
+
+    assert_equal ["aaa[Truncated]", "bb", "[Circular]"], truncator
+  end
+
+  def test_recursive_array_with_recursive_hashes
+    a = []
+    a << a
+
+    h = {}
+    h[:k] = h
+    a << h << "aaaa"
+
+    truncator = Telebugs::Truncator.new(@max_size).truncate(a)
+
+    assert_equal ["[Circular]", {k: "[Circular]"}, "aaa[Truncated]"], truncator
+    assert truncator.frozen?
+  end
+
+  def test_recursive_set_with_recursive_arrays
+    s = Set.new
+    s << s
+
+    h = {}
+    h[:k] = h
+    s << h << "aaaa"
+
+    truncator = Telebugs::Truncator.new(@max_size).truncate(s)
+
+    assert_equal Set.new(["[Circular]", {k: "[Circular]"}, "aaa[Truncated]"]), truncator
+    assert truncator.frozen?
+  end
+
+  def test_hash_with_long_strings
+    object = {
+      a: multiply_by_2_max_len("a", @max_size),
+      b: multiply_by_2_max_len("b", @max_size),
+      c: {d: multiply_by_2_max_len("d", @max_size), e: "e"}
+    }
+
+    truncator = Telebugs::Truncator.new(@max_size).truncate(object)
+
+    assert_equal(
+      {a: "aaa[Truncated]", b: "bbb[Truncated]", c: {d: "ddd[Truncated]", e: "e"}},
+      truncator
+    )
+    assert truncator.frozen?
+  end
+
+  def test_string_with_valid_unicode_characters
+    object = "€€€€€"
+    truncator = Telebugs::Truncator.new(@max_size).truncate(object)
+
+    assert_equal "€€€[Truncated]", truncator
+  end
+
+  def test_ascii_8bit_string_with_invalid_characters
+    encoded = Base64.encode64("\xD3\xE6\xBC\x9D\xBA").encode!("ASCII-8BIT")
+    object = Base64.decode64(encoded).freeze
+
+    truncator = Telebugs::Truncator.new(@max_size).truncate(object)
+
+    assert_equal "���[Truncated]", truncator
+    assert truncator.frozen?
+  end
+
+  def test_array_with_hashes_and_hash_like_objects_with_identical_keys
+    hashie = Class.new(Hash)
+
+    object = {
+      errors: [
+        {file: "a"},
+        {file: "a"},
+        hashie.new.merge(file: "bcde")
+      ]
+    }
+
+    truncator = Telebugs::Truncator.new(@max_size).truncate(object)
+
+    assert_equal(
+      {errors: [{file: "a"}, {file: "a"}, hashie.new.merge(file: "bcd[Truncated]")]},
+      truncator
+    )
+    assert truncator.frozen?
+  end
+end