1 files changed, 61 insertions, 6 deletions

diff --git a/app/models/ability.rb b/app/models/ability.rb
index c18bd21d754..6a63a8d46ba 100644
--- a/app/models/ability.rb
+++ b/app/models/ability.rb
@@ -54,7 +54,7 @@ class Ability
       end
     end
 
-    def allowed?(user, action, subject = :global, opts = {})
+    def allowed?(user, ability, subject = :global, opts = {})
       if subject.is_a?(Hash)
         opts = subject
         subject = :global
@@ -64,21 +64,76 @@ class Ability
 
       case opts[:scope]
       when :user
-        DeclarativePolicy.user_scope { policy.can?(action) }
+        DeclarativePolicy.user_scope { policy.allowed?(ability) }
       when :subject
-        DeclarativePolicy.subject_scope { policy.can?(action) }
+        DeclarativePolicy.subject_scope { policy.allowed?(ability) }
       else
-        policy.can?(action)
+        policy.allowed?(ability)
       end
+    ensure
+      # TODO: replace with runner invalidation:
+      # See: https://gitlab.com/gitlab-org/declarative-policy/-/merge_requests/24
+      # See: https://gitlab.com/gitlab-org/declarative-policy/-/merge_requests/25
+      forget_runner_result(policy.runner(ability)) if policy && ability_forgetting?
     end
 
     def policy_for(user, subject = :global)
-      cache = Gitlab::SafeRequestStore.active? ? Gitlab::SafeRequestStore : {}
-      DeclarativePolicy.policy_for(user, subject, cache: cache)
+      DeclarativePolicy.policy_for(user, subject, cache: ::Gitlab::SafeRequestStore.storage)
+    end
+
+    # This method is something of a band-aid over the problem. The problem is
+    # that some conditions may not be re-entrant, if facts change.
+    # (`BasePolicy#admin?` is a known offender, due to the effects of
+    # `admin_mode`)
+    #
+    # To deal with this we need to clear two elements of state: the offending
+    # conditions (selected by 'pattern') and the cached ability checks (cached
+    # on the `policy#runner(ability)`).
+    #
+    # Clearing the conditions (see `forget_all_but`) is fairly robust, provided
+    # the pattern is not _under_-selective. Clearing the runners is harder,
+    # since there is not good way to know which abilities any given condition
+    # may affect. The approach taken here (see `forget_runner_result`) is to
+    # discard all runner results generated during a `forgetting` block. This may
+    # be _under_-selective if a runner prior to this block cached a state value
+    # that might now be invalid.
+    #
+    # TODO: add some kind of reverse-dependency mapping in DeclarativePolicy
+    # See: https://gitlab.com/gitlab-org/declarative-policy/-/issues/14
+    def forgetting(pattern, &block)
+      was_forgetting = ability_forgetting?
+      ::Gitlab::SafeRequestStore[:ability_forgetting] = true
+      keys_before = ::Gitlab::SafeRequestStore.storage.keys
+
+      yield
+    ensure
+      ::Gitlab::SafeRequestStore[:ability_forgetting] = was_forgetting
+      forget_all_but(keys_before, matching: pattern)
     end
 
     private
 
+    def ability_forgetting?
+      ::Gitlab::SafeRequestStore[:ability_forgetting]
+    end
+
+    def forget_all_but(keys_before, matching:)
+      keys_after = ::Gitlab::SafeRequestStore.storage.keys
+
+      added_keys = keys_after - keys_before
+      added_keys.each do |key|
+        if key.is_a?(String) && key.start_with?('/dp') && key =~ matching
+          ::Gitlab::SafeRequestStore.delete(key)
+        end
+      end
+    end
+
+    def forget_runner_result(runner)
+      # TODO: add support in DP for this
+      # See: https://gitlab.com/gitlab-org/declarative-policy/-/issues/15
+      runner.instance_variable_set(:@state, nil)
+    end
+
     def apply_filters_if_needed(elements, user, filters)
       filters.each do |ability, filter|
         elements = filter.call(elements) unless allowed?(user, ability)