From acbc1a4fe9addca5dacd41fd405e0eae087f85ba Mon Sep 17 00:00:00 2001
From: GitLab Bot <gitlab-bot@gitlab.com>
Date: Sat, 26 Apr 2025 18:07:17 +0000
Subject: [PATCH] Add latest changes from gitlab-org/gitlab@master

---
 lib/gitlab/doctor/encryption_keys.rb          | 33 +++++++-----
 .../lib/gitlab/doctor/encryption_keys_spec.rb | 52 +++++++++++++++++++
 2 files changed, 73 insertions(+), 12 deletions(-)
 create mode 100644 spec/lib/gitlab/doctor/encryption_keys_spec.rb

diff --git a/lib/gitlab/doctor/encryption_keys.rb b/lib/gitlab/doctor/encryption_keys.rb
index 86cebd9d18a..7c12fafcf46 100644
--- a/lib/gitlab/doctor/encryption_keys.rb
+++ b/lib/gitlab/doctor/encryption_keys.rb
@@ -1,13 +1,10 @@
 # frozen_string_literal: true
 
-require 'base64'
-
 module Gitlab
   module Doctor
     class EncryptionKeys
       attr_reader :logger
 
-      PRINT_PROGRESS_EVERY = 1000
       KEY_TYPES = Gitlab::Encryption::KeyProvider::KEY_PROVIDERS.keys.grep(/active_record/)
       Key = Struct.new(:type, :id, :truncated_secret)
 
@@ -67,24 +64,36 @@ module Gitlab
 
         encryption_keys_usage = Hash.new { |hash, key| hash[key] = 0 }
 
-        model.find_each.with_index do |instance, index|
-          encrypted_attributes.each do |attribute_name|
-            encryption_keys_usage[encryption_key(instance, attribute_name)] += 1
+        encrypted_attributes.each do |attribute_name|
+          encryption_keys(model, attribute_name).each do |key_id, count|
+            encryption_keys_usage[key_id] += count
           end
-
-          logger.info "Checked #{index + 1}/#{total_count} #{model.name.pluralize}" if index % PRINT_PROGRESS_EVERY == 0
         end
-        logger.info "Checked #{total_count} #{model.name.pluralize}\n"
+        logger.info "Checked #{total_count} #{model.name}"
 
         encryption_keys_usage
       end
 
-      def encryption_key(instance, attribute_name)
-        Base64.decode64(Gitlab::Json.parse(instance.ciphertext_for(attribute_name))['h']['i'])
+      def encryption_keys(model, attr)
+        Hash[
+          model
+            .connection
+            .execute(
+              <<~SQL
+              SELECT #{attr}->'h'->'i' as key_id, COUNT(*) as usage_count
+              FROM #{model.table_name}
+              WHERE #{attr} IS NOT NULL
+              GROUP BY key_id
+              SQL
+            )
+            .filter_map { |a| a['key_id'] && [a['key_id'] && Base64.decode64(a['key_id']), a['usage_count']] }
+        ]
       end
 
       def models_with_encrypted_attributes
-        ApplicationRecord.descendants.select { |d| d.encrypted_attributes.present? }
+        Gitlab::Database.database_base_models.values.flat_map do |klass|
+          klass.descendants.select { |d| d.encrypted_attributes.present? }
+        end
       end
     end
   end
diff --git a/spec/lib/gitlab/doctor/encryption_keys_spec.rb b/spec/lib/gitlab/doctor/encryption_keys_spec.rb
new file mode 100644
index 00000000000..6e0ccf0745f
--- /dev/null
+++ b/spec/lib/gitlab/doctor/encryption_keys_spec.rb
@@ -0,0 +1,52 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::Doctor::EncryptionKeys, feature_category: :shared do
+  let(:logger) { instance_double(Logger).as_null_object }
+
+  subject(:doctor_encryption_secrets) { described_class.new(logger).run! }
+
+  it 'outputs current encryption secrets IDs, and truncated actual secrets' do
+    expect(logger).to receive(:info)
+      .with(/- active_record_encryption_primary_key: ID => `\w{4}`; truncated secret => `\w{3}...\w{3}`/)
+    expect(logger).to receive(:info)
+      .with(/- active_record_encryption_deterministic_key: ID => `\w{4}`; truncated secret => `\w{3}...\w{3}`/)
+
+    doctor_encryption_secrets
+  end
+
+  context 'when no encrypted attributes exist' do
+    it 'outputs "NONE"' do
+      expect(logger).to receive(:info).with(/Encryption keys usage for DependencyProxy::GroupSetting: NONE/)
+
+      doctor_encryption_secrets
+    end
+  end
+
+  context 'when encrypted attributes exist' do
+    # This will work in Rails 7.1.4, see https://github.com/rails/rails/issues/52003#issuecomment-2149673942
+    #
+    # let!(:key_provider1) { ActiveRecord::Encryption::DerivedSecretKeyProvider.new(SecureRandom.base64(32)) }
+    #
+    # before do
+    #   ActiveRecord::Encryption.with_encryption_context(key_provider: key_provider1) do
+    #     create(:dependency_proxy_group_setting)
+    #   end
+    # end
+    #
+    # Until then, we can only use the default key provider
+    let!(:key_provider1) { ActiveRecord::Encryption.key_provider }
+
+    before do
+      create(:dependency_proxy_group_setting)
+    end
+
+    it 'detects decryptable secrets' do
+      expect(logger).to receive(:info).with(/Encryption keys usage for DependencyProxy::GroupSetting:/)
+      expect(logger).to receive(:info).with(/- `#{key_provider1.encryption_key.id}` => 2/)
+
+      doctor_encryption_secrets
+    end
+  end
+end