Encrypt static_object_token_encrypted field via background migration

Contributes to https://gitlab.com/gitlab-org/gitlab/-/issues/348163

Uses a background migration to generate encrypted versions of the
static_object_token values.

Changelog: changed

app/models/user.rb

@@ -48,7 +48,7 @@ class User < ApplicationRecord
 
   add_authentication_token_field :incoming_email_token, token_generator: -> { SecureRandom.hex.to_i(16).to_s(36) }
   add_authentication_token_field :feed_token
-  add_authentication_token_field :static_object_token
+  add_authentication_token_field :static_object_token, encrypted: :optional
 
   default_value_for :admin, false
   default_value_for(:external) { Gitlab::CurrentSettings.user_default_external }

db/post_migrate/20211210140000_add_temporary_static_object_token_index.rb

+# frozen_string_literal: true
+
+class AddTemporaryStaticObjectTokenIndex < Gitlab::Database::Migration[1.0]
+  disable_ddl_transaction!
+
+  INDEX_NAME = 'index_users_with_static_object_token'
+
+  def up
+    add_concurrent_index :users, :id, where: "static_object_token IS NOT NULL AND static_object_token_encrypted IS NULL", name: INDEX_NAME
+  end
+
+  def down
+    remove_concurrent_index :users, :id, name: INDEX_NAME
+  end
+end

db/post_migrate/20211210140629_encrypt_static_object_token.rb

+# frozen_string_literal: true
+
+class EncryptStaticObjectToken < Gitlab::Database::Migration[1.0]
+  BATCH_SIZE = 10_000
+  MIGRATION = 'EncryptStaticObjectToken'
+
+  disable_ddl_transaction!
+
+  def up
+    queue_background_migration_jobs_by_range_at_intervals(
+      define_batchable_model('users').where.not(static_object_token: nil).where(static_object_token_encrypted: nil),
+      MIGRATION,
+      2.minutes,
+      batch_size: BATCH_SIZE,
+      track_jobs: true
+    )
+  end
+
+  def down
+    # no ops
+  end
+end

db/schema_migrations/20211210140000

+f02c1b7412d2bb6d8a20639704ad55cdbcc14bfccf0509b659c3ef9614bcfa2b
\ No newline at end of file

db/schema_migrations/20211210140629

+7940b0f692b62bcabbe98440082e2245fd28caba2c9e052e85e82acea0a98d23
\ No newline at end of file

db/structure.sql

@@ -27795,6 +27795,8 @@ CREATE INDEX index_users_star_projects_on_project_id ON users_star_projects USIN
 
 CREATE UNIQUE INDEX index_users_star_projects_on_user_id_and_project_id ON users_star_projects USING btree (user_id, project_id);
 
+CREATE INDEX index_users_with_static_object_token ON users USING btree (id) WHERE ((static_object_token IS NOT NULL) AND (static_object_token_encrypted IS NULL));
+
 CREATE UNIQUE INDEX index_verification_codes_on_phone_and_visitor_id_code ON ONLY verification_codes USING btree (visitor_id_code, phone, created_at);
 
 COMMENT ON INDEX index_verification_codes_on_phone_and_visitor_id_code IS 'JiHu-specific index';

lib/gitlab/background_migration/encrypt_static_object_token.rb

