phi_attrs 0.1.4 → 0.2.0

data/Rakefile

@@ -1,6 +1,19 @@
+# frozen_string_literal: true
+
 require 'bundler/gem_tasks'
 require 'rspec/core/rake_task'
 
 RSpec::Core::RakeTask.new(:spec)
 
+# Set Chandler options
+if defined? Chandler
+  Chandler::Tasks.configure do |config|
+    config.changelog_path = 'CHANGELOG.md'
+    config.github_repository = 'apsislabs/phi_attrs'
+  end
+
+  # Add chandler as a prerequisite for `rake release`
+  task 'release:rubygem_push' => 'chandler:push'
+end
+
 task default: :spec

data/bin/console

@@ -11,4 +11,10 @@ require 'phi_attrs'
 # Pry.start
 
 require 'irb'
+require 'irb/completion'
+
+PhiAttrs.configure do |conf|
+  conf.log_path = File.join('log', 'phi_access_console.log')
+end
+
 IRB.start(__FILE__)

data/bin/helpers/docker

@@ -0,0 +1,12 @@
+#!/usr/bin/env bash
+
+# runOnDocker SERVICE
+runOnDocker() {
+  if [[ $DOCKER_CONTAINER -ne 1 ]]; then
+    echo "On Host. Executing command on container"
+    docker-compose exec -T $1 $0
+    exit $?
+  fi
+}
+
+export -f runOnDocker

data/bin/rubo_fix

@@ -0,0 +1,6 @@
+#!/usr/bin/env bash
+source bin/helpers/docker
+runOnDocker rails
+
+echo "== Starting rubocop =="
+bundle exec rubocop --format worst --format simple --format offenses --auto-correct

data/bin/run_tests

@@ -0,0 +1,19 @@
+#!/usr/bin/env bash
+source bin/helpers/docker
+runOnDocker rails
+
+echo "== Starting unit tests =="
+bundle exec appraisal rspec
+if [ $? -ne 0 ]; then
+    echo -e "\n== RSpec failed; push aborted! ==\n"
+    exit 1
+fi
+
+echo "== Starting rubocop =="
+bundle exec rubocop --format worst --format simple --format offenses
+if [ $? -ne 0 ]; then
+    echo -e "\n== Rubocop failed; push aborted! ==\n"
+    echo -e "To auto-correct errors run:"
+    echo -e "\tbin/rubo_fix"
+    exit 1
+fi

data/bin/setup

@@ -3,5 +3,7 @@ set -euo pipefail
 IFS=$'\n\t'
 set -vx
 
+gem install bundler -v $BUNDLER_VERSION
+
 bundle check || bundle install
 bundle exec appraisal install

data/docker-compose.yml

@@ -7,10 +7,12 @@ services:
       - bundle_cache:/bundle
       - .:/app
     environment:
+      - BUNDLER_VERSION=2.0.1
       - BUNDLE_JOBS=5
       - BUNDLE_PATH=/bundle
       - BUNDLE_BIN=/bundle/bin
       - GEM_HOME=/bundle
+      - DOCKER_CONTAINER=1
     command:
       - docker/start.sh
 

data/gemfiles/rails_5.1.gemfile

@@ -0,0 +1,10 @@
+# frozen_string_literal: true
+
+# This file was generated by Appraisal
+
+source 'https://rubygems.org'
+
+gem 'rails', '5.1.3'
+gem 'tzinfo-data'
+
+gemspec path: '../'

data/lib/phi_attrs.rb

@@ -1,3 +1,5 @@
+# frozen_string_literal: true
+
 require 'rails'
 require 'active_support'
 require 'request_store'
@@ -11,26 +13,35 @@ require 'phi_attrs/exceptions'
 require 'phi_attrs/phi_record'
 
 module PhiAttrs
-  def phi_model(with: nil, except: nil)
-    include PhiRecord
-    logger = ActiveSupport::Logger.new(PhiAttrs.log_path)
-    logger.formatter = Formatter.new
-    file_logger = ActiveSupport::TaggedLogging.new(logger)
+  def self.log_phi_access(user, message)
+    PhiAttrs::Logger.tagged(PHI_ACCESS_LOG_TAG, user) do
+      PhiAttrs::Logger.info(message)
+    end
+  end
 
