From 37add27a00d38e4edaaec945ed9f44a123523884 Mon Sep 17 00:00:00 2001
From: Grzegorz Bizon <grzesiek.bizon@gmail.com>
Date: Tue, 27 Nov 2018 14:34:05 +0100
Subject: [PATCH] Improve token authenticable tests and exceptions

---
 .../concerns/token_authenticatable_strategies/base.rb     | 8 ++++----
 .../token_authenticatable_strategies/encrypted.rb         | 2 +-
 lib/gitlab/utils.rb                                       | 3 ++-
 spec/lib/gitlab/crypto_helper_spec.rb                     | 5 +++--
 4 files changed, 10 insertions(+), 8 deletions(-)

diff --git a/app/models/concerns/token_authenticatable_strategies/base.rb b/app/models/concerns/token_authenticatable_strategies/base.rb
index 23ee34962c7..4c63c0dd629 100644
--- a/app/models/concerns/token_authenticatable_strategies/base.rb
+++ b/app/models/concerns/token_authenticatable_strategies/base.rb
@@ -47,17 +47,17 @@ module TokenAuthenticatableStrategies
       options[:fallback] == true
     end
 
-    def self.fabricate(instance, field, options)
+    def self.fabricate(model, field, options)
       if options[:digest] && options[:encrypted]
         raise ArgumentError, 'Incompatible options set!'
       end
 
       if options[:digest]
-        TokenAuthenticatableStrategies::Digest.new(instance, field, options)
+        TokenAuthenticatableStrategies::Digest.new(model, field, options)
       elsif options[:encrypted]
-        TokenAuthenticatableStrategies::Encrypted.new(instance, field, options)
+        TokenAuthenticatableStrategies::Encrypted.new(model, field, options)
       else
-        TokenAuthenticatableStrategies::Insecure.new(instance, field, options)
+        TokenAuthenticatableStrategies::Insecure.new(model, field, options)
       end
     end
 
diff --git a/app/models/concerns/token_authenticatable_strategies/encrypted.rb b/app/models/concerns/token_authenticatable_strategies/encrypted.rb
index 8e052a3ef68..c23d78b050a 100644
--- a/app/models/concerns/token_authenticatable_strategies/encrypted.rb
+++ b/app/models/concerns/token_authenticatable_strategies/encrypted.rb
@@ -46,7 +46,7 @@ module TokenAuthenticatableStrategies
       raise ArgumentError unless token.present?
 
       instance[encrypted_field] = Gitlab::CryptoHelper.aes256_gcm_encrypt(token)
-      instance[token_field] = nil
+      fallback_strategy.set_token(instance, nil) if fallback?
       token
     end
 
diff --git a/lib/gitlab/utils.rb b/lib/gitlab/utils.rb
index 41ed1dc5605..96d2ed88b83 100644
--- a/lib/gitlab/utils.rb
+++ b/lib/gitlab/utils.rb
@@ -17,7 +17,8 @@ module Gitlab
     end
 
     def ensure_utf8_size(str, bytes:)
-      raise ArgumentError if str.empty? || bytes.negative?
+      raise ArgumentError, 'Empty string provided!' if str.empty?
+      raise ArgumentError, 'Negative string size provided!' if bytes.negative?
 
       truncated = str.each_char.each_with_object(+'') do |char, object|
         if object.bytesize + char.bytesize > bytes
diff --git a/spec/lib/gitlab/crypto_helper_spec.rb b/spec/lib/gitlab/crypto_helper_spec.rb
index bd3a38bfbaa..05cc6cf15de 100644
--- a/spec/lib/gitlab/crypto_helper_spec.rb
+++ b/spec/lib/gitlab/crypto_helper_spec.rb
@@ -5,7 +5,7 @@ describe Gitlab::CryptoHelper do
     it 'generates SHA256 digest Base46 encoded' do
       digest = described_class.sha256('some-value')
 
-      expect(digest).to match %r{^[A-Za-z0-9+/=]+$}
+      expect(digest).to match %r{\A[A-Za-z0-9+/=]+\z}
       expect(digest).to eq digest.strip
     end
   end
@@ -14,7 +14,8 @@ describe Gitlab::CryptoHelper do
     it 'is Base64 encoded string without new line character' do
       encrypted = described_class.aes256_gcm_encrypt('some-value')
 
-      expect(encrypted).to match %r{^[A-Za-z0-9+/=]+$}
+      expect(encrypted).to match %r{\A[A-Za-z0-9+/=]+\z}
+      expect(encrypted).not_to include "\n"
     end
   end