2 files changed, 136 insertions, 4 deletions

diff --git a/app/models/concerns/token_authenticatable_strategies/base.rb b/app/models/concerns/token_authenticatable_strategies/base.rb
index 413721d3e6c..01fb194281a 100644
--- a/app/models/concerns/token_authenticatable_strategies/base.rb
+++ b/app/models/concerns/token_authenticatable_strategies/base.rb
@@ -2,6 +2,8 @@
 
 module TokenAuthenticatableStrategies
   class Base
+    attr_reader :klass, :token_field, :options
+
     def initialize(klass, token_field, options)
       @klass = klass
       @token_field = token_field
@@ -22,6 +24,7 @@ module TokenAuthenticatableStrategies
 
     def ensure_token(instance)
       write_new_token(instance) unless token_set?(instance)
+      get_token(instance)
     end
 
     # Returns a token, but only saves when the database is in read & write mode
@@ -36,6 +39,36 @@ module TokenAuthenticatableStrategies
       instance.save! if Gitlab::Database.read_write?
     end
 
+    def fallback?
+      unless options[:fallback].in?([true, false, nil])
+        raise ArgumentError, 'fallback: needs to be a boolean value!'
+      end
+
+      options[:fallback] == true
+    end
+
+    def migrating?
+      unless options[:migrating].in?([true, false, nil])
+        raise ArgumentError, 'migrating: needs to be a boolean value!'
+      end
+
+      options[:migrating] == true
+    end
+
+    def self.fabricate(model, field, options)
+      if options[:digest] && options[:encrypted]
+        raise ArgumentError, 'Incompatible options set!'
+      end
+
+      if options[:digest]
+        TokenAuthenticatableStrategies::Digest.new(model, field, options)
+      elsif options[:encrypted]
+        TokenAuthenticatableStrategies::Encrypted.new(model, field, options)
+      else
+        TokenAuthenticatableStrategies::Insecure.new(model, field, options)
+      end
+    end
+
     protected
 
     def write_new_token(instance)
@@ -65,9 +98,5 @@ module TokenAuthenticatableStrategies
     def token_set?(instance)
       raise NotImplementedError
     end
-
-    def token_field_name
-      @token_field
-    end
   end
 end
diff --git a/app/models/concerns/token_authenticatable_strategies/encrypted.rb b/app/models/concerns/token_authenticatable_strategies/encrypted.rb
new file mode 100644
index 00000000000..152491aa6e9
--- /dev/null
+++ b/app/models/concerns/token_authenticatable_strategies/encrypted.rb
@@ -0,0 +1,103 @@
+# frozen_string_literal: true
+
+module TokenAuthenticatableStrategies
+  class Encrypted < Base
+    def initialize(*)
+      super
+
+      if migrating? && fallback?
+        raise ArgumentError, '`fallback` and `migrating` options are not compatible!'
+      end
+    end
+
+    def find_token_authenticatable(token, unscoped = false)
+      return if token.blank?
+
+      if fully_encrypted?
+        return find_by_encrypted_token(token, unscoped)
+      end
+
+      if fallback?
+        find_by_encrypted_token(token, unscoped) ||
+          find_by_plaintext_token(token, unscoped)
+      elsif migrating?
+        find_by_plaintext_token(token, unscoped)
+      else
+        raise ArgumentError, 'Unknown encryption phase!'
+      end
+    end
+
+    def ensure_token(instance)
+      # TODO, tech debt, because some specs are testing migrations, but are still
+      # using factory bot to create resources, it might happen that a database
+      # schema does not have "#{token_name}_encrypted" field yet, however a bunch
+      # of models call `ensure_#{token_name}` in `before_save`.
+      #
+      # In that case we are using insecure strategy, but this should only happen
+      # in tests, because otherwise `encrypted_field` is going to exist.
+      #
+      # Another use case is when we are caching resources / columns, like we do
+      # in case of ApplicationSetting.
+
+      return super if instance.has_attribute?(encrypted_field)
+
+      if fully_encrypted?
+        raise ArgumentError, 'Using encrypted strategy when encrypted field is missing!'
+      else
+        insecure_strategy.ensure_token(instance)
+      end
+    end
+
+    def get_token(instance)
+      return insecure_strategy.get_token(instance) if migrating?
+
+      encrypted_token = instance.read_attribute(encrypted_field)
+      token = Gitlab::CryptoHelper.aes256_gcm_decrypt(encrypted_token)
+
+      token || (insecure_strategy.get_token(instance) if fallback?)
+    end
+
+    def set_token(instance, token)
+      raise ArgumentError unless token.present?
+
+      instance[encrypted_field] = Gitlab::CryptoHelper.aes256_gcm_encrypt(token)
+      instance[token_field] = token if migrating?
+      instance[token_field] = nil if fallback?
+      token
+    end
+
+    def fully_encrypted?
+      !migrating? && !fallback?
+    end
+
+    protected
+
+    def find_by_plaintext_token(token, unscoped)
+      insecure_strategy.find_token_authenticatable(token, unscoped)
+    end
+
+    def find_by_encrypted_token(token, unscoped)
+      encrypted_value = Gitlab::CryptoHelper.aes256_gcm_encrypt(token)
+      relation(unscoped).find_by(encrypted_field => encrypted_value)
+    end
+
+    def insecure_strategy
+      @insecure_strategy ||= TokenAuthenticatableStrategies::Insecure
+        .new(klass, token_field, options)
+    end
+
+    def token_set?(instance)
+      raw_token = instance.read_attribute(encrypted_field)
+
+      unless fully_encrypted?
+        raw_token ||= insecure_strategy.get_token(instance)
+      end
+
+      raw_token.present?
+    end
+
+    def encrypted_field
+      @encrypted_field ||= "#{@token_field}_encrypted"
+    end
+  end
+end