From 89faa9ca045f332c9f06a1f159bb95f9949ebf06 Mon Sep 17 00:00:00 2001
From: Andrew Kane <andrew@ankane.org>
Date: Wed, 23 Oct 2024 16:50:50 -0700
Subject: [PATCH] Improved attributes, attribute_names, and has_attribute? when
 ciphertext attribute not loaded

---
 CHANGELOG.md         |  1 +
 lib/lockbox/model.rb | 35 ++++++++++++++++++++++++++++++++++-
 test/model_test.rb   | 11 ++++-------
 3 files changed, 39 insertions(+), 8 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index d568747..8b1ecf3 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,5 +1,6 @@
 ## 2.0.0 (unreleased)
 
+- Improved `attributes`, `attribute_names`, and `has_attribute?` when ciphertext attribute not loaded
 - Dropped support for Active Record < 7 and Ruby < 3.1
 - Dropped support for Mongoid < 8
 
diff --git a/lib/lockbox/model.rb b/lib/lockbox/model.rb
index 9ea0e71..ee7b766 100644
--- a/lib/lockbox/model.rb
+++ b/lib/lockbox/model.rb
@@ -148,7 +148,40 @@ def attributes
                     end
                   end
                 end
-                super
+
+                # remove attributes that do not have a ciphertext attribute
+                attributes = super
+                self.class.lockbox_attributes.each do |k, lockbox_attribute|
+                  if !attributes.include?(lockbox_attribute[:encrypted_attribute].to_s)
+                    attributes.delete(k.to_s)
+                    attributes.delete(lockbox_attribute[:attribute])
+                  end
+                end
+                attributes
+              end
+
+              # remove attribute names that do not have a ciphertext attribute
+              def attribute_names
+                # hash preserves key order
+                names_set = super.to_h { |v| [v, true] }
+                self.class.lockbox_attributes.each do |k, lockbox_attribute|
+                  if !names_set.include?(lockbox_attribute[:encrypted_attribute].to_s)
+                    names_set.delete(k.to_s)
+                    names_set.delete(lockbox_attribute[:attribute])
+                  end
+                end
+                names_set.keys
+              end
+
+              # check the ciphertext attribute for encrypted attributes
+              def has_attribute?(attr_name)
+                attr_name = attr_name.to_s
+                _, lockbox_attribute = self.class.lockbox_attributes.find { |_, la| la[:attribute] == attr_name }
+                if lockbox_attribute
+                  super(lockbox_attribute[:encrypted_attribute])
+                else
+                  super
+                end
               end
 
               # needed for in-place modifications
diff --git a/test/model_test.rb b/test/model_test.rb
index 7fcfe37..f599c8e 100644
--- a/test/model_test.rb
+++ b/test/model_test.rb
@@ -225,13 +225,10 @@ def test_attributes_not_loaded
     assert !user.has_attribute?("name")
     assert !user.has_attribute?(:name)
 
-    # TODO try to make virtual attribute behavior consistent
-    # this may be difficult, as virtual attributes are set to self.class._default_attributes
-    # which gets merged with query attributes in initialize method of active_record/core.rb
-    # assert_equal ["id"], user.attributes.keys
-    # assert_equal ["id"], user.attribute_names
-    # assert !user.has_attribute?("email")
-    # assert !user.has_attribute?(:email)
+    assert_equal ["id"], user.attributes.keys
+    assert_equal ["id"], user.attribute_names
+    assert !user.has_attribute?("email")
+    assert !user.has_attribute?(:email)
 
     user = User.select("id AS email_ciphertext").last
     assert_raises(Lockbox::DecryptionError) do