From 364791c9f02c5e1d44776fa8e48453dd447d4630 Mon Sep 17 00:00:00 2001
From: Lin Jen-Shin <godfat@godfat.org>
Date: Fri, 8 Mar 2019 23:28:59 +0800
Subject: [PATCH] Update TokenAuthenticatable so methods can be overridden

---
 app/models/application_setting.rb            |  8 ++++++--
 app/models/concerns/token_authenticatable.rb | 19 +++++++++++++------
 2 files changed, 19 insertions(+), 8 deletions(-)

diff --git a/app/models/application_setting.rb b/app/models/application_setting.rb
index fcd3cb4140a..56cc7c3784f 100644
--- a/app/models/application_setting.rb
+++ b/app/models/application_setting.rb
@@ -7,11 +7,15 @@ class ApplicationSetting < ActiveRecord::Base
   include IgnorableColumn
   include ChronicDurationAttribute
 
-  include ApplicationSettingImplementation
-
   add_authentication_token_field :runners_registration_token, encrypted: -> { Feature.enabled?(:application_settings_tokens_optional_encryption) ? :optional : :required }
   add_authentication_token_field :health_check_access_token
 
+  # Include here so it can override methods from
+  # `add_authentication_token_field`
+  # We don't prepend for now because otherwise we'll need to
+  # fix a lot of tests using allow_any_instance_of
+  include ApplicationSettingImplementation
+
   serialize :restricted_visibility_levels # rubocop:disable Cop/ActiveRecordSerialize
   serialize :import_sources # rubocop:disable Cop/ActiveRecordSerialize
   serialize :disabled_oauth_sign_in_sources, Array # rubocop:disable Cop/ActiveRecordSerialize
diff --git a/app/models/concerns/token_authenticatable.rb b/app/models/concerns/token_authenticatable.rb
index f5bb559ceda..8c769be0489 100644
--- a/app/models/concerns/token_authenticatable.rb
+++ b/app/models/concerns/token_authenticatable.rb
@@ -26,34 +26,41 @@ module TokenAuthenticatable
         end
       end
 
-      define_method(token_field) do
+      mod = token_authenticatable_module
+
+      mod.define_method(token_field) do
         strategy.get_token(self)
       end
 
-      define_method("set_#{token_field}") do |token|
+      mod.define_method("set_#{token_field}") do |token|
         strategy.set_token(self, token)
       end
 
-      define_method("ensure_#{token_field}") do
+      mod.define_method("ensure_#{token_field}") do
         strategy.ensure_token(self)
       end
 
       # Returns a token, but only saves when the database is in read & write mode
-      define_method("ensure_#{token_field}!") do
+      mod.define_method("ensure_#{token_field}!") do
         strategy.ensure_token!(self)
       end
 
       # Resets the token, but only saves when the database is in read & write mode
-      define_method("reset_#{token_field}!") do
+      mod.define_method("reset_#{token_field}!") do
         strategy.reset_token!(self)
       end
 
-      define_method("#{token_field}_matches?") do |other_token|
+      mod.define_method("#{token_field}_matches?") do |other_token|
         token = read_attribute(token_field)
         token.present? && ActiveSupport::SecurityUtils.variable_size_secure_compare(other_token, token)
       end
     end
 
+    def token_authenticatable_module
+      @token_authenticatable_module ||=
+        const_set(:TokenAuthenticatable, Module.new).tap(&method(:include))
+    end
+
     def token_authenticatable_fields
       @token_authenticatable_fields ||= []
     end