-    PhiAttrs::Logger.logger = file_logger
+  module Model
+    def phi_model(with: nil, except: nil)
+      include PhiRecord
+    end
   end
 
-  @@log_path = nil
+  module Controller
+    extend ActiveSupport::Concern
 
-  def self.configure
-    yield self if block_given?
-  end
+    included do
+      before_action :record_i18n_data
+    end
 
-  def self.log_path
-    @@log_path
-  end
+    private
+
+    def record_i18n_data
+      RequestStore.store[:phi_attrs_controller] = self.class.name
+      RequestStore.store[:phi_attrs_action] = params[:action]
+
+      return if PhiAttrs.current_user_method.nil?
+      return unless respond_to?(PhiAttrs.current_user_method, true)
 
-  def self.log_path=(value)
-    @@log_path = value
+      RequestStore.store[:phi_attrs_current_user] = send(PhiAttrs.current_user_method)
+    end
   end
 end

data/lib/phi_attrs/configure.rb

@@ -1,17 +1,35 @@
+# frozen_string_literal: true
+
 module PhiAttrs
-  module Configure
-    @@log_path = nil
+  @@log_path = nil
+  @@current_user_method = nil
+  @@translation_prefix = 'phi'
+
+  def self.configure
+    yield self if block_given?
+  end
 
-    def configure
-      yield self if block_given?
-    end
+  def self.log_path
+    @@log_path
+  end
+
+  def self.log_path=(value)
+    @@log_path = value
+  end
 
-    def log_path
-      @@log_path
-    end
+  def self.translation_prefix
+    @@translation_prefix
+  end
+
+  def self.translation_prefix=(value)
+    @@translation_prefix = value
+  end
+
+  def self.current_user_method
+    @@current_user_method
+  end
 
-    def log_path=(value)
-      @@log_path = value
-    end
+  def self.current_user_method=(value)
+    @@current_user_method = value
   end
 end

data/lib/phi_attrs/exceptions.rb

@@ -1,8 +1,12 @@
+# frozen_string_literal: true
+
 module PhiAttrs
   module Exceptions
     class PhiAccessException < StandardError
+      TAG = 'UNAUTHORIZED ACCESS'
+
       def initialize(msg)
-        PhiAttrs::Logger.tagged('UNAUTHORIZED ACCESS') { PhiAttrs::Logger.error(msg) }
+        PhiAttrs::Logger.tagged(TAG) { PhiAttrs::Logger.error(msg) }
         super(msg)
       end
     end

data/lib/phi_attrs/formatter.rb

@@ -1,10 +1,12 @@
+# frozen_string_literal: true
+
 module PhiAttrs
-  Format = "%s %5s: %s\n".freeze
+  FORMAT = "%s %5s: %s\n"
 
   # https://github.com/ruby/ruby/blob/trunk/lib/logger.rb#L587
   class Formatter < ::Logger::Formatter
-    def call(severity, timestamp, progname, msg)
-      Format % [format_datetime(timestamp), severity, msg2str(msg)]
+    def call(severity, timestamp, _progname, msg)
+      format(FORMAT, format_datetime(timestamp), severity, msg2str(msg))
     end
   end
 end

data/lib/phi_attrs/logger.rb

@@ -1,7 +1,19 @@
+# frozen_string_literal: true
+
 module PhiAttrs
+  PHI_ACCESS_LOG_TAG = 'PHI Access Log'
+
   class Logger
     class << self
-      cattr_accessor :logger
+      def logger
+        unless @logger
+          logger = ActiveSupport::Logger.new(PhiAttrs.log_path)
+          logger.formatter = Formatter.new
+          @logger = ActiveSupport::TaggedLogging.new(logger)
+        end
+        @logger
+      end
+
       delegate :debug, :info, :warn, :error, :fatal, :tagged, to: :logger
     end
   end

data/lib/phi_attrs/phi_record.rb

@@ -1,7 +1,7 @@
+# frozen_string_literal: true
+
 # Namespace for classes and modules that handle PHI Attribute Access Logging
 module PhiAttrs
-  PHI_ACCESS_LOG_TAG = 'PHI Access Log'.freeze
-
   # Module for extending ActiveRecord models to handle PHI access logging
   # and restrict access to attributes.
   #
