Refactor token authenticatable encrypted strategy

2 files changed, 20 insertions, 13 deletions

diff --git a/app/models/concerns/token_authenticatable_strategies/base.rb b/app/models/concerns/token_authenticatable_strategies/base.rb
index 9aa04ba3d62..ef5ed0e577e 100644
--- a/app/models/concerns/token_authenticatable_strategies/base.rb
+++ b/app/models/concerns/token_authenticatable_strategies/base.rb
@@ -1,6 +1,8 @@
 # frozen_string_literal: true
 
 module TokenAuthenticatableStrategies
+  attr_reader :klass, :token_field, :options
+
   class Base
     def initialize(klass, token_field, options)
       @klass = klass
@@ -36,6 +38,10 @@ module TokenAuthenticatableStrategies
       instance.save! if Gitlab::Database.read_write?
     end
 
+    def fallback?
+      options[:fallback] == true
+    end
+
     protected
 
     def write_new_token(instance)
diff --git a/app/models/concerns/token_authenticatable_strategies/encrypted.rb b/app/models/concerns/token_authenticatable_strategies/encrypted.rb
index c68ac399594..822f0b1935c 100644
--- a/app/models/concerns/token_authenticatable_strategies/encrypted.rb
+++ b/app/models/concerns/token_authenticatable_strategies/encrypted.rb
@@ -7,45 +7,46 @@ module TokenAuthenticatableStrategies
     def find_token_authenticatable(token, unscoped = false)
       return unless token
 
+      encrypted_value = Gitlab::CryptoHelper.aes256_gcm_encrypt(token)
       token_authenticatable = relation(unscoped)
-        .find_by(token_field_name => Gitlab::CryptoHelper.aes256_gcm_encrypt(token))
+        .find_by(encrypted_field => encrypted_value)
 
-      if @options[:fallback]
-        token_authenticatable ||= fallback_strategy.find_token_authenticatable(token)
+      if fallback?
+        token_authenticatable ||= fallback_strategy
+          .find_token_authenticatable(token)
       end
 
       token_authenticatable
     end
 
     def get_token(instance)
-      raw_token = instance.read_attribute(token_field_name)
+      raw_token = instance.read_attribute(encrypted_field)
       token = Gitlab::CryptoHelper.aes256_gcm_decrypt(raw_token)
-      token ||= fallback_strategy.get_token(instance) if @options[:fallback]
+      token ||= fallback_strategy.get_token(instance) if fallback?
     end
 
     def set_token(instance, token)
-      raise ArgumentError unless token
+      raise ArgumentError unless token.present?
 
-      instance[token_field_name] = Gitlab::CryptoHelper.aes256_gcm_encrypt(token)
-      # instance[@token_field] = nil if @options[:fallback] # TODO this seems wrong
+      instance[encrypted_field] = Gitlab::CryptoHelper.aes256_gcm_encrypt(token)
     end
 
     protected
 
     def fallback_strategy
       @fallback_strategy ||= TokenAuthenticatableStrategies::Insecure
-        .new(@klass, @token_field, @options)
+        .new(klass, token_field, options)
     end
 
     def token_set?(instance)
-      raw_token = instance.read_attribute(token_field_name)
-      raw_token ||= instance.read_attribute(@token_field) if @options[:fallback]
+      raw_token = instance.read_attribute(encrypted_field)
+      raw_token ||= instance.read_attribute(token_field) if fallback?
 
       raw_token.present?
     end
 
-    def token_field_name
-      "#{@token_field}_encrypted"
+    def encrypted_field
+      @encrypted_field ||= "#{@token_field}_encrypted"
     end
   end
 end