1 files changed, 27 insertions, 38 deletions

diff --git a/lib/gitlab/config/loader/multi_doc_yaml.rb b/lib/gitlab/config/loader/multi_doc_yaml.rb
index 346adc79896..34080d26b7c 100644
--- a/lib/gitlab/config/loader/multi_doc_yaml.rb
+++ b/lib/gitlab/config/loader/multi_doc_yaml.rb
@@ -4,59 +4,48 @@ module Gitlab
   module Config
     module Loader
       class MultiDocYaml
-        TooManyDocumentsError = Class.new(Loader::FormatError)
-        DataTooLargeError = Class.new(Loader::FormatError)
-        NotHashError = Class.new(Loader::FormatError)
+        include Gitlab::Utils::StrongMemoize
 
-        MULTI_DOC_DIVIDER = /^---$/.freeze
+        MULTI_DOC_DIVIDER = /^---\s+/.freeze
 
         def initialize(config, max_documents:, additional_permitted_classes: [])
+          @config = config
           @max_documents = max_documents
-          @safe_config = load_config(config, additional_permitted_classes)
+          @additional_permitted_classes = additional_permitted_classes
         end
 
-        def load!
-          raise TooManyDocumentsError, 'The parsed YAML has too many documents' if too_many_documents?
-          raise DataTooLargeError, 'The parsed YAML is too big' if too_big?
-          raise NotHashError, 'Invalid configuration format' unless all_hashes?
-
-          safe_config.map(&:deep_symbolize_keys)
+        def valid?
+          documents.all?(&:valid?)
         end
 
-        private
-
-        attr_reader :safe_config, :max_documents
-
-        def load_config(config, additional_permitted_classes)
-          config.split(MULTI_DOC_DIVIDER).filter_map do |document|
-            YAML.safe_load(document,
-              permitted_classes: [Symbol, *additional_permitted_classes],
-              permitted_symbols: [],
-              aliases: true
-            )
-          end
-        rescue Psych::Exception => e
-          raise Loader::FormatError, e.message
+        def load_raw!
+          documents.map(&:load_raw!)
         end
 
-        def all_hashes?
-          safe_config.all?(Hash)
+        def load!
+          documents.map(&:load!)
         end
 
-        def too_many_documents?
-          safe_config.count > max_documents
-        end
+        private
+
+        attr_reader :config, :max_documents, :additional_permitted_classes
+
+        # Valid YAML files can start with either a leading delimiter or no delimiter.
+        # To avoid counting a leading delimiter towards the document limit,
+        # this method splits the file by one more than the maximum number of permitted documents.
+        # It then discards the first document if it is blank.
+        def documents
+          docs = config
+                  .split(MULTI_DOC_DIVIDER, max_documents_including_leading_delimiter)
+                  .map { |d| Yaml.new(d, additional_permitted_classes: additional_permitted_classes) }
 
-        def too_big?
-          !deep_sizes.all?(&:valid?)
+          docs.shift if docs.first.blank?
+          docs
         end
+        strong_memoize_attr :documents
 
-        def deep_sizes
-          safe_config.map do |config|
-            Gitlab::Utils::DeepSize.new(config,
-              max_size: Gitlab::CurrentSettings.current_application_settings.max_yaml_size_bytes,
-              max_depth: Gitlab::CurrentSettings.current_application_settings.max_yaml_depth)
-          end
+        def max_documents_including_leading_delimiter
+          max_documents + 1
         end
       end
     end