@@ -15,13 +15,15 @@ module PhiAttrs
       class_attribute :__phi_include_methods
       class_attribute :__phi_extend_methods
       class_attribute :__phi_methods_wrapped
-      class_attribute :__phi_methods_extended
+      class_attribute :__phi_methods_to_extend
+      class_attribute :__instances_with_extended_phi
 
       after_initialize :wrap_phi
 
       # These have to default to an empty array
       self.__phi_methods_wrapped = []
-      self.__phi_methods_extended = []
+      self.__phi_methods_to_extend = []
+      self.__instances_with_extended_phi = Set.new
     end
 
     class_methods do
@@ -70,31 +72,205 @@ module PhiAttrs
       # @example
       #   Foo.allow_phi!('user@example.com', 'viewing patient record')
       #
-      def allow_phi!(user_id, reason)
-        RequestStore.store[:phi_access] ||= {}
+      def allow_phi!(user_id = nil, reason = nil)
+        raise ArgumentError, 'block not allowed. use allow_phi with block' if block_given?
 
-        RequestStore.store[:phi_access][name] ||= {
-          phi_access_allowed: true,
-          user_id: user_id,
-          reason: reason
-        }
+        user_id ||= current_user
+        reason ||= i18n_reason
+        raise ArgumentError, 'user_id and reason cannot be blank' if user_id.blank? || reason.blank?
+
+        __phi_stack.push({
+                           phi_access_allowed: true,
+                           user_id: user_id,
+                           reason: reason
+                         })
 
         PhiAttrs::Logger.tagged(PHI_ACCESS_LOG_TAG, name) do
-          PhiAttrs::Logger.info("PHI Access Enabled for #{user_id}: #{reason}")
+          PhiAttrs::Logger.info("PHI Access Enabled for '#{user_id}': #{reason}")
+        end
+      end
+
+      # Enable PHI access for any instance of this class in the block given only.
+      #
+      # @param [String] user_id                       A unique identifier for the person accessing the PHI
+      # @param [String] reason                        The reason for accessing PHI
+      # @param [collection of PhiRecord] allow_only   Specific PhiRecords to allow access to
+      # &block [block]                                The block in which PHI access is allowed for the class
+      #
+      # @example
+      #   Foo.allow_phi('user@example.com', 'viewing patient record') do
+      #     # PHI Access Allowed
+      #   end
+      #   # PHI Access Disallowed
+      #
+      # @example
+      #   Foo.allow_phi('user@example.com', 'exporting patient list', allow_only: list_of_foos) do
+      #     # PHI Access Allowed for `list_of_foo` only
+      #   end
+      #   # PHI Access Disallowed
+      #
+      def allow_phi(user_id = nil, reason = nil, allow_only: nil)
+        raise ArgumentError, 'block required. use allow_phi! without block' unless block_given?
+
+        if allow_only.present?
+          raise ArgumentError, 'allow_only must be iterable with each' unless allow_only.respond_to?(:each)
+          raise ArgumentError, "allow_only must all be `#{name}` objects" unless allow_only.all? { |t| t.is_a?(self) }
+          raise ArgumentError, 'allow_only must all have `allow_phi!` methods' unless allow_only.all? { |t| t.respond_to?(:allow_phi!) }
+        end
+
+        # Save this so we don't revoke access previously extended outside the block
+        frozen_instances = Hash[__instances_with_extended_phi.map { |obj| [obj, obj.instance_variable_get(:@__phi_relations_extended).clone] }]
+
+        if allow_only.nil?
+          allow_phi!(user_id, reason)
+        else
+          allow_only.each { |t| t.allow_phi!(user_id, reason) }
+        end
+
+        yield if block_given?
+
+        __instances_with_extended_phi.each do |obj|
+          if frozen_instances.include?(obj)
+            old_extensions = frozen_instances[obj]
+            new_extensions = obj.instance_variable_get(:@__phi_relations_extended) - old_extensions
+            obj.send(:revoke_extended_phi!, new_extensions) if new_extensions.any?
+          else
+            obj.send(:revoke_extended_phi!) # Instance is new to the set, so revoke everything
+          end
+        end
+
+        if allow_only.nil?
+          disallow_last_phi!
+        else
+          allow_only.each { |t| t.disallow_last_phi!(preserve_extensions: true) }
+          # We've handled any newly extended allowances ourselves above
         end
       end
 