+# frozen_string_literal: true
+
+module Gitlab
+  module BackgroundMigration
+    # Populates "static_object_token_encrypted" field with encrypted versions
+    # of values from "static_object_token" field
+    class EncryptStaticObjectToken
+      # rubocop:disable Style/Documentation
+      class User < ActiveRecord::Base
+        include ::EachBatch
+        self.table_name = 'users'
+        scope :with_static_object_token, -> { where.not(static_object_token: nil) }
+        scope :without_static_object_token_encrypted, -> { where(static_object_token_encrypted: nil) }
+      end
+      # rubocop:enable Style/Documentation
+
+      BATCH_SIZE = 100
+
+      def perform(start_id, end_id)
+        ranged_query = User
+          .where(id: start_id..end_id)
+          .with_static_object_token
+          .without_static_object_token_encrypted
+
+        ranged_query.each_batch(of: BATCH_SIZE) do |sub_batch|
+          first, last = sub_batch.pluck(Arel.sql('min(id), max(id)')).first
+
+          batch_query = User.unscoped
+                          .where(id: first..last)
+                          .with_static_object_token
+                          .without_static_object_token_encrypted
+
+          user_tokens = batch_query.pluck(:id, :static_object_token)
+
+          user_encrypted_tokens = user_tokens.map do |(id, plaintext_token)|
+            next if plaintext_token.blank?
+
+            [id, Gitlab::CryptoHelper.aes256_gcm_encrypt(plaintext_token)]
+          end
+
+          encrypted_tokens_sql = user_encrypted_tokens.compact.map { |(id, token)| "(#{id}, '#{token}')" }.join(',')
+
+          if user_encrypted_tokens.present?
+            User.connection.execute(<<~SQL)
+              WITH cte(cte_id, cte_token) AS #{::Gitlab::Database::AsWithMaterialized.materialized_if_supported} (
+                SELECT *
+                FROM (VALUES #{encrypted_tokens_sql}) AS t (id, token)
+              )
+              UPDATE #{User.table_name}
+              SET static_object_token_encrypted = cte_token
+              FROM cte
+              WHERE cte_id = id
+            SQL
+          end
+
+          mark_job_as_succeeded(start_id, end_id)
+        end
+      end
+
+      private
+
+      def mark_job_as_succeeded(*arguments)
+        Gitlab::Database::BackgroundMigrationJob.mark_all_as_succeeded(
+          self.class.name.demodulize,
+          arguments
+        )
+      end
+    end
+  end
+end

spec/controllers/projects/repositories_controller_spec.rb

@@ -187,6 +187,25 @@ RSpec.describe Projects::RepositoriesController do
             expect(response).to have_gitlab_http_status(:found)
           end
         end
+
+        context 'when token is migrated' do
+          let(:user) { create(:user, static_object_token: '') }
+          let(:token) { 'Test' }
+
+          it 'calls the action normally' do
+            user.update_column(:static_object_token, token)
+
+            get :archive, params: { namespace_id: project.namespace, project_id: project, id: 'master', token: token }, format: 'zip'
+            expect(user.static_object_token).to eq(token)
+            expect(response).to have_gitlab_http_status(:ok)
+
+            user.update_column(:static_object_token_encrypted, Gitlab::CryptoHelper.aes256_gcm_encrypt(token))
+
+            get :archive, params: { namespace_id: project.namespace, project_id: project, id: 'master', token: token }, format: 'zip'
+            expect(user.static_object_token).to eq(token)
+            expect(response).to have_gitlab_http_status(:ok)
+          end
+        end
       end
 
       context 'when a token header is present' do

spec/lib/gitlab/background_migration/encrypt_static_object_token_spec.rb

+# frozen_string_literal: true
+require 'spec_helper'
+
+RSpec.describe Gitlab::BackgroundMigration::EncryptStaticObjectToken do
+  let(:users) { table(:users) }
+  let!(:user_without_tokens) { create_user!(name: 'notoken') }
+  let!(:user_with_plaintext_token_1) { create_user!(name: 'plaintext_1', token: 'token') }
+  let!(:user_with_plaintext_token_2) { create_user!(name: 'plaintext_2', token: 'TOKEN') }
+  let!(:user_with_plaintext_empty_token) { create_user!(name: 'plaintext_3', token: '') }
+  let!(:user_with_encrypted_token) { create_user!(name: 'encrypted', encrypted_token: 'encrypted') }
+  let!(:user_with_both_tokens) { create_user!(name: 'both', token: 'token2', encrypted_token: 'encrypted2') }
+
+  before do
+    allow(Gitlab::CryptoHelper).to receive(:aes256_gcm_encrypt).and_call_original
+    allow(Gitlab::CryptoHelper).to receive(:aes256_gcm_encrypt).with('token') { 'secure_token' }
+    allow(Gitlab::CryptoHelper).to receive(:aes256_gcm_encrypt).with('TOKEN') { 'SECURE_TOKEN' }
+  end
+
+  subject { described_class.new.perform(start_id, end_id) }
+
+  let(:start_id) { users.minimum(:id) }
+  let(:end_id) { users.maximum(:id) }
+
+  it 'backfills encrypted tokens to users with plaintext token only', :aggregate_failures do
+    subject
+
+    new_state = users.pluck(:id, :static_object_token, :static_object_token_encrypted).to_h do |row|
+      [row[0], [row[1], row[2]]]
+    end
+
+    expect(new_state.count).to eq(6)
+
+    expect(new_state[user_with_plaintext_token_1.id]).to match_array(%w[token secure_token])
+    expect(new_state[user_with_plaintext_token_2.id]).to match_array(%w[TOKEN SECURE_TOKEN])
+
+    expect(new_state[user_with_plaintext_empty_token.id]).to match_array(['', nil])
+    expect(new_state[user_without_tokens.id]).to match_array([nil, nil])
+    expect(new_state[user_with_both_tokens.id]).to match_array(%w[token2 encrypted2])
+    expect(new_state[user_with_encrypted_token.id]).to match_array([nil, 'encrypted'])
+  end
+
+  private
+
+  def create_user!(name:, token: nil, encrypted_token: nil)
+    email = "#{name}@example.com"
+
+    table(:users).create!(
+      name: name,
+      email: email,
+      username: name,
+      projects_limit: 0,
+      static_object_token: token,
+      static_object_token_encrypted: encrypted_token
+    )
+  end
+end

