Gracefully handles excluded fields from attributes during serialization

2 files changed, 54 insertions, 3 deletions

diff --git a/lib/gitlab/json_cache.rb b/lib/gitlab/json_cache.rb
index 24daad638f4..e4bc437d787 100644
--- a/lib/gitlab/json_cache.rb
+++ b/lib/gitlab/json_cache.rb
@@ -80,8 +80,23 @@ module Gitlab
       # when the new_record? method incorrectly returns false.
       #
       # See https://gitlab.com/gitlab-org/gitlab-ee/issues/9903#note_145329964
-      attributes = klass.attributes_builder.build_from_database(raw, {})
-      klass.allocate.init_with("attributes" => attributes, "new_record" => new_record?(raw, klass))
+      klass
+        .allocate
+        .init_with(
+          "attributes" => attributes_for(klass, raw),
+          "new_record" => new_record?(raw, klass)
+        )
+    end
+
+    def attributes_for(klass, raw)
+      # We have models that leave out some fields from the JSON export for
+      # security reasons, e.g. models that include the CacheMarkdownField.
+      # The ActiveRecord::AttributeSet we build from raw does know about
+      # these columns so we need manually set them.
+      missing_attributes = (klass.columns.map(&:name) - raw.keys)
+      missing_attributes.each { |column| raw[column] = nil }
+
+      klass.attributes_builder.build_from_database(raw, {})
     end
 
     def new_record?(raw, klass)
diff --git a/spec/lib/gitlab/json_cache_spec.rb b/spec/lib/gitlab/json_cache_spec.rb
index 2cae8ec031a..b7dc8234bdf 100644
--- a/spec/lib/gitlab/json_cache_spec.rb
+++ b/spec/lib/gitlab/json_cache_spec.rb
@@ -7,7 +7,7 @@ describe Gitlab::JsonCache do
   let(:namespace) { 'geo' }
   let(:key) { 'foo' }
   let(:expanded_key) { "#{namespace}:#{key}:#{Rails.version}" }
-  let(:broadcast_message) { create(:broadcast_message) }
+  set(:broadcast_message) { create(:broadcast_message) }
 
   subject(:cache) { described_class.new(namespace: namespace, backend: backend) }
 
@@ -321,6 +321,42 @@ describe Gitlab::JsonCache do
 
             expect(result).to be_new_record
           end
+
+          it 'gracefully handles bad cached entry' do
+            allow(backend).to receive(:read)
+              .with(expanded_key)
+              .and_return('{')
+
+            expect(cache.read(key, BroadcastMessage)).to be_nil
+          end
+
+          it 'gracefully handles an empty hash' do
+            allow(backend).to receive(:read)
+              .with(expanded_key)
+              .and_return('{}')
+
+            expect(cache.read(key, BroadcastMessage)).to be_a(BroadcastMessage)
+          end
+
+          it 'gracefully handles unknown attributes' do
+            allow(backend).to receive(:read)
+              .with(expanded_key)
+              .and_return(broadcast_message.attributes.merge(unknown_attribute: 1).to_json)
+
+            expect(cache.read(key, BroadcastMessage)).to be_nil
+          end
+
+          it 'gracefully handles excluded fields from attributes during serialization' do
+            backend.write(expanded_key, broadcast_message.to_json)
+
+            result = cache.fetch(key, as: BroadcastMessage) { 'block result' }
+
+            excluded_fields = BroadcastMessage.cached_markdown_fields.html_fields
+
+            (excluded_fields + ['cached_markdown_version']).each do |field|
+              expect(result.public_send(field)).to be_nil
+            end
+          end
         end
 
         it "returns the result of the block when 'as' option is nil" do