-      # Revoke PHI access for this class, if enabled by PhiRecord#allow_phi!
+      # Explicitly disallow phi access in a specific area of code. This does not
+      # play nicely with the mutating versions of `allow_phi!` and `disallow_phi!`
+      #
+      # At the moment, this doesn't work at all, as the instance won't
+      # necessarily look at the class-level stack when determining if PHI is allowed.
+      #
+      # &block [block] The block in which PHI access is explicitly disallowed.
+      #
+      # @example
+      #   # PHI Access Disallowed
+      #   Foo.disallow_phi
+      #     # PHI Access *Still* Disallowed
+      #   end
+      #   # PHI Access *Still, still* Disallowed
+      #   Foo.allow_phi!('user@example.com', 'viewing patient record')
+      #   # PHI Access Allowed
+      #   Foo.disallow_phi do
+      #     # PHI Access Disallowed
+      #   end
+      #   # PHI Access Allowed Again
+      def disallow_phi
+        raise ArgumentError, 'block required. use disallow_phi! without block' unless block_given?
+
+        __phi_stack.push({
+                           phi_access_allowed: false
+                         })
+
+        yield if block_given?
+
+        __phi_stack.pop
+      end
+
+      # Revoke all PHI access for this class, if enabled by PhiRecord#allow_phi!
       #
       # @example
       #   Foo.disallow_phi!
       #
       def disallow_phi!
-        RequestStore.store[:phi_access].delete(name) if RequestStore.store[:phi_access].present?
+        raise ArgumentError, 'block not allowed. use disallow_phi with block' if block_given?
+
+        message = __phi_stack.present? ? "PHI access disabled for #{__user_id_string(__phi_stack)}" : 'PHI access disabled. No class level access was granted.'
+
+        __reset_phi_stack
+
         PhiAttrs::Logger.tagged(PHI_ACCESS_LOG_TAG, name) do
-          PhiAttrs::Logger.info('PHI access disabled')
+          PhiAttrs::Logger.info(message)
         end
       end
+
+      # Revoke last PHI access for this class, if enabled by PhiRecord#allow_phi!
+      #
+      # @example
+      #   Foo.disallow_last_phi!
+      #
+      def disallow_last_phi!
+        raise ArgumentError, 'block not allowed' if block_given?
+
+        removed_access = __phi_stack.pop
+        message = removed_access.present? ? "PHI access disabled for #{removed_access[:user_id]}" : 'PHI access disabled. No class level access was granted.'
+
+        PhiAttrs::Logger.tagged(PHI_ACCESS_LOG_TAG, name) do
+          PhiAttrs::Logger.info(message)
+        end
+      end
+
+      # Whether PHI access is allowed for this class
+      #
+      # @example
+      #   Foo.phi_allowed?
+      #
+      # @return [Boolean] whether PHI access is allowed for this instance
+      #
+      def phi_allowed?
+        __phi_stack.present? && __phi_stack[-1][:phi_access_allowed]
+      end
+
+      def __phi_stack
+        RequestStore.store[:phi_access] ||= {}
+        RequestStore.store[:phi_access][name] ||= []
+      end
+
+      def __reset_phi_stack
+        RequestStore.store[:phi_access] ||= {}
+        RequestStore.store[:phi_access][name] = []
+      end
+
+      def __user_id_string(access_list)
+        access_list ||= []
+        access_list.map { |c| "'#{c[:user_id]}'" }.join(',')
+      end
+
+      def current_user
+        RequestStore.store[:phi_attrs_current_user]
+      end
+
+      def i18n_reason
+        controller = RequestStore.store[:phi_attrs_controller]
+        action = RequestStore.store[:phi_attrs_action]
+
+        return nil if controller.blank? || action.blank?
+
+        i18n_path = [PhiAttrs.translation_prefix] + __path_to_controller_and_action(controller, action)
+        i18n_path.push(*__path_to_class)
+        i18n_key = i18n_path.join('.')
+
+        return I18n.t(i18n_key) if I18n.exists?(i18n_key)
+
+        locale = I18n.locale || I18n.default_locale
+
+        PhiAttrs::Logger.warn "No #{locale} PHI Reason found for #{i18n_key}"
+      end
+
+      def __path_to_controller_and_action(controller, action)
+        module_paths = controller.underscore.split('/')
+        class_name_parts = module_paths.pop.split('_')
+        class_name_parts.pop if class_name_parts[-1] == 'controller'
+        module_paths.push(class_name_parts.join('_'), action)
+      end
+
+      def __path_to_class
+        module_paths = name.underscore.split('/')
+        class_name_parts = module_paths.pop.split('_')
+        module_paths.push(class_name_parts.join('_'))
+      end
     end
 
     # Get all method names to be wrapped with PHI access logging