spec/migrations/20211210140629_encrypt_static_object_token_spec.rb

+# frozen_string_literal: true
+require 'spec_helper'
+
+require_migration!
+
+RSpec.describe EncryptStaticObjectToken, :migration do
+  let_it_be(:background_migration_jobs) { table(:background_migration_jobs) }
+  let_it_be(:users) { table(:users) }
+
+  let!(:user_without_tokens) { create_user!(name: 'notoken') }
+  let!(:user_with_plaintext_token_1) { create_user!(name: 'plaintext_1', token: 'token') }
+  let!(:user_with_plaintext_token_2) { create_user!(name: 'plaintext_2', token: 'TOKEN') }
+  let!(:user_with_encrypted_token) { create_user!(name: 'encrypted', encrypted_token: 'encrypted') }
+  let!(:user_with_both_tokens) { create_user!(name: 'both', token: 'token2', encrypted_token: 'encrypted2') }
+
+  before do
+    stub_const("#{described_class}::BATCH_SIZE", 1)
+  end
+
+  around do |example|
+    freeze_time { Sidekiq::Testing.fake! { example.run } }
+  end
+
+  it 'schedules background migrations' do
+    migrate!
+
+    expect(background_migration_jobs.count).to eq(2)
+    expect(background_migration_jobs.first.arguments).to match_array([user_with_plaintext_token_1.id, user_with_plaintext_token_1.id])
+    expect(background_migration_jobs.second.arguments).to match_array([user_with_plaintext_token_2.id, user_with_plaintext_token_2.id])
+
+    expect(BackgroundMigrationWorker.jobs.size).to eq(2)
+    expect(described_class::MIGRATION).to be_scheduled_delayed_migration(2.minutes, user_with_plaintext_token_1.id, user_with_plaintext_token_1.id)
+    expect(described_class::MIGRATION).to be_scheduled_delayed_migration(4.minutes, user_with_plaintext_token_2.id, user_with_plaintext_token_2.id)
+  end
+
+  private
+
+  def create_user!(name:, token: nil, encrypted_token: nil)
+    email = "#{name}@example.com"
+
+    table(:users).create!(
+      name: name,
+      email: email,
+      username: name,
+      projects_limit: 0,
+      static_object_token: token,
+      static_object_token_encrypted: encrypted_token
+    )
+  end
+end

spec/models/user_spec.rb

@@ -1773,6 +1773,29 @@ RSpec.describe User do
       expect(static_object_token).not_to be_blank
       expect(user.reload.static_object_token).to eq static_object_token
     end
+
+    it 'generates an encrypted version of the token' do
+      user = create(:user, static_object_token: nil)
+
+      expect(user[:static_object_token]).to be_nil
+      expect(user[:static_object_token_encrypted]).to be_nil
+
+      user.static_object_token
+
+      expect(user[:static_object_token]).to be_nil
+      expect(user[:static_object_token_encrypted]).to be_present
+    end
+
+    it 'prefers an encoded version of the token' do
+      user = create(:user, static_object_token: nil)
+
+      token = user.static_object_token
+
+      user.update_column(:static_object_token, 'Test')
+
+      expect(user.static_object_token).not_to eq('Test')
+      expect(user.static_object_token).to eq(token)
+    end
   end
 
   describe 'enabled_static_object_token' do