abceda6cc5
It keeps track of the memory being used when loading the YAML file as well as the depth of nesting. Track exception when YAML is too big
79 lines
1.4 KiB
Ruby
79 lines
1.4 KiB
Ruby
# frozen_string_literal: true
|
|
|
|
require 'objspace'
|
|
|
|
module Gitlab
|
|
module Utils
|
|
class DeepSize
|
|
Error = Class.new(StandardError)
|
|
TooMuchDataError = Class.new(Error)
|
|
|
|
DEFAULT_MAX_SIZE = 1.megabyte
|
|
DEFAULT_MAX_DEPTH = 100
|
|
|
|
def initialize(root, max_size: DEFAULT_MAX_SIZE, max_depth: DEFAULT_MAX_DEPTH)
|
|
@root = root
|
|
@max_size = max_size
|
|
@max_depth = max_depth
|
|
@size = 0
|
|
@depth = 0
|
|
|
|
evaluate
|
|
end
|
|
|
|
def valid?
|
|
!too_big? && !too_deep?
|
|
end
|
|
|
|
private
|
|
|
|
def evaluate
|
|
add_object(@root)
|
|
rescue Error
|
|
# NOOP
|
|
end
|
|
|
|
def too_big?
|
|
@size > @max_size
|
|
end
|
|
|
|
def too_deep?
|
|
@depth > @max_depth
|
|
end
|
|
|
|
def add_object(object)
|
|
@size += ObjectSpace.memsize_of(object)
|
|
raise TooMuchDataError if @size > @max_size
|
|
|
|
add_array(object) if object.is_a?(Array)
|
|
add_hash(object) if object.is_a?(Hash)
|
|
end
|
|
|
|
def add_array(object)
|
|
with_nesting do
|
|
object.each do |n|
|
|
add_object(n)
|
|
end
|
|
end
|
|
end
|
|
|
|
def add_hash(object)
|
|
with_nesting do
|
|
object.each do |key, value|
|
|
add_object(key)
|
|
add_object(value)
|
|
end
|
|
end
|
|
end
|
|
|
|
def with_nesting
|
|
@depth += 1
|
|
raise TooMuchDataError if too_deep?
|
|
|
|
yield
|
|
|
|
@depth -= 1
|
|
end
|
|
end
|
|
end
|
|
end
|