@@ -125,29 +301,112 @@ module PhiAttrs
     #   foo = Foo.find(1)
     #   foo.allow_phi!('user@example.com', 'viewing patient record')
     #
-    def allow_phi!(user_id, reason)
+    def allow_phi!(user_id = nil, reason = nil)
+      raise ArgumentError, 'block not allowed. use allow_phi with block' if block_given?
+
+      user_id ||= self.class.current_user
+      reason ||= self.class.i18n_reason
+      raise ArgumentError, 'user_id and reason cannot be blank' if user_id.blank? || reason.blank?
+
       PhiAttrs::Logger.tagged(*phi_log_keys) do
-        @__phi_access_allowed = true
-        @__phi_user_id = user_id
-        @__phi_access_reason = reason
+        @__phi_access_stack.push({
+                                   phi_access_allowed: true,
+                                   user_id: user_id,
+                                   reason: reason
+                                 })
 
         PhiAttrs::Logger.info("PHI Access Enabled for '#{user_id}': #{reason}")
       end
     end
 
-    # Revoke PHI access for a single instance of this class
+    # Enable PHI access for a single instance of this class inside the block.
+    # Nested calls to allow_phi will log once per nested call
+    #
+    # @param [String] user_id   A unique identifier for the person accessing the PHI
+    # @param [String] reason    The reason for accessing PHI
+    # @yield                    The block in which phi access is allowed
+    #
+    # @example
+    #   foo = Foo.find(1)
+    #   foo.allow_phi('user@example.com', 'viewing patient record') do
+    #    # PHI Access Allowed Here
+    #   end
+    #   # PHI Access Disallowed Here
+    #
+    def allow_phi(user_id = nil, reason = nil)
+      raise ArgumentError, 'block required. use allow_phi! without block' unless block_given?
+
+      extended_instances = @__phi_relations_extended.clone
+      allow_phi!(user_id, reason)
+
+      yield if block_given?
+
+      new_extensions = @__phi_relations_extended - extended_instances
+      disallow_last_phi!(preserve_extensions: true)
+      revoke_extended_phi!(new_extensions) if new_extensions.any?
+    end
+
+    # Revoke all PHI access for a single instance of this class.
     #
     # @example
     #   foo = Foo.find(1)
     #   foo.disallow_phi!
     #
     def disallow_phi!
+      raise ArgumentError, 'block not allowed. use disallow_phi with block' if block_given?
+
+      PhiAttrs::Logger.tagged(*phi_log_keys) do
+        removed_access_for = self.class.__user_id_string(@__phi_access_stack)
+
+        revoke_extended_phi!
+        @__phi_access_stack = []
+
+        message = removed_access_for.present? ? "PHI access disabled for #{removed_access_for}" : 'PHI access disabled. No instance level access was granted.'
+        PhiAttrs::Logger.info(message)
+      end
+    end
+
+    # Dissables PHI access for a single instance of this class inside the block.
+    # Nested calls to allow_phi will log once per nested call
+    #
+    # @param [String] user_id   A unique identifier for the person accessing the PHI
+    # @param [String] reason    The reason for accessing PHI
+    # @yield                    The block in which phi access is allowed
+    #
+    # @example
+    #   foo = Foo.find(1)
+    #   foo.allow_phi('user@example.com', 'viewing patient record') do
+    #    # PHI Access Allowed Here
+    #   end
+    #   # PHI Access Disallowed Here
+    #
+    def disallow_phi
+      raise ArgumentError, 'block required. use disallow_phi! without block' unless block_given?
+
+      add_disallow_flag!
+      add_disallow_flag_to_extended_phi!
+
+      yield if block_given?
+
+      remove_disallow_flag_from_extended_phi!
+      remove_disallow_flag!
+    end
+
+    # Revoke last PHI access for a single instance of this class.
+    #
+    # @example
+    #   foo = Foo.find(1)
+    #   foo.disallow_last_phi!
+    #
+    def disallow_last_phi!(preserve_extensions: false)
+      raise ArgumentError, 'block not allowed' if block_given?
+
       PhiAttrs::Logger.tagged(*phi_log_keys) do
