From 77493161b04cf257c3c39e75122feeaa9d73eb2b Mon Sep 17 00:00:00 2001
From: Wyatt Kirby <wyatt@apsis.io>
Date: Thu, 10 Oct 2024 15:09:20 -0400
Subject: [PATCH] Update to use ensure block for closing PHI access

---
 lib/phi_attrs/phi_record.rb | 42 ++++++++++++++++++-------------------
 1 file changed, 21 insertions(+), 21 deletions(-)

diff --git a/lib/phi_attrs/phi_record.rb b/lib/phi_attrs/phi_record.rb
index 43a02c4..8bcddbb 100644
--- a/lib/phi_attrs/phi_record.rb
+++ b/lib/phi_attrs/phi_record.rb
@@ -142,32 +142,32 @@ def get_phi(user_id = nil, reason = nil, allow_only: nil)
         # Save this so we don't revoke access previously extended outside the block
         frozen_instances = __instances_with_extended_phi.index_with { |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
+        begin
+          if allow_only.nil?
+            allow_phi!(user_id, reason)
+          else
+            allow_only.each { |t| t.allow_phi!(user_id, reason) }
+          end
 
-        result = yield if block_given?
+          return yield if block_given?
+        ensure
+          __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
 
-        __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?
+          if allow_only.nil?
+            disallow_last_phi!
           else
-            obj.send(:revoke_extended_phi!) # Instance is new to the set, so revoke everything
+            allow_only.each { |t| t.disallow_last_phi!(preserve_extensions: true) }
+            # We've handled any newly extended allowances ourselves above
           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
-
-        result
       end
 
       # Explicitly disallow phi access in a specific area of code. This does not