-        @__phi_access_allowed = false
-        @__phi_user_id = nil
-        @__phi_access_reason = nil
+        removed_access = @__phi_access_stack.pop
 
-        PhiAttrs::Logger.info('PHI access disabled')
+        revoke_extended_phi! unless preserve_extensions
+        message = removed_access.present? ? "PHI access disabled for #{removed_access[:user_id]}" : 'PHI access disabled. No instance level access was granted.'
+        PhiAttrs::Logger.info(message)
       end
     end
 
@@ -160,7 +419,23 @@ module PhiAttrs
     # @return [Boolean] whether PHI access is allowed for this instance
     #
     def phi_allowed?
-      @__phi_access_allowed || RequestStore.store.dig(:phi_access, self.class.name, :phi_access_allowed)
+      !phi_context.nil? && phi_context[:phi_access_allowed]
+    end
+
+    protected
+
+    # Adds a disallow phi flag to instance internal stack.
+    # @private since subject to change
+    def add_disallow_flag!
+      @__phi_access_stack.push({
+                                 phi_access_allowed: false
+                               })
+    end
+
+    # removes the last item in instance internal stack.
+    # @private since subject to change
+    def remove_disallow_flag!
+      @__phi_access_stack.pop
     end
 
     private
@@ -172,12 +447,13 @@ module PhiAttrs
     #
     def wrap_phi
       # Disable PHI access by default
-      @__phi_access_allowed = false
-      @__phi_access_logged = false
+      @__phi_access_stack = []
+      @__phi_methods_extended = Set.new
+      @__phi_relations_extended = Set.new
 
       # Wrap attributes with PHI Logger and Access Control
       __phi_wrapped_methods.each { |m| phi_wrap_method(m) }
-      __phi_extended_methods.each { |m| phi_extend_access(m) }
+      __phi_extended_methods.each { |m| phi_wrap_extension(m) }
     end
 
     # Log Key for an instance of this class. If the instance is persisted in the
@@ -204,7 +480,7 @@ module PhiAttrs
     # @return [String] the user_id passed in to allow_phi!
     #
     def phi_allowed_by
-      @__phi_user_id || RequestStore.store.dig(:phi_access, self.class.name, :user_id)
+      phi_context[:user_id]
     end
 
     # The access reason for allowing access to this instance.
@@ -215,7 +491,41 @@ module PhiAttrs
     # @return [String] the reason passed in to allow_phi!
     #
     def phi_access_reason
-      @__phi_access_reason || RequestStore.store.dig(:phi_access, self.class.name, :reason)
+      phi_context[:reason]
+    end
+
+    def phi_context
+      instance_phi_context || class_phi_context
+    end
+
+    def instance_phi_context
+      @__phi_access_stack && @__phi_access_stack[-1]
+    end
+
+    def class_phi_context
+      self.class.__phi_stack[-1]
+    end
+
+    # The unique identifiers for everything with access allowed on this instance.
+    #
+    # @private
+    #
+    # @return String of all the user_id's passed in to allow_phi!
+    #
+    def all_phi_allowed_by
+      self.class.__user_id_string(all_phi_context)
+    end
+
+    def all_phi_context
+      (@__phi_access_stack || []) + (self.class.__phi_stack || [])
+    end
+
+    def all_phi_context_logged?
+      all_phi_context.all? { |v| v[:logged] }
+    end
+
+    def set_all_phi_context_logged
+      all_phi_context.each { |c| c[:logged] = true }
     end
 
     # Core logic for wrapping methods in PHI access logging and access restriction.
@@ -260,9 +570,9 @@ module PhiAttrs
         PhiAttrs::Logger.tagged(*phi_log_keys) do
           raise PhiAttrs::Exceptions::PhiAccessException, "Attempted PHI access for #{self.class.name} #{@__phi_user_id}" unless phi_allowed?
 
-          unless @__phi_access_logged
-            PhiAttrs::Logger.info("'#{phi_allowed_by}' accessing #{self.class.name}. Triggered by method: #{method_name}")
-            @__phi_access_logged = true
+          unless all_phi_context_logged?
+            PhiAttrs::Logger.info("#{self.class.name} access by [#{all_phi_allowed_by}]. Triggered by method: #{method_name}")
+            set_all_phi_context_logged
           end
 
           send(unwrapped_method, *args, &block)
@@ -283,18 +593,24 @@ module PhiAttrs
     #
     # @private
     #
-    def phi_extend_access(method_name)
-      return if self.class.__phi_methods_extended.include? method_name
+    def phi_wrap_extension(method_name)
+      return if self.class.__phi_methods_to_extend.include? method_name
 
-      wrapped_method = :"__#{method_name}_phi_access_extended"
-      unwrapped_method = :"__#{method_name}_phi_access_original"
+      wrapped_method = wrapped_extended_name(method_name)
+      unwrapped_method = unwrapped_extended_name(method_name)
 
       self.class.send(:define_method, wrapped_method) do |*args, &block|
-        # get the unwrapped relation
         relation = send(unwrapped_method, *args, &block)
 
-        if phi_allowed? && relation.class.included_modules.include?(PhiRecord)
-          relation.allow_phi!(phi_allowed_by, phi_access_reason) unless relation.phi_allowed?
+        if phi_allowed?
+          if relation.present? && relation_klass(relation).included_modules.include?(PhiRecord)
+            relations = relation.is_a?(Enumerable) ? relation : [relation]
+            relations.each do |r|
+              r.allow_phi!(phi_allowed_by, phi_access_reason) unless @__phi_relations_extended.include?(r)
+            end
+            @__phi_relations_extended.merge(relations)
+            self.class.__instances_with_extended_phi.add(self)
+          end
         end
 
         relation
@@ -304,7 +620,47 @@ module PhiAttrs
       self.class.send(:alias_method, unwrapped_method, method_name)
       self.class.send(:alias_method, method_name, wrapped_method)
 
-      self.class.__phi_methods_extended << method_name
+      self.class.__phi_methods_to_extend << method_name
+    end
+
+    # Revoke PHI access for all `extend`ed relations (or only those given)
+    def revoke_extended_phi!(relations = nil)
+      relations ||= @__phi_relations_extended
+      relations.each do |relation|
+        relation.disallow_last_phi! if relation.present? && relation_klass(relation).included_modules.include?(PhiRecord)
+      end
+      @__phi_relations_extended.subtract(relations)
+    end
+
+    # Adds a disallow PHI access to the stack for block syntax for all `extend`ed relations (or only those given)
+    def add_disallow_flag_to_extended_phi!(relations = nil)
+      relations ||= @__phi_relations_extended
+      relations.each do |relation|
+        relation.add_disallow_flag! if relation.present? && relation_klass(relation).included_modules.include?(PhiRecord)
+      end
+    end
+
+    # Adds a disallow PHI access to the stack for all for all `extend`ed relations (or only those given)
+    def remove_disallow_flag_from_extended_phi!(relations = nil)
+      relations ||= @__phi_relations_extended
+      relations.each do |relation|
+        relation.remove_disallow_flag! if relation.present? && relation_klass(relation).included_modules.include?(PhiRecord)
+      end
+    end
+
+    def relation_klass(rel)
+      return rel.klass if rel.is_a?(ActiveRecord::Relation)
+      return rel.first.class if rel.is_a?(Enumerable)
+
+      return rel.class
+    end
+
+    def wrapped_extended_name(method_name)
+      :"__#{method_name}_phi_access_extended"
+    end
+
+    def unwrapped_extended_name(method_name)
+      :"__#{method_name}_phi_access_original"
     end
   end
 end